feat: implement initial facial expression recognition training pipeline with TTA support

.gitignore

@@ -0,0 +1,140 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+pip-wheel-metadata/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+.python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# celery beat schedule file
+celerybeat-schedule
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+####
+saved/checkpoints/*
+saved/logs/*
+debug
+ssh.sh
+vim-markdown-preview.html
+deploy.prototxt.txt
+res10_300x300_ssd_iter_140000.caffemodel
+run.sh
+saved
+pretrained_ckpt
+version.py
+*.swp
+.idea/

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2021 Luan Pham
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

app.py

@@ -0,0 +1,71 @@
+import gradio as gr
+import cv2
+import numpy as np
+from rmn import RMN
+
+# Khởi tạo model RMN.
+# Quá trình này sẽ mất một chút thời gian ở lần chạy đầu tiên để tải weights.
+m = RMN()
+
+def detect_emotion(image):
+    """
+    Hàm xử lý ảnh đầu vào, gọi model RMN và trả về ảnh đã vẽ bounding box.
+    """
+    if image is None:
+        return None
+        
+    # Chuyển đổi từ PIL / Gradio RGB sang OpenCV BGR để tương thích với RMN
+    image_bgr = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
+    
+    # Dự đoán cảm xúc
+    results = m.detect_emotion_for_single_frame(image_bgr)
+    
+    # RMN cung cấp sẵn hàm draw() để vẽ bounding box và label lên ảnh
+    image_with_boxes = m.draw(image_bgr, results)
+    
+    # Chuyển ngược lại từ BGR sang RGB để Gradio hiển thị đúng màu
+    output_image = cv2.cvtColor(image_with_boxes, cv2.COLOR_BGR2RGB)
+    
+    # Có thể format lại text kết quả để hiển thị ra màn hình
+    results_text = "No face detected"
+    if results:
+        results_text = "\n".join([f"Face {i+1}: {res['emo_label']} (Prob: {res['emo_proba']:.2f})" for i, res in enumerate(results)])
+        
+    return output_image, results_text
+
+# Khởi tạo giao diện Gradio
+with gr.Blocks(title="Facial Expression Recognition") as demo:
+    gr.Markdown("# 🎭 Facial Expression Recognition (RMN)")
+    gr.Markdown(
+        "Upload an image or use your webcam to detect facial expressions using the **Residual Masking Network (RMN)**. "
+        "This model achieves state-of-the-art results on the FER2013 dataset."
+    )
+    
+    with gr.Row():
+        with gr.Column():
+            # Đầu vào ảnh
+            input_image = gr.Image(label="Input Image", type="numpy")
+            btn = gr.Button("Detect Emotion", variant="primary")
+            
+        with gr.Column():
+            # Đầu ra ảnh và text
+            output_image = gr.Image(label="Detected Result", type="numpy")
+            output_text = gr.Textbox(label="Detection Details", lines=3)
+            
+    # Kết nối nút bấm với hàm xử lý
+    btn.click(fn=detect_emotion, inputs=input_image, outputs=[output_image, output_text])
+    
+    # Cho phép demo các ảnh có sẵn
+    gr.Examples(
+        examples=[
+            # Bạn có thể thêm đường dẫn các ảnh mẫu vào đây nếu có
+            # ["examples/happy.jpg"],
+            # ["examples/sad.jpg"]
+        ],
+        inputs=input_image
+    )
+
+# Chạy ứng dụng
+if __name__ == "__main__":
+    # share=True tạo ra một đường link public tạm thời (rất tiện để khoe nhanh)
+    demo.launch(server_name="0.0.0.0", server_port=7860, share=True)

demo.py

@@ -0,0 +1,5 @@
+from rmn import RMN
+
+if __name__ == "__main__":
+    rmn = RMN()
+    rmn.video_demo()

models/__init__.py

@@ -0,0 +1,3 @@
+from .resmasking import resmasking_dropout1
+
+__all__ = ["resmasking_dropout1"]

models/masking.py

@@ -0,0 +1,375 @@
+import traceback
+
+import torch
+import torch.nn as nn
+
+from .resnet import BasicBlock, conv1x1
+
+
+def up_pooling(in_channels, out_channels, kernel_size=2, stride=2):
+    return nn.Sequential(
+        nn.ConvTranspose2d(
+            in_channels, out_channels, kernel_size=kernel_size, stride=stride
+        ),
+        nn.BatchNorm2d(out_channels),
+        nn.ReLU(inplace=True),
+    )
+
+
+class Masking4(nn.Module):
+    def __init__(self, in_channels, out_channels, block=BasicBlock):
+        assert in_channels == out_channels
+        super(Masking4, self).__init__()
+        filters = [
+            in_channels,
+            in_channels * 2,
+            in_channels * 4,
+            in_channels * 8,
+            in_channels * 16,
+        ]
+
+        self.downsample1 = nn.Sequential(
+            conv1x1(filters[0], filters[1], 1),
+            nn.BatchNorm2d(filters[1]),
+        )
+
+        self.downsample2 = nn.Sequential(
+            conv1x1(filters[1], filters[2], 1),
+            nn.BatchNorm2d(filters[2]),
+        )
+
+        self.downsample3 = nn.Sequential(
+            conv1x1(filters[2], filters[3], 1),
+            nn.BatchNorm2d(filters[3]),
+        )
+
+        self.downsample4 = nn.Sequential(
+            conv1x1(filters[3], filters[4], 1),
+            nn.BatchNorm2d(filters[4]),
+        )
+
+        """
+        self.conv1 = block(filters[0], filters[1], downsample=conv1x1(filters[0], filters[1], 1))
+        self.conv2 = block(filters[1], filters[2], downsample=conv1x1(filters[1], filters[2], 1))
+        self.conv3 = block(filters[2], filters[3], downsample=conv1x1(filters[2], filters[3], 1))
+        """
+
+        self.conv1 = block(filters[0], filters[1], downsample=self.downsample1)
+        self.conv2 = block(filters[1], filters[2], downsample=self.downsample2)
+        self.conv3 = block(filters[2], filters[3], downsample=self.downsample3)
+        self.conv4 = block(filters[3], filters[4], downsample=self.downsample4)
+
+        self.down_pooling = nn.MaxPool2d(kernel_size=2)
+
+        self.downsample5 = nn.Sequential(
+            conv1x1(filters[4], filters[3], 1),
+            nn.BatchNorm2d(filters[3]),
+        )
+
+        self.downsample6 = nn.Sequential(
+            conv1x1(filters[3], filters[2], 1),
+            nn.BatchNorm2d(filters[2]),
+        )
+
+        self.downsample7 = nn.Sequential(
+            conv1x1(filters[2], filters[1], 1),
+            nn.BatchNorm2d(filters[1]),
+        )
+
+        self.downsample8 = nn.Sequential(
+            conv1x1(filters[1], filters[0], 1),
+            nn.BatchNorm2d(filters[0]),
+        )
+
+        """
+        self.up_pool4 = up_pooling(filters[3], filters[2])
+        self.conv4 = block(filters[3], filters[2], downsample=conv1x1(filters[3], filters[2], 1))
+        self.up_pool5 = up_pooling(filters[2], filters[1])
+        self.conv5 = block(filters[2], filters[1], downsample=conv1x1(filters[2], filters[1], 1))
+
+        self.conv6 = block(filters[1], filters[0], downsample=conv1x1(filters[1], filters[0], 1))
+        """
+
+        self.up_pool5 = up_pooling(filters[4], filters[3])
+        self.conv5 = block(filters[4], filters[3], downsample=self.downsample5)
+        self.up_pool6 = up_pooling(filters[3], filters[2])
+        self.conv6 = block(filters[3], filters[2], downsample=self.downsample6)
+        self.up_pool7 = up_pooling(filters[2], filters[1])
+        self.conv7 = block(filters[2], filters[1], downsample=self.downsample7)
+        self.conv8 = block(filters[1], filters[0], downsample=self.downsample8)
+
+        # init weight
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(m.weight, mode="fan_out", nonlinearity="relu")
+            elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+
+        # Zero-initialize the last BN in each residual branch,
+        # so that the residual branch starts with zeros, and each residual block behaves like an identity.
+        # This improves the model by 0.2~0.3% according to https://arxiv.org/abs/1706.02677
+        for m in self.modules():
+            if isinstance(m, BasicBlock):
+                nn.init.constant_(m.bn2.weight, 0)
+
+    def forward(self, x):
+        x1 = self.conv1(x)
+        p1 = self.down_pooling(x1)
+        x2 = self.conv2(p1)
+        p2 = self.down_pooling(x2)
+        x3 = self.conv3(p2)
+        p3 = self.down_pooling(x3)
+        x4 = self.conv4(p3)
+
+        x5 = self.up_pool5(x4)
+        x5 = torch.cat([x5, x3], dim=1)
+        x5 = self.conv5(x5)
+
+        x6 = self.up_pool6(x5)
+        x6 = torch.cat([x6, x2], dim=1)
+        x6 = self.conv6(x6)
+
+        x7 = self.up_pool7(x6)
+        x7 = torch.cat([x7, x1], dim=1)
+        x7 = self.conv7(x7)
+
+        x8 = self.conv8(x7)
+
+        output = torch.softmax(x8, dim=1)
+        # output = torch.sigmoid(x8)
+        return output
+
+
+class Masking3(nn.Module):
+    def __init__(self, in_channels, out_channels, block=BasicBlock):
+        assert in_channels == out_channels
+        super(Masking3, self).__init__()
+        filters = [in_channels, in_channels * 2, in_channels * 4, in_channels * 8]
+
+        self.downsample1 = nn.Sequential(
+            conv1x1(filters[0], filters[1], 1),
+            nn.BatchNorm2d(filters[1]),
+        )
+
+        self.downsample2 = nn.Sequential(
+            conv1x1(filters[1], filters[2], 1),
+            nn.BatchNorm2d(filters[2]),
+        )
+
+        self.downsample3 = nn.Sequential(
+            conv1x1(filters[2], filters[3], 1),
+            nn.BatchNorm2d(filters[3]),
+        )
+
+        """
+        self.conv1 = block(filters[0], filters[1], downsample=conv1x1(filters[0], filters[1], 1))
+        self.conv2 = block(filters[1], filters[2], downsample=conv1x1(filters[1], filters[2], 1))
+        self.conv3 = block(filters[2], filters[3], downsample=conv1x1(filters[2], filters[3], 1))
+        """
+
+        self.conv1 = block(filters[0], filters[1], downsample=self.downsample1)
+        self.conv2 = block(filters[1], filters[2], downsample=self.downsample2)
+        self.conv3 = block(filters[2], filters[3], downsample=self.downsample3)
+
+        self.down_pooling = nn.MaxPool2d(kernel_size=2)
+
+        self.downsample4 = nn.Sequential(
+            conv1x1(filters[3], filters[2], 1),
+            nn.BatchNorm2d(filters[2]),
+        )
+
+        self.downsample5 = nn.Sequential(
+            conv1x1(filters[2], filters[1], 1),
+            nn.BatchNorm2d(filters[1]),
+        )
+
+        self.downsample6 = nn.Sequential(
+            conv1x1(filters[1], filters[0], 1),
+            nn.BatchNorm2d(filters[0]),
+        )
+
+        """
+        self.up_pool4 = up_pooling(filters[3], filters[2])
+        self.conv4 = block(filters[3], filters[2], downsample=conv1x1(filters[3], filters[2], 1))
+        self.up_pool5 = up_pooling(filters[2], filters[1])
+        self.conv5 = block(filters[2], filters[1], downsample=conv1x1(filters[2], filters[1], 1))
+
+        self.conv6 = block(filters[1], filters[0], downsample=conv1x1(filters[1], filters[0], 1))
+        """
+
+        self.up_pool4 = up_pooling(filters[3], filters[2])
+        self.conv4 = block(filters[3], filters[2], downsample=self.downsample4)
+        self.up_pool5 = up_pooling(filters[2], filters[1])
+        self.conv5 = block(filters[2], filters[1], downsample=self.downsample5)
+
+        self.conv6 = block(filters[1], filters[0], downsample=self.downsample6)
+
+        # init weight
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(m.weight, mode="fan_out", nonlinearity="relu")
+            elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+
+        # Zero-initialize the last BN in each residual branch,
+        # so that the residual branch starts with zeros, and each residual block behaves like an identity.
+        # This improves the model by 0.2~0.3% according to https://arxiv.org/abs/1706.02677
+        for m in self.modules():
+            if isinstance(m, BasicBlock):
+                nn.init.constant_(m.bn2.weight, 0)
+
+    def forward(self, x):
+        x1 = self.conv1(x)
+        p1 = self.down_pooling(x1)
+        x2 = self.conv2(p1)
+        p2 = self.down_pooling(x2)
+        x3 = self.conv3(p2)
+
+        x4 = self.up_pool4(x3)
+        x4 = torch.cat([x4, x2], dim=1)
+
+        x4 = self.conv4(x4)
+
+        x5 = self.up_pool5(x4)
+        x5 = torch.cat([x5, x1], dim=1)
+        x5 = self.conv5(x5)
+
+        x6 = self.conv6(x5)
+
+        output = torch.softmax(x6, dim=1)
+        # output = torch.sigmoid(x6)
+        return output
+
+
+class Masking2(nn.Module):
+    def __init__(self, in_channels, out_channels, block=BasicBlock):
+        assert in_channels == out_channels
+        super(Masking2, self).__init__()
+        filters = [in_channels, in_channels * 2, in_channels * 4, in_channels * 8]
+
+        self.downsample1 = nn.Sequential(
+            conv1x1(filters[0], filters[1], 1),
+            nn.BatchNorm2d(filters[1]),
+        )
+
+        self.downsample2 = nn.Sequential(
+            conv1x1(filters[1], filters[2], 1),
+            nn.BatchNorm2d(filters[2]),
+        )
+
+        """
+        self.conv1 = block(filters[0], filters[1], downsample=conv1x1(filters[0], filters[1], 1))
+        self.conv2 = block(filters[1], filters[2], downsample=conv1x1(filters[1], filters[2], 1))
+        """
+        self.conv1 = block(filters[0], filters[1], downsample=self.downsample1)
+        self.conv2 = block(filters[1], filters[2], downsample=self.downsample2)
+
+        self.down_pooling = nn.MaxPool2d(kernel_size=2)
+
+        self.downsample3 = nn.Sequential(
+            conv1x1(filters[2], filters[1], 1),
+            nn.BatchNorm2d(filters[1]),
+        )
+
+        self.downsample4 = nn.Sequential(
+            conv1x1(filters[1], filters[0], 1),
+            nn.BatchNorm2d(filters[0]),
+        )
+
+        """
+        self.up_pool3 = up_pooling(filters[2], filters[1])
+        self.conv3 = block(filters[2], filters[1], downsample=conv1x1(filters[2], filters[1], 1))
+        self.conv4 = block(filters[1], filters[0], downsample=conv1x1(filters[1], filters[0], 1))
+        """
+        self.up_pool3 = up_pooling(filters[2], filters[1])
+        self.conv3 = block(filters[2], filters[1], downsample=self.downsample3)
+        self.conv4 = block(filters[1], filters[0], downsample=self.downsample4)
+
+        # init weight
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(m.weight, mode="fan_out", nonlinearity="relu")
+            elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+
+        # Zero-initialize the last BN in each residual branch,
+        # so that the residual branch starts with zeros, and each residual block behaves like an identity.
+        # This improves the model by 0.2~0.3% according to https://arxiv.org/abs/1706.02677
+        for m in self.modules():
+            if isinstance(m, BasicBlock):
+                nn.init.constant_(m.bn2.weight, 0)
+
+    def forward(self, x):
+        x1 = self.conv1(x)
+        p1 = self.down_pooling(x1)
+        x2 = self.conv2(p1)
+
+        x3 = self.up_pool3(x2)
+        x3 = torch.cat([x3, x1], dim=1)
+        x3 = self.conv3(x3)
+
+        x4 = self.conv4(x3)
+
+        output = torch.softmax(x4, dim=1)
+        # output = torch.sigmoid(x4)
+        return output
+
+
+class Masking1(nn.Module):
+    def __init__(self, in_channels, out_channels, block=BasicBlock):
+        assert in_channels == out_channels
+        super(Masking1, self).__init__()
+        filters = [in_channels, in_channels * 2, in_channels * 4, in_channels * 8]
+
+        self.downsample1 = nn.Sequential(
+            conv1x1(filters[0], filters[1], 1),
+            nn.BatchNorm2d(filters[1]),
+        )
+
+        self.conv1 = block(filters[0], filters[1], downsample=self.downsample1)
+
+        self.downsample2 = nn.Sequential(
+            conv1x1(filters[1], filters[0], 1),
+            nn.BatchNorm2d(filters[0]),
+        )
+
+        self.conv2 = block(filters[1], filters[0], downsample=self.downsample2)
+
+        # init weight
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(m.weight, mode="fan_out", nonlinearity="relu")
+            elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+
+        # Zero-initialize the last BN in each residual branch,
+        # so that the residual branch starts with zeros, and each residual block behaves like an identity.
+        # This improves the model by 0.2~0.3% according to https://arxiv.org/abs/1706.02677
+        for m in self.modules():
+            if isinstance(m, BasicBlock):
+                nn.init.constant_(m.bn2.weight, 0)
+
+    def forward(self, x):
+        x1 = self.conv1(x)
+        x2 = self.conv2(x1)
+        output = torch.softmax(x2, dim=1)
+        # output = torch.sigmoid(x2)
+        return output
+
+
+def masking(in_channels, out_channels, depth, block=BasicBlock):
+    if depth == 1:
+        return Masking1(in_channels, out_channels, block)
+    elif depth == 2:
+        return Masking2(in_channels, out_channels, block)
+    elif depth == 3:
+        return Masking3(in_channels, out_channels, block)
+    elif depth == 4:
+        return Masking4(in_channels, out_channels, block)
+    else:
+        traceback.print_exc()
+        raise Exception("depth need to be from 0-3")

models/resmasking.py

@@ -0,0 +1,60 @@
+import torch
+import torch.nn as nn
+
+from .masking import masking
+from .resnet import BasicBlock, ResNet
+
+
+class ResMasking(ResNet):
+    def __init__(self, in_channels=3, num_classes=7):
+        super(ResMasking, self).__init__(
+            block=BasicBlock,
+            layers=[3, 4, 6, 3],
+            in_channels=in_channels,
+            num_classes=1000,
+        )
+
+        self.fc = nn.Linear(512, num_classes)
+
+        self.mask1 = masking(64, 64, depth=4)
+        self.mask2 = masking(128, 128, depth=3)
+        self.mask3 = masking(256, 256, depth=2)
+        self.mask4 = masking(512, 512, depth=1)
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu(x)
+        x = self.maxpool(x)
+
+        x = self.layer1(x)
+        m = self.mask1(x)
+        x = x * (1 + m)
+
+        x = self.layer2(x)
+        m = self.mask2(x)
+        x = x * (1 + m)
+
+        x = self.layer3(x)
+        m = self.mask3(x)
+        x = x * (1 + m)
+
+        x = self.layer4(x)
+        m = self.mask4(x)
+        x = x * (1 + m)
+
+        x = self.avgpool(x)
+        x = torch.flatten(x, 1)
+        x = self.fc(x)
+        return x
+
+
+def resmasking_dropout1(in_channels=3, num_classes=7, weight_path=""):
+    del weight_path
+
+    model = ResMasking(in_channels=in_channels, num_classes=num_classes)
+    model.fc = nn.Sequential(
+        nn.Dropout(0.4),
+        nn.Linear(512, num_classes),
+    )
+    return model

models/resnet.py

@@ -0,0 +1,203 @@
+import torch
+import torch.nn as nn
+
+__all__ = [
+    "ResNet",
+    "BasicBlock",
+    "conv1x1"
+]
+
+
+def conv3x3(in_planes, out_planes, stride=1, groups=1, dilation=1):
+    """3x3 convolution with padding"""
+    return nn.Conv2d(
+        in_planes,
+        out_planes,
+        kernel_size=3,
+        stride=stride,
+        padding=dilation,
+        groups=groups,
+        bias=False,
+        dilation=dilation,
+    )
+
+
+def conv1x1(in_planes, out_planes, stride=1):
+    """1x1 convolution"""
+    return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)
+
+
+class BasicBlock(nn.Module):
+    expansion = 1
+    __constants__ = ["downsample"]
+
+    def __init__(
+        self,
+        inplanes,
+        planes,
+        stride=1,
+        downsample=None,
+        groups=1,
+        base_width=64,
+        dilation=1,
+        norm_layer=None,
+    ):
+        super(BasicBlock, self).__init__()
+        if norm_layer is None:
+            norm_layer = nn.BatchNorm2d
+        if groups != 1 or base_width != 64:
+            raise ValueError("BasicBlock only supports groups=1 and base_width=64")
+        if dilation > 1:
+            raise NotImplementedError("Dilation > 1 not supported in BasicBlock")
+        # Both self.conv1 and self.downsample layers downsample the input when stride != 1
+        self.conv1 = conv3x3(inplanes, planes, stride)
+        self.bn1 = norm_layer(planes)
+        self.relu = nn.ReLU(inplace=True)
+        self.conv2 = conv3x3(planes, planes)
+        self.bn2 = norm_layer(planes)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        identity = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+
+        if self.downsample is not None:
+            identity = self.downsample(x)
+
+        out += identity
+        out = self.relu(out)
+
+        return out
+
+
+class ResNet(nn.Module):
+    def __init__(
+        self,
+        block,
+        layers,
+        num_classes=1000,
+        zero_init_residual=False,
+        groups=1,
+        width_per_group=64,
+        replace_stride_with_dilation=None,
+        norm_layer=None,
+        in_channels=3,
+    ):
+        super(ResNet, self).__init__()
+        if norm_layer is None:
+            norm_layer = nn.BatchNorm2d
+        self._norm_layer = norm_layer
+
+        self.inplanes = 64
+        self.dilation = 1
+        if replace_stride_with_dilation is None:
+            # each element in the tuple indicates if we should replace
+            # the 2x2 stride with a dilated convolution instead
+            replace_stride_with_dilation = [False, False, False]
+        if len(replace_stride_with_dilation) != 3:
+            raise ValueError(
+                "replace_stride_with_dilation should be None "
+                "or a 3-element tuple, got {}".format(replace_stride_with_dilation)
+            )
+        self.groups = groups
+        self.base_width = width_per_group
+
+        # NOTE: strictly set the in_channels = 3 to load the pretrained model
+        self.conv1 = nn.Conv2d(
+            3, self.inplanes, kernel_size=7, stride=2, padding=3, bias=False
+        )
+        # self.conv1 = nn.Conv2d(in_channels, self.inplanes, kernel_size=7, stride=2, padding=3, bias=False)
+        self.bn1 = norm_layer(self.inplanes)
+        self.relu = nn.ReLU(inplace=True)
+        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+        self.layer1 = self._make_layer(block, 64, layers[0])
+        self.layer2 = self._make_layer(
+            block, 128, layers[1], stride=2, dilate=replace_stride_with_dilation[0]
+        )
+        self.layer3 = self._make_layer(
+            block, 256, layers[2], stride=2, dilate=replace_stride_with_dilation[1]
+        )
+        self.layer4 = self._make_layer(
+            block, 512, layers[3], stride=2, dilate=replace_stride_with_dilation[2]
+        )
+        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
+
+        # NOTE: strictly set the num_classes = 1000 to load the pretrained model
+        self.fc = nn.Linear(512 * block.expansion, 1000)
+
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(m.weight, mode="fan_out", nonlinearity="relu")
+            elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+
+        if zero_init_residual:
+            for m in self.modules():
+                if isinstance(m, BasicBlock):
+                    nn.init.constant_(m.bn2.weight, 0)
+
+    def _make_layer(self, block, planes, blocks, stride=1, dilate=False):
+        norm_layer = self._norm_layer
+        downsample = None
+        previous_dilation = self.dilation
+        if dilate:
+            self.dilation *= stride
+            stride = 1
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                conv1x1(self.inplanes, planes * block.expansion, stride),
+                norm_layer(planes * block.expansion),
+            )
+
+        layers = []
+        layers.append(
+            block(
+                self.inplanes,
+                planes,
+                stride,
+                downsample,
+                self.groups,
+                self.base_width,
+                previous_dilation,
+                norm_layer,
+            )
+        )
+        self.inplanes = planes * block.expansion
+        for _ in range(1, blocks):
+            layers.append(
+                block(
+                    self.inplanes,
+                    planes,
+                    groups=self.groups,
+                    base_width=self.base_width,
+                    dilation=self.dilation,
+                    norm_layer=norm_layer,
+                )
+            )
+
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu(x)
+        x = self.maxpool(x)
+
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.layer4(x)
+
+        x = self.avgpool(x)
+        x = torch.flatten(x, 1)
+        x = self.fc(x)
+
+        return x

models/utils.py

@@ -0,0 +1,4 @@
+try:
+    from torch.hub import load_state_dict_from_url  # noqa: F401
+except ImportError:
+    from torch.utils.model_zoo import load_url as load_state_dict_from_url  # noqa: F401

requirements-hf.txt

@@ -0,0 +1,9 @@
+# Các thư viện dành riêng cho việc triển khai (deploy)
+gradio==4.36.1
+opencv-python-headless==4.9.0.80
+torch
+torchvision
+numpy
+
+# Đường dẫn tới thư viện rmn (Nếu tải trực tiếp lên HF thì nó sẽ cài từ pypi)
+rmn==3.1.2

rmn/__init__.py

@@ -0,0 +1,340 @@
+import os
+
+import cv2
+import numpy as np
+import torch
+from torchvision.transforms import transforms
+
+from models import resmasking_dropout1
+
+
+def show(img, name="disp", width=1000):
+    """
+    name: name of window, should be name of img
+    img: source of img, should in type ndarray
+    """
+    cv2.namedWindow(name, cv2.WINDOW_GUI_NORMAL)
+    cv2.resizeWindow(name, width, 1000)
+    cv2.imshow(name, img)
+    cv2.waitKey(0)
+    cv2.destroyAllWindows()
+
+
+checkpoint_url = "https://github.com/phamquiluan/ResidualMaskingNetwork/releases/download/v0.0.1/Z_resmasking_dropout1_rot30_2019Nov30_13.32"
+local_checkpoint_path = "pretrained_ckpt"
+
+prototxt_url = "https://github.com/phamquiluan/ResidualMaskingNetwork/releases/download/v0.0.1/deploy.prototxt.txt"
+local_prototxt_path = "deploy.prototxt.txt"
+
+ssd_checkpoint_url = "https://github.com/phamquiluan/ResidualMaskingNetwork/releases/download/v0.0.1/res10_300x300_ssd_iter_140000.caffemodel"
+local_ssd_checkpoint_path = "res10_300x300_ssd_iter_140000.caffemodel"
+
+
+def download_checkpoint(remote_url, local_path):
+    import requests
+    from tqdm import tqdm
+
+    response = requests.get(remote_url, stream=True)
+    total_size_in_bytes = int(response.headers.get("content-length", 0))
+    block_size = 1024  # 1 Kibibyte
+
+    progress_bar = tqdm(
+        desc=f"Downloading {local_path}..",
+        total=total_size_in_bytes,
+        unit="iB",
+        unit_scale=True,
+    )
+
+    with open(local_path, "wb") as ref:
+        for data in response.iter_content(block_size):
+            progress_bar.update(len(data))
+            ref.write(data)
+
+    progress_bar.close()
+    if total_size_in_bytes != 0 and progress_bar.n != total_size_in_bytes:
+        print("ERROR, something went wrong")
+
+
+for remote_path, local_path in [
+    (checkpoint_url, local_checkpoint_path),
+    (prototxt_url, local_prototxt_path),
+    (ssd_checkpoint_url, local_ssd_checkpoint_path),
+]:
+    if not os.path.exists(local_path):
+        print(f"{local_path} does not exists!")
+        download_checkpoint(remote_url=remote_path, local_path=local_path)
+
+
+def ensure_color(image):
+    if len(image.shape) == 2:
+        return np.dstack([image] * 3)
+    elif image.shape[2] == 1:
+        return np.dstack([image] * 3)
+    return image
+
+
+def ensure_gray(image):
+    try:
+        image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
+    except cv2.error:
+        pass
+    return image
+
+
+def get_ssd_face_detector():
+    ssd_face_detector = cv2.dnn.readNetFromCaffe(
+        prototxt=local_prototxt_path,
+        caffeModel=local_ssd_checkpoint_path,
+    )
+    return ssd_face_detector
+
+
+transform = transforms.Compose(
+    transforms=[transforms.ToPILImage(), transforms.ToTensor()]
+)
+
+FER_2013_EMO_DICT = {
+    0: "angry",
+    1: "disgust",
+    2: "fear",
+    3: "happy",
+    4: "sad",
+    5: "surprise",
+    6: "neutral",
+}
+
+is_cuda = torch.cuda.is_available()
+
+image_size = (224, 224)
+
+
+def get_emo_model():
+    emo_model = resmasking_dropout1(in_channels=3, num_classes=7)
+    if is_cuda:
+        emo_model.cuda(0)
+    state = torch.load(local_checkpoint_path, map_location="cpu")
+    emo_model.load_state_dict(state["net"])
+    emo_model.eval()
+    return emo_model
+
+
+def convert_to_square(xmin, ymin, xmax, ymax):
+    # convert to square location
+    center_x = (xmin + xmax) // 2
+    center_y = (ymin + ymax) // 2
+
+    square_length = ((xmax - xmin) + (ymax - ymin)) // 2 // 2
+    square_length *= 1.1
+
+    xmin = int(center_x - square_length)
+    ymin = int(center_y - square_length)
+    xmax = int(center_x + square_length)
+    ymax = int(center_y + square_length)
+    return xmin, ymin, xmax, ymax
+
+
+class RMN:
+    def __init__(self, face_detector=True):
+        if face_detector:
+            self.face_detector = get_ssd_face_detector()
+        self.emo_model = get_emo_model()
+
+    @torch.no_grad()
+    def detect_emotion_for_single_face_image(self, face_image):
+        """
+        Params:
+        -----------
+        face_image : np.ndarray
+            a cropped face image
+
+        Return:
+        -----------
+        emo_label : str
+            dominant emotion label
+
+        emo_proba : float
+            dominant emotion proba
+
+        proba_list : list
+            all emotion label and their proba
+        """
+        assert isinstance(face_image, np.ndarray)
+        face_image = ensure_color(face_image)
+        face_image = cv2.resize(face_image, image_size)
+
+        face_image = transform(face_image)
+        if is_cuda:
+            face_image = face_image.cuda(0)
+
+        face_image = torch.unsqueeze(face_image, dim=0)
+
+        output = torch.squeeze(self.emo_model(face_image), 0)
+        proba = torch.softmax(output, 0)
+
+        # get dominant emotion
+        emo_proba, emo_idx = torch.max(proba, dim=0)
+        emo_idx = emo_idx.item()
+        emo_proba = emo_proba.item()
+        emo_label = FER_2013_EMO_DICT[emo_idx]
+
+        # get proba for each emotion
+        proba = proba.tolist()
+        proba_list = []
+        for emo_idx, emo_name in FER_2013_EMO_DICT.items():
+            proba_list.append({emo_name: proba[emo_idx]})
+
+        return emo_label, emo_proba, proba_list
+
+    @torch.no_grad()
+    def video_demo(self):
+        vid = cv2.VideoCapture(0)
+
+        while True:
+            ret, frame = vid.read()
+            if frame is None or ret is not True:
+                continue
+
+            try:
+                frame = np.fliplr(frame).astype(np.uint8)
+
+                results = self.detect_emotion_for_single_frame(frame)
+                frame = self.draw(frame, results)
+
+                cv2.rectangle(frame, (1, 1), (220, 25), (223, 128, 255), cv2.FILLED)
+                cv2.putText(
+                    frame,
+                    f"press q to exit",
+                    (20, 20),
+                    cv2.FONT_HERSHEY_SIMPLEX,
+                    0.8,
+                    (0, 0, 0),
+                    2,
+                )
+                cv2.imshow("disp", frame)
+                if cv2.waitKey(1) == ord("q"):
+                    break
+
+            except Exception as err:
+                print(err)
+                continue
+
+        cv2.destroyAllWindows()
+
+    @staticmethod
+    def draw(frame, results):
+        """
+        Params:
+        ---------
+        frame : np.ndarray
+
+        results : list of dict.keys('xmin', 'xmax', 'ymin', 'ymax', 'emo_label', 'emo_proba')
+
+        Returns:
+        ---------
+        frame : np.ndarray
+        """
+        for r in results:
+            xmin = r["xmin"]
+            xmax = r["xmax"]
+            ymin = r["ymin"]
+            ymax = r["ymax"]
+            emo_label = r["emo_label"]
+            emo_proba = r["emo_proba"]
+
+            label_size, base_line = cv2.getTextSize(
+                f"{emo_label}: 000", cv2.FONT_HERSHEY_SIMPLEX, 0.8, 2
+            )
+
+            # draw face
+            cv2.rectangle(frame, (xmin, ymin), (xmax, ymax), (179, 255, 179), 2)
+
+            cv2.rectangle(
+                frame,
+                (xmax, ymin + 1 - label_size[1]),
+                (xmax + label_size[0], ymin + 1 + base_line),
+                (223, 128, 255),
+                cv2.FILLED,
+            )
+            cv2.putText(
+                frame,
+                f"{emo_label} {int(emo_proba * 100)}",
+                (xmax, ymin + 1),
+                cv2.FONT_HERSHEY_SIMPLEX,
+                0.8,
+                (0, 0, 0),
+                2,
+            )
+
+        return frame
+
+    def detect_faces(self, frame):
+        h, w = frame.shape[:2]
+        blob = cv2.dnn.blobFromImage(
+            cv2.resize(frame, (300, 300)),
+            1.0,
+            (300, 300),
+            (104.0, 177.0, 123.0),
+            False,
+            False,
+        )
+        self.face_detector.setInput(blob)
+        faces = self.face_detector.forward()
+
+        face_results = []
+        for i in range(0, faces.shape[2]):
+            confidence = faces[0, 0, i, 2]
+            if confidence < 0.5:
+                continue
+            xmin, ymin, xmax, ymax = (
+                faces[0, 0, i, 3:7] * np.array([w, h, w, h])
+            ).astype("int")
+            xmin, ymin, xmax, ymax = convert_to_square(xmin, ymin, xmax, ymax)
+            if xmax <= xmin or ymax <= ymin:
+                continue
+
+            face_results.append(
+                {
+                    "xmin": xmin,
+                    "ymin": ymin,
+                    "xmax": xmax,
+                    "ymax": ymax,
+                }
+            )
+        return face_results
+
+    @torch.no_grad()
+    def detect_emotion_for_single_frame(self, frame):
+        gray = ensure_gray(frame)
+
+        results = []
+        face_results = self.detect_faces(frame)
+        print(f"num faces: {len(face_results)}")
+
+        for face in face_results:
+            xmin = face["xmin"]
+            ymin = face["ymin"]
+            xmax = face["xmax"]
+            ymax = face["ymax"]
+
+            face_image = gray[ymin:ymax, xmin:xmax]
+
+            if face_image.shape[0] < 10 or face_image.shape[1] < 10:
+                continue
+            (
+                emo_label,
+                emo_proba,
+                proba_list,
+            ) = self.detect_emotion_for_single_face_image(face_image)
+
+            results.append(
+                {
+                    "xmin": xmin,
+                    "ymin": ymin,
+                    "xmax": xmax,
+                    "ymax": ymax,
+                    "emo_label": emo_label,
+                    "emo_proba": emo_proba,
+                    "proba_list": proba_list,
+                }
+            )
+        return results

train.py

@@ -0,0 +1,533 @@
+import datetime
+import os
+import random
+
+import cv2
+import imgaug
+import matplotlib.pyplot as plt
+import numpy as np
+import pandas as pd
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.optim.lr_scheduler import ReduceLROnPlateau
+from torch.utils.data import DataLoader, Dataset
+from torch.utils.tensorboard import SummaryWriter
+from torchvision import transforms
+from tqdm import tqdm
+from sklearn.metrics import ConfusionMatrixDisplay, confusion_matrix
+
+from models import resmasking_dropout1
+from utils.augmenters.augment import seg
+from utils.metrics.metrics import accuracy
+from utils.radam import RAdam
+
+
+FER_2013_EMO_DICT = {
+    0: "angry",
+    1: "disgust",
+    2: "fear",
+    3: "happy",
+    4: "sad",
+    5: "surprise",
+    6: "neutral",
+}
+
+
+class FER2013Dataset(Dataset):
+    def __init__(
+        self,
+        stage,
+        data_path,
+        image_size,
+        use_tta=False,
+        tta_size=48,
+    ):
+        self._stage = stage
+        self._data_path = data_path
+        self._image_size = (image_size, image_size)
+        self._use_tta = use_tta
+        self._tta_size = tta_size
+
+        csv_path = os.path.join(self._data_path, f"{self._stage}.csv")
+        self._data = pd.read_csv(csv_path)
+
+        self._pixels = self._data["pixels"].tolist()
+        self._targets = self._data["emotion"].tolist()
+
+        self._transform = transforms.Compose(
+            [
+                transforms.ToPILImage(),
+                transforms.ToTensor(),
+            ]
+        )
+
+    def is_tta(self):
+        return self._use_tta
+
+    def __len__(self):
+        return len(self._pixels)
+
+    def __getitem__(self, idx):
+        pixels = self._pixels[idx]
+        pixels = list(map(int, pixels.split(" ")))
+        image = np.asarray(pixels).reshape(48, 48).astype(np.uint8)
+
+        image = cv2.resize(image, self._image_size)
+        image = np.dstack([image] * 3)
+
+        if self._stage == "train":
+            image = seg(image=image)
+
+        target = int(self._targets[idx])
+
+        if self._stage == "test" and self._use_tta:
+            images = [seg(image=image) for _ in range(self._tta_size)]
+            images = [self._transform(tta_image) for tta_image in images]
+            return images, target
+
+        image = self._transform(image)
+        return image, target
+
+
+def set_seed(seed):
+    random.seed(seed)
+    np.random.seed(seed)
+    imgaug.seed(seed)
+
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.benchmark = False
+
+
+def resolve_device(device_name):
+    if device_name.startswith("cuda") and torch.cuda.is_available():
+        return torch.device(device_name)
+    return torch.device("cpu")
+
+
+def build_datasets(data_path, image_size, use_tta, tta_size):
+    train_set = FER2013Dataset(
+        stage="train",
+        data_path=data_path,
+        image_size=image_size,
+    )
+    val_set = FER2013Dataset(
+        stage="val",
+        data_path=data_path,
+        image_size=image_size,
+    )
+    test_set = FER2013Dataset(
+        stage="test",
+        data_path=data_path,
+        image_size=image_size,
+        use_tta=use_tta,
+        tta_size=tta_size,
+    )
+    return train_set, val_set, test_set
+
+
+def build_dataloaders(train_set, val_set, test_set, batch_size, num_workers, use_tta):
+    pin_memory = torch.cuda.is_available()
+
+    train_loader = DataLoader(
+        train_set,
+        batch_size=batch_size,
+        num_workers=num_workers,
+        pin_memory=pin_memory,
+        shuffle=True,
+    )
+    val_loader = DataLoader(
+        val_set,
+        batch_size=batch_size,
+        num_workers=num_workers,
+        pin_memory=pin_memory,
+        shuffle=False,
+    )
+
+    test_loader = None
+    if not use_tta:
+        test_loader = DataLoader(
+            test_set,
+            batch_size=1,
+            num_workers=num_workers,
+            pin_memory=pin_memory,
+            shuffle=False,
+        )
+
+    return train_loader, val_loader, test_loader
+
+
+def visualize_dataset(data_path):
+    stages = ("train", "val", "test")
+    class_ids = list(FER_2013_EMO_DICT.keys())
+    class_labels = [FER_2013_EMO_DICT[class_id] for class_id in class_ids]
+    counts = pd.Series(0, index=class_ids, dtype=np.int64)
+
+    for stage in stages:
+        csv_path = os.path.join(data_path, f"{stage}.csv")
+        df = pd.read_csv(csv_path)
+        stage_counts = df["emotion"].value_counts()
+        counts = counts.add(stage_counts, fill_value=0).astype(int)
+
+    plt.figure(figsize=(8, 4))
+    plt.bar(class_labels, counts.values)
+    plt.title("Class Distribution")
+    plt.xlabel("Emotion")
+    plt.ylabel("Samples")
+    plt.xticks(rotation=20)
+    plt.tight_layout()
+    plt.show()
+
+
+def train_one_epoch(model, loader, criterion, optimizer, device):
+    model.train()
+    total_loss = 0.0
+    total_acc = 0.0
+
+    for images, targets in tqdm(loader, total=len(loader), leave=False):
+        images = images.to(device, non_blocking=True)
+        targets = targets.to(device, non_blocking=True)
+
+        outputs = model(images)
+        loss = criterion(outputs, targets)
+        acc = accuracy(outputs, targets)[0]
+
+        total_loss += loss.item()
+        total_acc += acc.item()
+
+        optimizer.zero_grad()
+        loss.backward()
+        optimizer.step()
+
+    num_batches = len(loader)
+    return total_loss / num_batches, total_acc / num_batches
+
+
+def validate_one_epoch(model, loader, criterion, device):
+    model.eval()
+    total_loss = 0.0
+    total_acc = 0.0
+
+    with torch.no_grad():
+        for images, targets in tqdm(loader, total=len(loader), leave=False):
+            images = images.to(device, non_blocking=True)
+            targets = targets.to(device, non_blocking=True)
+
+            outputs = model(images)
+            loss = criterion(outputs, targets)
+            acc = accuracy(outputs, targets)[0]
+
+            total_loss += loss.item()
+            total_acc += acc.item()
+
+    num_batches = len(loader)
+    return total_loss / num_batches, total_acc / num_batches
+
+
+def eval_test_without_tta(model, loader, device):
+    model.eval()
+    total_acc = 0.0
+
+    with torch.no_grad():
+        for images, targets in tqdm(loader, total=len(loader), leave=False):
+            images = images.to(device, non_blocking=True)
+            targets = targets.to(device, non_blocking=True)
+
+            outputs = model(images)
+            acc = accuracy(outputs, targets)[0]
+            total_acc += acc.item()
+
+    return total_acc / len(loader)
+
+
+def eval_test_with_tta(model, dataset, device):
+    model.eval()
+    total_acc = 0.0
+
+    with torch.no_grad():
+        for idx in tqdm(range(len(dataset)), total=len(dataset), leave=False):
+            images, target = dataset[idx]
+
+            images = torch.stack(images, dim=0).to(device, non_blocking=True)
+            target = torch.LongTensor([target]).to(device, non_blocking=True)
+
+            outputs = model(images)
+            outputs = F.softmax(outputs, dim=1)
+            outputs = torch.sum(outputs, dim=0, keepdim=True)
+
+            acc = accuracy(outputs, target)[0]
+            total_acc += acc.item()
+
+    return total_acc / len(dataset)
+
+
+def save_checkpoint(path, model, train_params, metrics):
+    state = {
+        "net": model.state_dict(),
+        "config": train_params,
+        **metrics,
+    }
+    torch.save(state, path)
+
+
+def train(
+    data_path="data",
+    image_size=224,
+    lr=1e-4,
+    weight_decay=1e-3,
+    batch_size=48,
+    num_workers=2,
+    device_name="cuda:0",
+    max_epoch_num=50,
+    max_plateau_count=8,
+    plateau_patience=2,
+    log_dir="log",
+    checkpoint_dir="checkpoint",
+    seed=1234,
+    use_tta=True,
+    tta_size=10,
+):
+    train_params = {
+        "data_path": data_path,
+        "image_size": image_size,
+        "lr": lr,
+        "weight_decay": weight_decay,
+        "batch_size": batch_size,
+        "num_workers": num_workers,
+        "device_name": device_name,
+        "max_epoch_num": max_epoch_num,
+        "max_plateau_count": max_plateau_count,
+        "plateau_patience": plateau_patience,
+        "log_dir": log_dir,
+        "checkpoint_dir": checkpoint_dir,
+        "seed": seed,
+        "use_tta": use_tta,
+        "tta_size": tta_size,
+    }
+
+    set_seed(seed)
+    device = resolve_device(device_name)
+
+    cwd = os.getcwd()
+    log_root = os.path.join(cwd, log_dir)
+    ckpt_root = os.path.join(cwd, checkpoint_dir)
+    os.makedirs(log_root, exist_ok=True)
+    os.makedirs(ckpt_root, exist_ok=True)
+
+    start_time = datetime.datetime.now().replace(microsecond=0)
+    run_name = f"resmasking_dropout1_train_{start_time.strftime('%d%m%Y_%H%M%S')}"
+
+    writer = SummaryWriter(os.path.join(log_root, run_name))
+    checkpoint_path = os.path.join(ckpt_root, f"{run_name}.pt")
+
+    print("Start training")
+    print(train_params)
+    print(f"Device: {device}")
+
+    model = resmasking_dropout1(
+        in_channels=3,
+        num_classes=7,
+    ).to(device)
+
+    train_set, val_set, test_set = build_datasets(
+        data_path=data_path,
+        image_size=image_size,
+        use_tta=use_tta,
+        tta_size=tta_size,
+    )
+    train_loader, val_loader, test_loader = build_dataloaders(
+        train_set=train_set,
+        val_set=val_set,
+        test_set=test_set,
+        batch_size=batch_size,
+        num_workers=num_workers,
+        use_tta=use_tta,
+    )
+
+    criterion = nn.CrossEntropyLoss().to(device)
+    optimizer = RAdam(
+        params=model.parameters(),
+        lr=lr,
+        weight_decay=weight_decay,
+    )
+    scheduler = ReduceLROnPlateau(
+        optimizer,
+        patience=plateau_patience,
+        min_lr=1e-6,
+    )
+
+    train_losses = []
+    val_losses = []
+    train_accs = []
+    val_accs = []
+
+    best_val_loss = 1e9
+    best_val_acc = 0.0
+    best_train_loss = 1e9
+    best_train_acc = 0.0
+
+    plateau_count = 0
+    current_epoch = 0
+
+    while current_epoch < max_epoch_num and plateau_count <= max_plateau_count:
+        current_epoch += 1
+
+        train_loss, train_acc = train_one_epoch(
+            model=model,
+            loader=train_loader,
+            criterion=criterion,
+            optimizer=optimizer,
+            device=device,
+        )
+        val_loss, val_acc = validate_one_epoch(
+            model=model,
+            loader=val_loader,
+            criterion=criterion,
+            device=device,
+        )
+
+        train_losses.append(train_loss)
+        val_losses.append(val_loss)
+        train_accs.append(train_acc)
+        val_accs.append(val_acc)
+
+        if current_epoch == 1 or val_acc > best_val_acc:
+            plateau_count = 0
+            best_val_acc = val_acc
+            best_val_loss = val_loss
+            best_train_acc = train_acc
+            best_train_loss = train_loss
+
+            save_checkpoint(
+                path=checkpoint_path,
+                model=model,
+                train_params=train_params,
+                metrics={
+                    "best_val_loss": best_val_loss,
+                    "best_val_acc": best_val_acc,
+                    "best_train_loss": best_train_loss,
+                    "best_train_acc": best_train_acc,
+                    "train_losses": train_losses,
+                    "val_losses": val_losses,
+                    "train_accs": train_accs,
+                    "val_accs": val_accs,
+                    "current_epoch": current_epoch,
+                },
+            )
+        else:
+            plateau_count += 1
+
+        scheduler.step(100.0 - val_acc)
+
+        writer.add_scalar("Accuracy/Train", train_acc, current_epoch)
+        writer.add_scalar("Accuracy/Val", val_acc, current_epoch)
+        writer.add_scalar("Loss/Train", train_loss, current_epoch)
+        writer.add_scalar("Loss/Val", val_loss, current_epoch)
+
+        consume_time = str(datetime.datetime.now() - start_time)
+        message = (
+            f"E{current_epoch:03d}  "
+            f"{train_loss:.3f}/{val_loss:.3f}/{best_val_loss:.3f} "
+            f"{train_acc:.3f}/{val_acc:.3f}/{best_val_acc:.3f} "
+            f"| p{plateau_count:02d}  Time {consume_time[:-7]}"
+        )
+        print(message)
+
+    best_state = torch.load(checkpoint_path, map_location=device)
+    model.load_state_dict(best_state["net"])
+
+    if use_tta:
+        test_acc = eval_test_with_tta(model, test_set, device)
+    else:
+        test_acc = eval_test_without_tta(model, test_loader, device)
+
+    best_state["test_acc"] = test_acc
+    torch.save(best_state, checkpoint_path)
+
+    consume_time = str(datetime.datetime.now() - start_time)
+    writer.add_text(
+        "Summary",
+        f"Converged after {current_epoch} epochs, consume {consume_time[:-7]}",
+    )
+    writer.add_text("Results", f"Best validation accuracy: {best_val_acc:.3f}")
+    writer.add_text("Results", f"Best training accuracy: {best_train_acc:.3f}")
+    writer.add_text("Results", f"Test accuracy: {test_acc:.3f}")
+    writer.close()
+
+    print(f"Best checkpoint saved at: {checkpoint_path}")
+    print(f"Test accuracy: {test_acc:.3f}")
+
+
+def eval_test_with_tta_conf_matrix(checkpoint_path):
+
+    checkpoint = torch.load(checkpoint_path)
+    config = checkpoint.get("config")
+
+    data_path = config.get("data_path")
+    image_size = config.get("image_size")
+    device_name = config.get("device_name")
+    use_tta = config.get("use_tta")
+    tta_size = config.get("tta_size")
+
+    if not use_tta:
+        raise Exception
+
+    device = resolve_device(device_name)
+
+    model = resmasking_dropout1(
+        in_channels=3,
+        num_classes=7,
+    ).to(device)
+    model.load_state_dict(checkpoint["net"])
+    model.eval()
+
+    _, _, test_set = build_datasets(
+        data_path=data_path,
+        image_size=image_size,
+        use_tta=True,
+        tta_size=tta_size,
+    )
+
+    y_true = []
+    y_pred = []
+
+    total_acc = 0.0
+    with torch.no_grad():
+        for idx in tqdm(range(len(test_set)), total=len(test_set), leave=False):
+            images, target = test_set[idx]
+            images = torch.stack(images, dim=0).to(device, non_blocking=True)
+            target = torch.LongTensor([target]).to(device, non_blocking=True)
+
+            outputs = model(images)
+            outputs = F.softmax(outputs, dim=1)
+            outputs = torch.sum(outputs, dim=0, keepdim=True)
+            pred = torch.argmax(outputs, dim=1).item()
+
+            acc = accuracy(outputs, target)[0]
+            total_acc += acc.item()
+
+            y_true.append(int(target))
+            y_pred.append(int(pred))
+
+    class_ids = list(FER_2013_EMO_DICT.keys())
+    class_labels = [FER_2013_EMO_DICT[class_id] for class_id in class_ids]
+    cm = confusion_matrix(y_true, y_pred, labels=class_ids)
+
+    acc = total_acc / len(test_set) if len(test_set) > 0 else 0.0
+    print(f"Test accuracy: {acc:.3f}")
+
+    fig, ax = plt.subplots(figsize=(8, 8))
+    disp = ConfusionMatrixDisplay(confusion_matrix=cm, display_labels=class_labels)
+    disp.plot(ax=ax, cmap="Blues", colorbar=True, values_format="d")
+    ax.set_title("Confusion Matrix")
+    plt.xticks(rotation=45, ha="right")
+    plt.tight_layout()
+    plt.show()
+
+
+if __name__ == "__main__":
+    # visualize_dataset(data_path="data")
+    # eval_test_with_tta_conf_matrix("checkpoint/branch1.pt")
+    train()

utils/augmenters/augment.py

@@ -0,0 +1,8 @@
+from imgaug import augmenters as iaa
+
+seg = iaa.Sequential(
+    [
+        iaa.Fliplr(p=0.5),
+        iaa.Affine(rotate=(-30, 30)),
+    ]
+)

utils/metrics/metrics.py

@@ -0,0 +1,10 @@
+import torch
+
+
+def accuracy(output, target):
+    with torch.no_grad():
+        batch_size = target.size(0)
+        pred = torch.argmax(output, dim=1)
+        correct = pred.eq(target).float().sum(0)
+        acc = correct * 100 / batch_size
+    return [acc]

utils/radam.py

@@ -0,0 +1,94 @@
+import math
+
+import torch
+from torch.optim.optimizer import Optimizer
+
+
+class RAdam(Optimizer):
+    def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8, weight_decay=0):
+        defaults = dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay)
+        self.buffer = [[None, None, None] for ind in range(10)]
+        super(RAdam, self).__init__(params, defaults)
+
+    def __setstate__(self, state):
+        super(RAdam, self).__setstate__(state)
+
+    def step(self, closure=None):
+
+        loss = None
+        if closure is not None:
+            loss = closure()
+
+        for group in self.param_groups:
+
+            for p in group["params"]:
+                if p.grad is None:
+                    continue
+                grad = p.grad.data.float()
+                if grad.is_sparse:
+                    raise RuntimeError("RAdam does not support sparse gradients")
+
+                p_data_fp32 = p.data.float()
+
+                state = self.state[p]
+
+                if len(state) == 0:
+                    state["step"] = 0
+                    state["exp_avg"] = torch.zeros_like(p_data_fp32)
+                    state["exp_avg_sq"] = torch.zeros_like(p_data_fp32)
+                else:
+                    state["exp_avg"] = state["exp_avg"].type_as(p_data_fp32)
+                    state["exp_avg_sq"] = state["exp_avg_sq"].type_as(p_data_fp32)
+
+                exp_avg, exp_avg_sq = state["exp_avg"], state["exp_avg_sq"]
+                beta1, beta2 = group["betas"]
+
+                exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)
+                exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1)
+
+                state["step"] += 1
+                buffered = self.buffer[int(state["step"] % 10)]
+                if state["step"] == buffered[0]:
+                    N_sma, step_size = buffered[1], buffered[2]
+                else:
+                    buffered[0] = state["step"]
+                    beta2_t = beta2 ** state["step"]
+                    N_sma_max = 2 / (1 - beta2) - 1
+                    N_sma = N_sma_max - 2 * state["step"] * beta2_t / (1 - beta2_t)
+                    buffered[1] = N_sma
+
+                    # more conservative since it's an approximated value
+                    if N_sma >= 5:
+                        step_size = (
+                            group["lr"]
+                            * math.sqrt(
+                                (1 - beta2_t)
+                                * (N_sma - 4)
+                                / (N_sma_max - 4)
+                                * (N_sma - 2)
+                                / N_sma
+                                * N_sma_max
+                                / (N_sma_max - 2)
+                            )
+                            / (1 - beta1 ** state["step"])
+                        )
+                    else:
+                        step_size = group["lr"] / (1 - beta1 ** state["step"])
+                    buffered[2] = step_size
+
+                if group["weight_decay"] != 0:
+                    p_data_fp32.add_(
+                        p_data_fp32, alpha=-group["weight_decay"] * group["lr"]
+                    )
+
+                # more conservative since it's an approximated value
+                if N_sma >= 5:
+                    denom = exp_avg_sq.sqrt().add_(group["eps"])
+                    p_data_fp32.addcdiv_(exp_avg, denom, value=-step_size)
+                else:
+                    p_data_fp32.add_(exp_avg, alpha=-step_size)
+
+                p.data.copy_(p_data_fp32)
+
+        return loss
+