init

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+*.dat filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,4 @@
+venv
+__pycache__
+samples
+output_videos

README.md

@@ -0,0 +1,109 @@
+---
+license: mit
+---
+
+# LipCoordNet: Enhanced Lip Reading with Landmark Coordinates
+
+## Introduction
+
+LipCoordNet is an advanced neural network model designed for accurate lip reading by incorporating lip landmark coordinates as a supplementary input to the traditional image sequence input. This enhancement to the original LipNet architecture aims to improve the precision of sentence predictions by providing additional geometric context to the model.
+
+## Features
+
+- **Dual Input System**: Utilizes both raw image sequences and corresponding lip landmark coordinates for improved context.
+- **Enhanced Spatial Resolution**: Improved spatial analysis of lip movements through detailed landmark tracking.
+- **State-of-the-Art Performance**: Outperforms the original LipNet, as well as VIPL's [PyTorch implementation of LipNet](https://github.com/VIPL-Audio-Visual-Speech-Understanding/LipNet-PyTorch).
+  | Scenario | Image Size (W x H) | CER | WER |
+  | :-------------------------------: | :----------------: | :--: | :---: |
+  | Unseen speakers (Original) | 100 x 50 | 6.7% | 13.6% |
+  | Overlapped speakers (Original) | 100 x 50 | 2.0% | 5.6% |
+  | Unseen speakers (VIPL LipNet) | 128 x 64 | 6.7% | 13.3% |
+  | Overlapped speakers (VIPL LipNet) | 128 x 64 | 1.9% | 4.6% |
+  | Overlapped speakers (LipCoordNet) | 128 x 64 | 0.6% | 1.7% |
+
+## Getting Started
+
+### Prerequisites
+
+- Python 3.10 or later
+- Pytorch 2.0 or later
+- OpenCV
+- NumPy
+- dlib (for landmark detection)
+- The detailed list of dependencies can be found in `requirements.txt`.
+
+### Installation
+
+1. Clone the repository:
+
+   ```bash
+   git clone https://huggingface.co/SilentSpeak/LipCoordNet
+   ```
+
+2. Navigate to the project directory:
+   ```bash
+   cd LipCoordNet
+   ```
+3. Install the required dependencies:
+   ```bash
+   pip install -r requirements.txt
+   ```
+
+### Usage
+
+To train the LipCoordNet model with your dataset, first update the options.py file with the appropriate paths to your dataset and pretrained weights (comment out the weights if you want to start from scratch). Then, run the following command:
+
+```bash
+python train.py
+```
+
+To perform sentence prediction using the pre-trained model:
+
+```bash
+python inference.py --input_video <path_to_video>
+```
+
+note: ffmpeg is required to convert video to image sequence and run the inference script.
+
+## Model Architecture
+
+![LipCoordNet model architecture](./assets/LipCoordNet_model_architecture.png)
+
+## Training
+
+This model is built on top of the [LipNet-Pytorch](https://github.com/VIPL-Audio-Visual-Speech-Understanding/LipNet-PyTorch) project on GitHub. The training process if similar to the original LipNet model, with the addition of landmark coordinates as a supplementary input. We used the pretrained weights from the original LipNet model as a starting point for training our model, froze the weights for the original LipNet layers, and trained the new layers for the landmark coordinates.
+
+The dataset used to train this model is the [EGCLLC dataset](https://huggingface.co/datasets/SilentSpeak/EGCLLC). The dataset is not included in this repository, but can be downloaded from the link above.
+
+Total training time: 2 days
+Total epochs: 51
+Training hardware: NVIDIA GeForce RTX 3080 12GB
+
+![LipCoordNet training curves](./assets/training_graphs.png)
+
+For an interactive view of the training curves, please refer to the tensorboard logs in the `runs` directory.
+Use this command to view the logs:
+
+```bash
+tensorboard --logdir runs
+```
+
+## Evaluation
+
+We achieved a lowest WER of 1.7%, CER of 0.6% and a loss of 0.0256 on the validation dataset.
+
+## License
+
+This project is licensed under the MIT License.
+
+## Acknowledgments
+
+This model, LipCoordNet, has been developed with reference to the LipNet-PyTorch implementation available at [VIPL-Audio-Visual-Speech-Understanding](https://github.com/VIPL-Audio-Visual-Speech-Understanding/LipNet-PyTorch). We extend our gratitude to the contributors of this repository for providing a solid foundation and insightful examples that greatly facilitated the development of our enhanced lip reading model. Their work has been instrumental in advancing the field of audio-visual speech understanding and has provided the community with valuable resources to build upon.
+
+Alvarez Casado, C., Bordallo Lopez, M. Real-time face alignment: evaluation methods, training strategies and implementation optimization. Springer Journal of Real-time image processing, 2021
+
+Assael, Y., Shillingford, B., Whiteson, S., & Freitas, N. (2017). LipNet: End-to-End Sentence-level Lipreading. GPU Technology Conference.
+
+## Contact
+
+Project Link: https://github.com/ffeew/LipCoordNet

cvtransforms.py

@@ -0,0 +1,13 @@
+import random
+
+
+def HorizontalFlip(batch_img, p=0.5):
+    # (T, H, W, C)
+    if random.random() > p:
+        batch_img = batch_img[:, :, ::-1, ...]
+    return batch_img
+
+
+def ColorNormalize(batch_img):
+    batch_img = batch_img / 255.0
+    return batch_img

dataset.py

@@ -0,0 +1,190 @@
+import numpy as np
+import cv2
+import os
+from torch.utils.data import Dataset
+from cvtransforms import *
+import torch
+import editdistance
+import json
+
+
+class MyDataset(Dataset):
+    letters = [
+        " ",
+        "A",
+        "B",
+        "C",
+        "D",
+        "E",
+        "F",
+        "G",
+        "H",
+        "I",
+        "J",
+        "K",
+        "L",
+        "M",
+        "N",
+        "O",
+        "P",
+        "Q",
+        "R",
+        "S",
+        "T",
+        "U",
+        "V",
+        "W",
+        "X",
+        "Y",
+        "Z",
+    ]
+
+    def __init__(
+        self,
+        video_path,
+        anno_path,
+        coords_path,
+        file_list,
+        vid_pad,
+        txt_pad,
+        phase,
+    ):
+        self.anno_path = anno_path
+        self.coords_path = coords_path
+        self.vid_pad = vid_pad
+        self.txt_pad = txt_pad
+        self.phase = phase
+
+        with open(file_list, "r") as f:
+            self.videos = [
+                os.path.join(video_path, line.strip()) for line in f.readlines()
+            ]
+
+        self.data = []
+        for vid in self.videos:
+            items = vid.split("/")
+            self.data.append((vid, items[-4], items[-1]))
+
+    def __getitem__(self, idx):
+        (vid, spk, name) = self.data[idx]
+        vid = self._load_vid(vid)
+        anno = self._load_anno(
+            os.path.join(self.anno_path, spk, "align", name + ".align")
+        )
+        coord = self._load_coords(os.path.join(self.coords_path, spk, name + ".json"))
+
+        if self.phase == "train":
+            vid = HorizontalFlip(vid)
+
+        vid = ColorNormalize(vid)
+
+        vid_len = vid.shape[0]
+        anno_len = anno.shape[0]
+        vid = self._padding(vid, self.vid_pad)
+        anno = self._padding(anno, self.txt_pad)
+        coord = self._padding(coord, self.vid_pad)
+
+        return {
+            "vid": torch.FloatTensor(vid.transpose(3, 0, 1, 2)),
+            "txt": torch.LongTensor(anno),
+            "coord": torch.FloatTensor(coord),
+            "txt_len": anno_len,
+            "vid_len": vid_len,
+        }
+
+    def __len__(self):
+        return len(self.data)
+
+    def _load_vid(self, p):
+        files = os.listdir(p)
+        files = list(filter(lambda file: file.find(".jpg") != -1, files))
+        files = sorted(files, key=lambda file: int(os.path.splitext(file)[0]))
+        array = [cv2.imread(os.path.join(p, file)) for file in files]
+        array = list(filter(lambda im: not im is None, array))
+        array = [
+            cv2.resize(im, (128, 64), interpolation=cv2.INTER_LANCZOS4) for im in array
+        ]
+        array = np.stack(array, axis=0).astype(np.float32)
+
+        return array
+
+    def _load_anno(self, name):
+        with open(name, "r") as f:
+            lines = [line.strip().split(" ") for line in f.readlines()]
+            txt = [line[2] for line in lines]
+            txt = list(filter(lambda s: not s.upper() in ["SIL", "SP"], txt))
+        return MyDataset.txt2arr(" ".join(txt).upper(), 1)
+
+    def _load_coords(self, name):
+        # obtained from the resized image in the lip coordinate extraction
+        img_width = 600
+        img_height = 500
+        with open(name, "r") as f:
+            coords_data = json.load(f)
+
+        coords = []
+        for frame in sorted(coords_data.keys(), key=int):
+            frame_coords = coords_data[frame]
+
+            # Normalize the coordinates
+            normalized_coords = []
+            for x, y in zip(frame_coords[0], frame_coords[1]):
+                normalized_x = x / img_width
+                normalized_y = y / img_height
+                normalized_coords.append((normalized_x, normalized_y))
+
+            coords.append(normalized_coords)
+        coords_array = np.array(coords, dtype=np.float32)
+        return coords_array
+
+    def _padding(self, array, length):
+        array = [array[_] for _ in range(array.shape[0])]
+        size = array[0].shape
+        for i in range(length - len(array)):
+            array.append(np.zeros(size))
+        return np.stack(array, axis=0)
+
+    @staticmethod
+    def txt2arr(txt, start):
+        arr = []
+        for c in list(txt):
+            arr.append(MyDataset.letters.index(c) + start)
+        return np.array(arr)
+
+    @staticmethod
+    def arr2txt(arr, start):
+        txt = []
+        for n in arr:
+            if n >= start:
+                txt.append(MyDataset.letters[n - start])
+        return "".join(txt).strip()
+
+    @staticmethod
+    def ctc_arr2txt(arr, start):
+        pre = -1
+        txt = []
+        for n in arr:
+            if pre != n and n >= start:
+                if (
+                    len(txt) > 0
+                    and txt[-1] == " "
+                    and MyDataset.letters[n - start] == " "
+                ):
+                    pass
+                else:
+                    txt.append(MyDataset.letters[n - start])
+            pre = n
+        return "".join(txt).strip()
+
+    @staticmethod
+    def wer(predict, truth):
+        word_pairs = [(p[0].split(" "), p[1].split(" ")) for p in zip(predict, truth)]
+        wer = [1.0 * editdistance.eval(p[0], p[1]) / len(p[1]) for p in word_pairs]
+        return wer
+
+    @staticmethod
+    def cer(predict, truth):
+        cer = [
+            1.0 * editdistance.eval(p[0], p[1]) / len(p[1]) for p in zip(predict, truth)
+        ]
+        return cer

inference.py

@@ -0,0 +1,350 @@
+import argparse
+import os
+from model import LipCoordNet
+from dataset import MyDataset
+import torch
+import cv2
+import face_alignment
+import numpy as np
+import dlib
+import glob
+
+
+def get_position(size, padding=0.25):
+    x = [
+        0.000213256,
+        0.0752622,
+        0.18113,
+        0.29077,
+        0.393397,
+        0.586856,
+        0.689483,
+        0.799124,
+        0.904991,
+        0.98004,
+        0.490127,
+        0.490127,
+        0.490127,
+        0.490127,
+        0.36688,
+        0.426036,
+        0.490127,
+        0.554217,
+        0.613373,
+        0.121737,
+        0.187122,
+        0.265825,
+        0.334606,
+        0.260918,
+        0.182743,
+        0.645647,
+        0.714428,
+        0.793132,
+        0.858516,
+        0.79751,
+        0.719335,
+        0.254149,
+        0.340985,
+        0.428858,
+        0.490127,
+        0.551395,
+        0.639268,
+        0.726104,
+        0.642159,
+        0.556721,
+        0.490127,
+        0.423532,
+        0.338094,
+        0.290379,
+        0.428096,
+        0.490127,
+        0.552157,
+        0.689874,
+        0.553364,
+        0.490127,
+        0.42689,
+    ]
+
+    y = [
+        0.106454,
+        0.038915,
+        0.0187482,
+        0.0344891,
+        0.0773906,
+        0.0773906,
+        0.0344891,
+        0.0187482,
+        0.038915,
+        0.106454,
+        0.203352,
+        0.307009,
+        0.409805,
+        0.515625,
+        0.587326,
+        0.609345,
+        0.628106,
+        0.609345,
+        0.587326,
+        0.216423,
+        0.178758,
+        0.179852,
+        0.231733,
+        0.245099,
+        0.244077,
+        0.231733,
+        0.179852,
+        0.178758,
+        0.216423,
+        0.244077,
+        0.245099,
+        0.780233,
+        0.745405,
+        0.727388,
+        0.742578,
+        0.727388,
+        0.745405,
+        0.780233,
+        0.864805,
+        0.902192,
+        0.909281,
+        0.902192,
+        0.864805,
+        0.784792,
+        0.778746,
+        0.785343,
+        0.778746,
+        0.784792,
+        0.824182,
+        0.831803,
+        0.824182,
+    ]
+
+    x, y = np.array(x), np.array(y)
+
+    x = (x + padding) / (2 * padding + 1)
+    y = (y + padding) / (2 * padding + 1)
+    x = x * size
+    y = y * size
+    return np.array(list(zip(x, y)))
+
+
+def transformation_from_points(points1, points2):
+    points1 = points1.astype(np.float64)
+    points2 = points2.astype(np.float64)
+
+    c1 = np.mean(points1, axis=0)
+    c2 = np.mean(points2, axis=0)
+    points1 -= c1
+    points2 -= c2
+    s1 = np.std(points1)
+    s2 = np.std(points2)
+    points1 /= s1
+    points2 /= s2
+
+    U, S, Vt = np.linalg.svd(points1.T * points2)
+    R = (U * Vt).T
+    return np.vstack(
+        [
+            np.hstack(((s2 / s1) * R, c2.T - (s2 / s1) * R * c1.T)),
+            np.matrix([0.0, 0.0, 1.0]),
+        ]
+    )
+
+
+def load_video(file, device: str):
+    # create the samples directory if it doesn't exist
+    if not os.path.exists("samples"):
+        os.makedirs("samples")
+
+    p = os.path.join("samples")
+    output = os.path.join("samples", "%04d.jpg")
+    cmd = "ffmpeg -hide_banner -loglevel error -i {} -qscale:v 2 -r 25 {}".format(
+        file, output
+    )
+    os.system(cmd)
+
+    files = os.listdir(p)
+    files = sorted(files, key=lambda x: int(os.path.splitext(x)[0]))
+
+    array = [cv2.imread(os.path.join(p, file)) for file in files]
+
+    array = list(filter(lambda im: not im is None, array))
+
+    fa = face_alignment.FaceAlignment(
+        face_alignment.LandmarksType._2D, flip_input=False, device=device
+    )
+    points = [fa.get_landmarks(I) for I in array]
+
+    front256 = get_position(256)
+    video = []
+    for point, scene in zip(points, array):
+        if point is not None:
+            shape = np.array(point[0])
+            shape = shape[17:]
+            M = transformation_from_points(np.matrix(shape), np.matrix(front256))
+
+            img = cv2.warpAffine(scene, M[:2], (256, 256))
+            (x, y) = front256[-20:].mean(0).astype(np.int32)
+            w = 160 // 2
+            img = img[y - w // 2 : y + w // 2, x - w : x + w, ...]
+            img = cv2.resize(img, (128, 64))
+            video.append(img)
+
+    video = np.stack(video, axis=0).astype(np.float32)
+    video = torch.FloatTensor(video.transpose(3, 0, 1, 2)) / 255.0
+
+    return video
+
+
+def extract_lip_coordinates(detector, predictor, img_path):
+    image = cv2.imread(img_path)
+    image = cv2.resize(image, (600, 500))
+    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
+
+    rects = detector(gray)
+    retries = 3
+    while retries > 0:
+        try:
+            assert len(rects) == 1
+            break
+        except AssertionError as e:
+            retries -= 1
+
+    for rect in rects:
+        # apply the shape predictor to the face ROI
+        shape = predictor(gray, rect)
+        x = []
+        y = []
+        for n in range(48, 68):
+            x.append(shape.part(n).x)
+            y.append(shape.part(n).y)
+    return [x, y]
+
+
+def generate_lip_coordinates(frame_images_directory, detector, predictor):
+    frames = glob.glob(frame_images_directory + "/*.jpg")
+    frames.sort()
+
+    img = cv2.imread(frames[0])
+    height, width, layers = img.shape
+
+    coords = []
+    for frame in frames:
+        x_coords, y_coords = extract_lip_coordinates(detector, predictor, frame)
+        normalized_coords = []
+        for x, y in zip(x_coords, y_coords):
+            normalized_x = x / width
+            normalized_y = y / height
+            normalized_coords.append((normalized_x, normalized_y))
+        coords.append(normalized_coords)
+    coords_array = np.array(coords, dtype=np.float32)
+    coords_array = torch.from_numpy(coords_array)
+    return coords_array
+
+
+def ctc_decode(y):
+    y = y.argmax(-1)
+    t = y.size(0)
+    result = []
+    for i in range(t + 1):
+        result.append(MyDataset.ctc_arr2txt(y[:i], start=1))
+    return result
+
+
+def output_video(p, txt, output_path):
+    files = os.listdir(p)
+    files = sorted(files, key=lambda x: int(os.path.splitext(x)[0]))
+
+    font = cv2.FONT_HERSHEY_SIMPLEX
+
+    for file, line in zip(files, txt):
+        img = cv2.imread(os.path.join(p, file))
+        h, w, _ = img.shape
+        img = cv2.putText(
+            img, line, (w // 8, 11 * h // 12), font, 1.2, (0, 0, 0), 3, cv2.LINE_AA
+        )
+        img = cv2.putText(
+            img,
+            line,
+            (w // 8, 11 * h // 12),
+            font,
+            1.2,
+            (255, 255, 255),
+            0,
+            cv2.LINE_AA,
+        )
+        h = h // 2
+        w = w // 2
+        img = cv2.resize(img, (w, h))
+        cv2.imwrite(os.path.join(p, file), img)
+
+    # create the output_videos directory if it doesn't exist
+    if not os.path.exists(output_path):
+        os.makedirs(output_path)
+
+    output = os.path.join(output_path, "output.mp4")
+    cmd = "ffmpeg -hide_banner -loglevel error -y -i {}/%04d.jpg -r 25 {}".format(
+        p, output
+    )
+    os.system(cmd)
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--weights",
+        type=str,
+        default="pretrain/LipCoordNet_coords_loss_0.025581153109669685_wer_0.01746208431890914_cer_0.006488426950253695.pt",
+        help="path to the weights file",
+    )
+    parser.add_argument(
+        "--input_video",
+        type=str,
+        help="path to the input video frames",
+    )
+    parser.add_argument(
+        "--device",
+        type=str,
+        default="cuda",
+        help="device to run the model on",
+    )
+
+    parser.add_argument(
+        "--output_path",
+        type=str,
+        default="output_videos",
+        help="directory to save the output video",
+    )
+
+    args = parser.parse_args()
+
+    # validate if device is valid
+    if args.device not in ("cuda", "cpu"):
+        raise ValueError("Invalid device, must be either cuda or cpu")
+
+    device = args.device
+
+    # load model
+    model = LipCoordNet()
+    model.load_state_dict(torch.load(args.weights))
+    model = model.to(device)
+    model.eval()
+    detector = dlib.get_frontal_face_detector()
+    predictor = dlib.shape_predictor(
+        "lip_coordinate_extraction/shape_predictor_68_face_landmarks_GTX.dat"
+    )
+
+    # load video
+    video = load_video(args.input_video, device)
+
+    # generate lip coordinates
+    coords = generate_lip_coordinates("samples", detector, predictor)
+
+    pred = model(video[None, ...].to(device), coords[None, ...].to(device))
+    output = ctc_decode(pred[0])
+    print(output[-1])
+    output_video("samples", output, args.output_path)
+
+
+if __name__ == "__main__":
+    main()

model.py

@@ -0,0 +1,113 @@
+import torch
+import torch.nn as nn
+import torch.nn.init as init
+import math
+
+
+class LipCoordNet(torch.nn.Module):
+    def __init__(self, dropout_p=0.5, coord_input_dim=40, coord_hidden_dim=128):
+        super(LipCoordNet, self).__init__()
+        self.conv1 = nn.Conv3d(3, 32, (3, 5, 5), (1, 2, 2), (1, 2, 2))
+        self.pool1 = nn.MaxPool3d((1, 2, 2), (1, 2, 2))
+
+        self.conv2 = nn.Conv3d(32, 64, (3, 5, 5), (1, 1, 1), (1, 2, 2))
+        self.pool2 = nn.MaxPool3d((1, 2, 2), (1, 2, 2))
+
+        self.conv3 = nn.Conv3d(64, 96, (3, 3, 3), (1, 1, 1), (1, 1, 1))
+        self.pool3 = nn.MaxPool3d((1, 2, 2), (1, 2, 2))
+
+        self.gru1 = nn.GRU(96 * 4 * 8, 256, 1, bidirectional=True)
+        self.gru2 = nn.GRU(512, 256, 1, bidirectional=True)
+
+        self.FC = nn.Linear(512 + 2 * coord_hidden_dim, 27 + 1)
+        self.dropout_p = dropout_p
+
+        self.relu = nn.ReLU(inplace=True)
+        self.dropout = nn.Dropout(self.dropout_p)
+        self.dropout3d = nn.Dropout3d(self.dropout_p)
+
+        # New GRU layer for lip coordinates
+        self.coord_gru = nn.GRU(
+            coord_input_dim, coord_hidden_dim, 1, bidirectional=True
+        )
+
+        self._init()
+
+    def _init(self):
+        init.kaiming_normal_(self.conv1.weight, nonlinearity="relu")
+        init.constant_(self.conv1.bias, 0)
+
+        init.kaiming_normal_(self.conv2.weight, nonlinearity="relu")
+        init.constant_(self.conv2.bias, 0)
+
+        init.kaiming_normal_(self.conv3.weight, nonlinearity="relu")
+        init.constant_(self.conv3.bias, 0)
+
+        init.kaiming_normal_(self.FC.weight, nonlinearity="sigmoid")
+        init.constant_(self.FC.bias, 0)
+
+        for m in (self.gru1, self.gru2):
+            stdv = math.sqrt(2 / (96 * 3 * 6 + 256))
+            for i in range(0, 256 * 3, 256):
+                init.uniform_(
+                    m.weight_ih_l0[i : i + 256],
+                    -math.sqrt(3) * stdv,
+                    math.sqrt(3) * stdv,
+                )
+                init.orthogonal_(m.weight_hh_l0[i : i + 256])
+                init.constant_(m.bias_ih_l0[i : i + 256], 0)
+                init.uniform_(
+                    m.weight_ih_l0_reverse[i : i + 256],
+                    -math.sqrt(3) * stdv,
+                    math.sqrt(3) * stdv,
+                )
+                init.orthogonal_(m.weight_hh_l0_reverse[i : i + 256])
+                init.constant_(m.bias_ih_l0_reverse[i : i + 256], 0)
+
+    def forward(self, x, coords):
+        # branch 1
+        x = self.conv1(x)
+        x = self.relu(x)
+        x = self.dropout3d(x)
+        x = self.pool1(x)
+
+        x = self.conv2(x)
+        x = self.relu(x)
+        x = self.dropout3d(x)
+        x = self.pool2(x)
+
+        x = self.conv3(x)
+        x = self.relu(x)
+        x = self.dropout3d(x)
+        x = self.pool3(x)
+
+        # (B, C, T, H, W)->(T, B, C, H, W)
+        x = x.permute(2, 0, 1, 3, 4).contiguous()
+        # (B, C, T, H, W)->(T, B, C*H*W)
+        x = x.view(x.size(0), x.size(1), -1)
+
+        self.gru1.flatten_parameters()
+        self.gru2.flatten_parameters()
+
+        x, h = self.gru1(x)
+        x = self.dropout(x)
+        x, h = self.gru2(x)
+        x = self.dropout(x)
+
+        # branch 2
+        # Process lip coordinates through GRU
+        self.coord_gru.flatten_parameters()
+
+        # (B, T, N, C)->(T, B, C, N, C)
+        coords = coords.permute(1, 0, 2, 3).contiguous()
+        # (T, B, C, N, C)->(T, B, C, N*C)
+        coords = coords.view(coords.size(0), coords.size(1), -1)
+        coords, _ = self.coord_gru(coords)
+        coords = self.dropout(coords)
+
+        # combine the two branches
+        combined = torch.cat((x, coords), dim=2)
+
+        x = self.FC(combined)
+        x = x.permute(1, 0, 2).contiguous()
+        return x

options.py

@@ -0,0 +1,21 @@
+gpu = "0"
+random_seed = 0
+data_type = "coords"
+video_path = "lip_images/"
+train_list = f"data/{data_type}_train.txt"
+val_list = f"data/{data_type}_val.txt"
+anno_path = "GRID_alignments"
+coords_path = "lip_coordinates"
+vid_padding = 75
+txt_padding = 200
+batch_size = 40
+base_lr = 2e-5
+num_workers = 16
+max_epoch = 10000
+display = 50
+test_step = 1000
+save_prefix = f"weights/LipCoordNet_{data_type}"
+is_optimize = True
+pin_memory = True
+
+weights = "pretrain/LipCoordNet_coords_loss_0.025581153109669685_wer_0.01746208431890914_cer_0.006488426950253695.pt"

train.py

@@ -0,0 +1,237 @@
+import torch
+import torch.nn as nn
+from torch.utils.data import DataLoader
+import os
+from dataset import MyDataset
+import numpy as np
+import time
+from model import LipCoordNet
+import torch.optim as optim
+from tensorboardX import SummaryWriter
+import options as opt
+from tqdm import tqdm
+
+
+def dataset2dataloader(dataset, num_workers=opt.num_workers, shuffle=True):
+    return DataLoader(
+        dataset,
+        batch_size=opt.batch_size,
+        shuffle=shuffle,
+        num_workers=num_workers,
+        drop_last=False,
+        pin_memory=opt.pin_memory,
+    )
+
+
+def show_lr(optimizer):
+    lr = []
+    for param_group in optimizer.param_groups:
+        lr += [param_group["lr"]]
+    return np.array(lr).mean()
+
+
+def ctc_decode(y):
+    y = y.argmax(-1)
+    return [MyDataset.ctc_arr2txt(y[_], start=1) for _ in range(y.size(0))]
+
+
+def test(model, net):
+    with torch.no_grad():
+        dataset = MyDataset(
+            opt.video_path,
+            opt.anno_path,
+            opt.coords_path,
+            opt.val_list,
+            opt.vid_padding,
+            opt.txt_padding,
+            "test",
+        )
+
+        print("num_test_data:{}".format(len(dataset.data)))
+        model.eval()
+        loader = dataset2dataloader(dataset, shuffle=False)
+        loss_list = []
+        wer = []
+        cer = []
+        crit = nn.CTCLoss()
+        tic = time.time()
+        print("RUNNING VALIDATION")
+        pbar = tqdm(loader)
+        for i_iter, input in enumerate(pbar):
+            vid = input.get("vid").cuda(non_blocking=opt.pin_memory)
+            txt = input.get("txt").cuda(non_blocking=opt.pin_memory)
+            vid_len = input.get("vid_len").cuda(non_blocking=opt.pin_memory)
+            txt_len = input.get("txt_len").cuda(non_blocking=opt.pin_memory)
+            coord = input.get("coord").cuda(non_blocking=opt.pin_memory)
+
+            y = net(vid, coord)
+
+            loss = (
+                crit(
+                    y.transpose(0, 1).log_softmax(-1),
+                    txt,
+                    vid_len.view(-1),
+                    txt_len.view(-1),
+                )
+                .detach()
+                .cpu()
+                .numpy()
+            )
+            loss_list.append(loss)
+            pred_txt = ctc_decode(y)
+
+            truth_txt = [MyDataset.arr2txt(txt[_], start=1) for _ in range(txt.size(0))]
+            wer.extend(MyDataset.wer(pred_txt, truth_txt))
+            cer.extend(MyDataset.cer(pred_txt, truth_txt))
+            if i_iter % opt.display == 0:
+                v = 1.0 * (time.time() - tic) / (i_iter + 1)
+                eta = v * (len(loader) - i_iter) / 3600.0
+
+                print("".join(101 * "-"))
+                print("{:<50}|{:>50}".format("predict", "truth"))
+                print("".join(101 * "-"))
+                for predict, truth in list(zip(pred_txt, truth_txt))[:10]:
+                    print("{:<50}|{:>50}".format(predict, truth))
+                print("".join(101 * "-"))
+                print(
+                    "test_iter={},eta={},wer={},cer={}".format(
+                        i_iter, eta, np.array(wer).mean(), np.array(cer).mean()
+                    )
+                )
+                print("".join(101 * "-"))
+
+        return (np.array(loss_list).mean(), np.array(wer).mean(), np.array(cer).mean())
+
+
+def train(model, net):
+    dataset = MyDataset(
+        opt.video_path,
+        opt.anno_path,
+        opt.coords_path,
+        opt.train_list,
+        opt.vid_padding,
+        opt.txt_padding,
+        "train",
+    )
+
+    loader = dataset2dataloader(dataset)
+    optimizer = optim.Adam(
+        model.parameters(), lr=opt.base_lr, weight_decay=0.0, amsgrad=True
+    )
+
+    print("num_train_data:{}".format(len(dataset.data)))
+    crit = nn.CTCLoss()
+    tic = time.time()
+
+    train_wer = []
+    for epoch in range(opt.max_epoch):
+        print(f"RUNNING EPOCH {epoch}")
+        pbar = tqdm(loader)
+
+        for i_iter, input in enumerate(pbar):
+            model.train()
+            vid = input.get("vid").cuda(non_blocking=opt.pin_memory)
+            txt = input.get("txt").cuda(non_blocking=opt.pin_memory)
+            vid_len = input.get("vid_len").cuda(non_blocking=opt.pin_memory)
+            txt_len = input.get("txt_len").cuda(non_blocking=opt.pin_memory)
+            coord = input.get("coord").cuda(non_blocking=opt.pin_memory)
+
+            optimizer.zero_grad()
+            y = net(vid, coord)
+            loss = crit(
+                y.transpose(0, 1).log_softmax(-1),
+                txt,
+                vid_len.view(-1),
+                txt_len.view(-1),
+            )
+            loss.backward()
+
+            if opt.is_optimize:
+                optimizer.step()
+
+            tot_iter = i_iter + epoch * len(loader)
+
+            pred_txt = ctc_decode(y)
+
+            truth_txt = [MyDataset.arr2txt(txt[_], start=1) for _ in range(txt.size(0))]
+            train_wer.extend(MyDataset.wer(pred_txt, truth_txt))
+
+            if tot_iter % opt.display == 0:
+                v = 1.0 * (time.time() - tic) / (tot_iter + 1)
+                eta = (len(loader) - i_iter) * v / 3600.0
+
+                writer.add_scalar("train loss", loss, tot_iter)
+                writer.add_scalar("train wer", np.array(train_wer).mean(), tot_iter)
+                print("".join(101 * "-"))
+                print("{:<50}|{:>50}".format("predict", "truth"))
+                print("".join(101 * "-"))
+
+                for predict, truth in list(zip(pred_txt, truth_txt))[:3]:
+                    print("{:<50}|{:>50}".format(predict, truth))
+                print("".join(101 * "-"))
+                print(
+                    "epoch={},tot_iter={},eta={},loss={},train_wer={}".format(
+                        epoch, tot_iter, eta, loss, np.array(train_wer).mean()
+                    )
+                )
+                print("".join(101 * "-"))
+
+            if tot_iter % opt.test_step == 0:
+                (loss, wer, cer) = test(model, net)
+                print(
+                    "i_iter={},lr={},loss={},wer={},cer={}".format(
+                        tot_iter, show_lr(optimizer), loss, wer, cer
+                    )
+                )
+                writer.add_scalar("val loss", loss, tot_iter)
+                writer.add_scalar("wer", wer, tot_iter)
+                writer.add_scalar("cer", cer, tot_iter)
+                savename = "{}_loss_{}_wer_{}_cer_{}.pt".format(
+                    opt.save_prefix, loss, wer, cer
+                )
+                (path, name) = os.path.split(savename)
+                if not os.path.exists(path):
+                    os.makedirs(path)
+                torch.save(model.state_dict(), savename)
+                if not opt.is_optimize:
+                    exit()
+
+
+if __name__ == "__main__":
+    print("Loading options...")
+    os.environ["CUDA_VISIBLE_DEVICES"] = opt.gpu
+    writer = SummaryWriter()
+    model = LipCoordNet()
+    model = model.cuda()
+    net = nn.DataParallel(model).cuda()
+
+    if hasattr(opt, "weights"):
+        pretrained_dict = torch.load(opt.weights)
+        model_dict = model.state_dict()
+        pretrained_dict = {
+            k: v
+            for k, v in pretrained_dict.items()
+            if k in model_dict.keys() and v.size() == model_dict[k].size()
+        }
+
+        # freeze the pretrained layers
+        for k, param in pretrained_dict.items():
+            param.requires_grad = False
+
+        missed_params = [
+            k for k, v in model_dict.items() if not k in pretrained_dict.keys()
+        ]
+        print(
+            "loaded params/tot params:{}/{}".format(
+                len(pretrained_dict), len(model_dict)
+            )
+        )
+        print("miss matched params:{}".format(missed_params))
+        model_dict.update(pretrained_dict)
+        model.load_state_dict(model_dict)
+
+    torch.manual_seed(opt.random_seed)
+    torch.cuda.manual_seed_all(opt.random_seed)
+    torch.backends.cudnn.benchmark = True
+
+    train(model, net)