Summary

.gitignore

@@ -168,4 +168,5 @@ dmypy.json
 .pyre/
 
 # Custom
-*.pth
+*.pth
+*.pt

app.css

@@ -20,6 +20,14 @@ div.dl2 img {
     max-height: 200px;
 }
 
+div.dl5 {
+    max-height: 200px;
+}
+
+div.dl5 img {
+    max-height: 200px;
+}
+
 div.video1 div.video-container {
     height: 500px;
 }
@@ -32,6 +40,10 @@ div.video3 {
     height: 200px;
 }
 
+div.video4 {
+    height: 200px;
+}
+
 div.stat {
     height: 286px;
 }

app.py

@@ -18,6 +18,7 @@ from app.app_utils import preprocess_image_and_predict, preprocess_video_and_pre
 def clear_static_info():
     return (
         gr.Image(value=None, type="pil"),
+        gr.Image(value=None, scale=1, elem_classes="dl5"),
         gr.Image(value=None, scale=1, elem_classes="dl2"),
         gr.Label(value=None, num_top_classes=3, scale=1, elem_classes="dl3"),
     )
@@ -27,6 +28,7 @@ def clear_dynamic_info():
         gr.Video(value=None),
         gr.Video(value=None),
         gr.Video(value=None),
+        gr.Video(value=None),
         gr.Plot(value=None),
     )
 
@@ -45,8 +47,9 @@ with gr.Blocks(css="app.css") as demo:
                     )
             with gr.Column(scale=2, elem_classes="dl4"):
                 with gr.Row():
-                    output_video = gr.Video(label="Original video", scale=2, elem_classes="video2")
+                    output_video = gr.Video(label="Original video", scale=1, elem_classes="video2")
                     output_face = gr.Video(label="Pre-processed video", scale=1, elem_classes="video3")
+                    output_heatmaps = gr.Video(label="Heatmaps", scale=1, elem_classes="video4")
                 output_statistics = gr.Plot(label="Statistics of emotions", elem_classes="stat")
         gr.Examples(
             ["videos/video1.mp4",
@@ -58,7 +61,7 @@ with gr.Blocks(css="app.css") as demo:
         gr.Markdown(value=DESCRIPTION_STATIC)
         with gr.Row():
             with gr.Column(scale=2, elem_classes="dl1"):
-                input_image = gr.Image(type="pil")
+                input_image = gr.Image(label="Original image", type="pil")
                 with gr.Row():
                     clear_btn = gr.Button(
                         value="Clear", interactive=True, scale=1, elem_classes="clear"
@@ -67,7 +70,9 @@ with gr.Blocks(css="app.css") as demo:
                         value="Submit", interactive=True, scale=1, elem_classes="submit"
                     )
             with gr.Column(scale=1, elem_classes="dl4"):
-                output_image = gr.Image(scale=1, elem_classes="dl2")
+                with gr.Row():
+                    output_image = gr.Image(label="Face", scale=1, elem_classes="dl5")
+                    output_heatmap = gr.Image(label="Heatmap", scale=1, elem_classes="dl2")
                 output_label = gr.Label(num_top_classes=3, scale=1, elem_classes="dl3")
         gr.Examples(
             [
@@ -87,13 +92,13 @@ with gr.Blocks(css="app.css") as demo:
     submit.click(
         fn=preprocess_image_and_predict,
         inputs=[input_image],
-        outputs=[output_image, output_label],
+        outputs=[output_image, output_heatmap, output_label],
         queue=True,
     )
     clear_btn.click(
         fn=clear_static_info,
         inputs=[],
-        outputs=[input_image, output_image, output_label],
+        outputs=[input_image, output_image, output_heatmap, output_label],
         queue=True,
     )
 
@@ -102,7 +107,8 @@ with gr.Blocks(css="app.css") as demo:
         inputs=input_video,
         outputs=[
             output_video,
-            output_face, 
+            output_face,
+            output_heatmaps, 
             output_statistics
         ],
         queue=True,
@@ -113,7 +119,8 @@ with gr.Blocks(css="app.css") as demo:
         outputs=[
             input_video,
             output_video,
-            output_face, 
+            output_face,
+            output_heatmaps, 
             output_statistics
         ],
         queue=True,

app/app_utils.py

@@ -10,9 +10,10 @@ import numpy as np
 import mediapipe as mp
 from PIL import Image
 import cv2
+from pytorch_grad_cam.utils.image import show_cam_on_image
 
 # Importing necessary components for the Gradio app
-from app.model import pth_model_static, pth_model_dynamic, pth_processing
+from app.model import pth_model_static, pth_model_dynamic, cam, pth_processing
 from app.face_utils import get_box, display_info
 from app.config import DICT_EMO, config_data
 from app.plot import statistics_plot
@@ -49,8 +50,13 @@ def preprocess_image_and_predict(inp):
                     .numpy()[0]
                 )
                 confidences = {DICT_EMO[i]: float(prediction[i]) for i in range(7)}
+                grayscale_cam = cam(input_tensor=cur_face_n)
+                grayscale_cam = grayscale_cam[0, :]
+                cur_face_hm = cv2.resize(cur_face,(224,224))
+                cur_face_hm = np.float32(cur_face_hm) / 255
+                heatmap = show_cam_on_image(cur_face_hm, grayscale_cam, use_rgb=True)
 
-    return cur_face, confidences
+    return cur_face, heatmap, confidences
 
 
 def preprocess_video_and_predict(video):
@@ -60,14 +66,20 @@ def preprocess_video_and_predict(video):
     h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
     fps = np.round(cap.get(cv2.CAP_PROP_FPS))
 
-    path_save_video = 'result.mp4'
-    vid_writer = cv2.VideoWriter(path_save_video, cv2.VideoWriter_fourcc(*'mp4v'), fps, (224, 224))
+    path_save_video_face = 'result_face.mp4'
+    vid_writer_face = cv2.VideoWriter(path_save_video_face, cv2.VideoWriter_fourcc(*'mp4v'), fps, (224, 224))
+
+    path_save_video_hm = 'result_hm.mp4'
+    vid_writer_hm = cv2.VideoWriter(path_save_video_hm, cv2.VideoWriter_fourcc(*'mp4v'), fps, (224, 224))
 
     lstm_features = []
     count_frame = 1
+    count_face = 0
     probs = []
     frames = []
-    last_output = None 
+    last_output = None
+    last_heatmap = None 
+    cur_face = None
 
     with mp_face_mesh.FaceMesh(
     max_num_faces=1,
@@ -90,9 +102,16 @@ def preprocess_video_and_predict(video):
                     startX, startY, endX, endY  = get_box(fl, w, h)
                     cur_face = frame_copy[startY:endY, startX: endX]
 
-                    if (count_frame-1)%config_data.FRAME_DOWNSAMPLING == 0:
+                    if count_face%config_data.FRAME_DOWNSAMPLING == 0:
                         cur_face_copy = pth_processing(Image.fromarray(cur_face))
                         features = torch.nn.functional.relu(pth_model_static.extract_features(cur_face_copy)).detach().numpy()
+
+                        grayscale_cam = cam(input_tensor=cur_face_copy)
+                        grayscale_cam = grayscale_cam[0, :]
+                        cur_face_hm = cv2.resize(cur_face,(224,224), interpolation = cv2.INTER_AREA)
+                        cur_face_hm = np.float32(cur_face_hm) / 255
+                        heatmap = show_cam_on_image(cur_face_hm, grayscale_cam, use_rgb=False)
+                        last_heatmap = heatmap
         
                         if len(lstm_features) == 0:
                             lstm_features = [features]*10
@@ -103,32 +122,48 @@ def preprocess_video_and_predict(video):
                         lstm_f = torch.unsqueeze(lstm_f, 0)
                         output = pth_model_dynamic(lstm_f).detach().numpy()
                         last_output = output
+
+                        if count_face == 0:
+                            count_face += 1
+
                     else:
                         if last_output is not None:
                             output = last_output
+                            heatmap = last_heatmap
+
                         elif last_output is None:
-                            output = np.zeros((7))
+                            output = np.empty((1, 7))
+                            output[:] = np.nan
                             
                     probs.append(output[0])
                     frames.append(count_frame)
             else:
-                empty = np.empty((7))
-                empty[:] = np.nan
-                probs.append(empty)
-                frames.append(count_frame)                        
-
-            cur_face = cv2.cvtColor(cur_face, cv2.COLOR_RGB2BGR)
-            cur_face = cv2.resize(cur_face, (224,224), interpolation = cv2.INTER_AREA)
+                if last_output is not None:
+                    lstm_features = []
+                    empty = np.empty((7))
+                    empty[:] = np.nan
+                    probs.append(empty)
+                    frames.append(count_frame)                        
+
+            if cur_face is not None:
+                heatmap_f = display_info(heatmap, 'Frame: {}'.format(count_frame), box_scale=.3)
+
+                cur_face = cv2.cvtColor(cur_face, cv2.COLOR_RGB2BGR)
+                cur_face = cv2.resize(cur_face, (224,224), interpolation = cv2.INTER_AREA)
+                cur_face = display_info(cur_face, 'Frame: {}'.format(count_frame), box_scale=.3)
+                vid_writer_face.write(cur_face)
+                vid_writer_hm.write(heatmap_f)
 
-            cur_face = display_info(cur_face, 'Frame: {}'.format(count_frame), box_scale=.3)
             count_frame += 1
-            vid_writer.write(cur_face)
+            if count_face != 0:
+                count_face += 1
 
-        vid_writer.release()
+        vid_writer_face.release()
+        vid_writer_hm.release()
 
         stat = statistics_plot(frames, probs)
 
         if not stat:
-            return None, None, None
+            return None, None, None, None
         
-    return video, path_save_video, stat
+    return video, path_save_video_face, path_save_video_hm, stat

app/face_utils.py

@@ -34,7 +34,8 @@ def get_box(fl, w, h):
     return startX, startY, endX, endY
 
 def display_info(img, text, margin=1.0, box_scale=1.0):
-    img_h, img_w, _ = img.shape
+    img_copy = img.copy()
+    img_h, img_w, _ = img_copy.shape
     line_width = int(min(img_h, img_w) * 0.001)
     thickness = max(int(line_width / 3), 1)
 
@@ -45,15 +46,15 @@ def display_info(img, text, margin=1.0, box_scale=1.0):
     t_w, t_h = cv2.getTextSize(text, font_face, font_scale, None)[0]
 
     margin_n = int(t_h * margin)
-    sub_img = img[0 + margin_n: 0 + margin_n + t_h + int(2 * t_h * box_scale),
+    sub_img = img_copy[0 + margin_n: 0 + margin_n + t_h + int(2 * t_h * box_scale),
               img_w - t_w - margin_n - int(2 * t_h * box_scale): img_w - margin_n]
 
     white_rect = np.ones(sub_img.shape, dtype=np.uint8) * 255
 
-    img[0 + margin_n: 0 + margin_n + t_h + int(2 * t_h * box_scale),
+    img_copy[0 + margin_n: 0 + margin_n + t_h + int(2 * t_h * box_scale),
     img_w - t_w - margin_n - int(2 * t_h * box_scale):img_w - margin_n] = cv2.addWeighted(sub_img, 0.5, white_rect, .5, 1.0)
 
-    cv2.putText(img=img,
+    cv2.putText(img=img_copy,
                 text=text,
                 org=(img_w - t_w - margin_n - int(2 * t_h * box_scale) // 2,
                      0 + margin_n + t_h + int(2 * t_h * box_scale) // 2),
@@ -64,4 +65,4 @@ def display_info(img, text, margin=1.0, box_scale=1.0):
                 lineType=cv2.LINE_AA,
                 bottomLeftOrigin=False)
 
-    return img
+    return img_copy

app/model.py

@@ -10,9 +10,11 @@ import torch
 import requests
 from PIL import Image
 from torchvision import transforms
+from pytorch_grad_cam import GradCAM
 
 # Importing necessary components for the Gradio app
 from app.config import config_data
+from app.model_architectures import ResNet50, LSTMPyTorch
 
 
 def load_model(model_url, model_path):
@@ -21,17 +23,23 @@ def load_model(model_url, model_path):
             with open(model_path, "wb") as file:
                 for chunk in response.iter_content(chunk_size=8192):
                     file.write(chunk)
-        return torch.jit.load(model_path).eval()
+        return model_path
     except Exception as e:
         print(f"Error loading model: {e}")
         return None
 
+path_static = load_model(config_data.model_static_url, config_data.model_static_path)   
+pth_model_static = ResNet50(7, channels=3)
+pth_model_static.load_state_dict(torch.load(path_static))
+pth_model_static.eval()
 
-pth_model_static = load_model(config_data.model_static_url, config_data.model_static_path)
-
-pth_model_dynamic = load_model(config_data.model_dynamic_url, config_data.model_dynamic_path)
-
+path_dynamic = load_model(config_data.model_dynamic_url, config_data.model_dynamic_path) 
+pth_model_dynamic = LSTMPyTorch()
+pth_model_dynamic.load_state_dict(torch.load(path_dynamic))
+pth_model_dynamic.eval()
 
+target_layers = [pth_model_static.layer4]
+cam = GradCAM(model=pth_model_static, target_layers=target_layers)
 
 def pth_processing(fp):
     class PreprocessInput(torch.nn.Module):

app/model_architectures.py

@@ -0,0 +1,150 @@
+"""
+File: model.py
+Author: Elena Ryumina and Dmitry Ryumin
+Description: This module provides model architectures.
+License: MIT License
+"""
+
+import torch
+import torch.nn as  nn
+import torch.nn.functional as F
+import math
+
+class Bottleneck(nn.Module):
+    expansion = 4
+    def __init__(self, in_channels, out_channels, i_downsample=None, stride=1):
+        super(Bottleneck, self).__init__()
+        
+        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=stride, padding=0, bias=False)
+        self.batch_norm1 = nn.BatchNorm2d(out_channels, eps=0.001, momentum=0.99)
+        
+        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, padding='same', bias=False)
+        self.batch_norm2 = nn.BatchNorm2d(out_channels, eps=0.001, momentum=0.99)
+        
+        self.conv3 = nn.Conv2d(out_channels, out_channels*self.expansion, kernel_size=1, stride=1, padding=0, bias=False)
+        self.batch_norm3 = nn.BatchNorm2d(out_channels*self.expansion, eps=0.001, momentum=0.99)
+        
+        self.i_downsample = i_downsample
+        self.stride = stride
+        self.relu = nn.ReLU()
+        
+    def forward(self, x):
+        identity = x.clone()
+        x = self.relu(self.batch_norm1(self.conv1(x)))
+        
+        x = self.relu(self.batch_norm2(self.conv2(x)))
+        
+        x = self.conv3(x)
+        x = self.batch_norm3(x)
+        
+        #downsample if needed
+        if self.i_downsample is not None:
+            identity = self.i_downsample(identity)
+        #add identity
+        x+=identity
+        x=self.relu(x)
+        
+        return x
+
+class Conv2dSame(torch.nn.Conv2d):
+
+    def calc_same_pad(self, i: int, k: int, s: int, d: int) -> int:
+        return max((math.ceil(i / s) - 1) * s + (k - 1) * d + 1 - i, 0)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        ih, iw = x.size()[-2:]
+
+        pad_h = self.calc_same_pad(i=ih, k=self.kernel_size[0], s=self.stride[0], d=self.dilation[0])
+        pad_w = self.calc_same_pad(i=iw, k=self.kernel_size[1], s=self.stride[1], d=self.dilation[1])
+
+        if pad_h > 0 or pad_w > 0:
+            x = F.pad(
+                x, [pad_w // 2, pad_w - pad_w // 2, pad_h // 2, pad_h - pad_h // 2]
+            )
+        return F.conv2d(
+            x,
+            self.weight,
+            self.bias,
+            self.stride,
+            self.padding,
+            self.dilation,
+            self.groups,
+        )
+
+class ResNet(nn.Module):
+    def __init__(self, ResBlock, layer_list, num_classes, num_channels=3):
+        super(ResNet, self).__init__()
+        self.in_channels = 64
+
+        self.conv_layer_s2_same = Conv2dSame(num_channels, 64, 7, stride=2, groups=1, bias=False)
+        self.batch_norm1 = nn.BatchNorm2d(64, eps=0.001, momentum=0.99)
+        self.relu = nn.ReLU()
+        self.max_pool = nn.MaxPool2d(kernel_size = 3, stride=2)
+        
+        self.layer1 = self._make_layer(ResBlock, layer_list[0], planes=64, stride=1)
+        self.layer2 = self._make_layer(ResBlock, layer_list[1], planes=128, stride=2)
+        self.layer3 = self._make_layer(ResBlock, layer_list[2], planes=256, stride=2)
+        self.layer4 = self._make_layer(ResBlock, layer_list[3], planes=512, stride=2)
+        
+        self.avgpool = nn.AdaptiveAvgPool2d((1,1))
+        self.fc1 = nn.Linear(512*ResBlock.expansion, 512)
+        self.relu1 = nn.ReLU()
+        self.fc2 = nn.Linear(512, num_classes)
+
+    def extract_features(self, x):
+        x = self.relu(self.batch_norm1(self.conv_layer_s2_same(x)))
+        x = self.max_pool(x)
+        # print(x.shape)
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.layer4(x)
+        
+        x = self.avgpool(x)
+        x = x.reshape(x.shape[0], -1)
+        x = self.fc1(x)
+        return x
+        
+    def forward(self, x):
+        x = self.extract_features(x)
+        x = self.relu1(x)
+        x = self.fc2(x)
+        return x
+        
+    def _make_layer(self, ResBlock, blocks, planes, stride=1):
+        ii_downsample = None
+        layers = []
+        
+        if stride != 1 or self.in_channels != planes*ResBlock.expansion:
+            ii_downsample = nn.Sequential(
+                nn.Conv2d(self.in_channels, planes*ResBlock.expansion, kernel_size=1, stride=stride, bias=False, padding=0),
+                nn.BatchNorm2d(planes*ResBlock.expansion, eps=0.001, momentum=0.99)
+            )
+            
+        layers.append(ResBlock(self.in_channels, planes, i_downsample=ii_downsample, stride=stride))
+        self.in_channels = planes*ResBlock.expansion
+        
+        for i in range(blocks-1):
+            layers.append(ResBlock(self.in_channels, planes))
+            
+        return nn.Sequential(*layers)
+        
+def ResNet50(num_classes, channels=3):
+    return ResNet(Bottleneck, [3,4,6,3], num_classes, channels)
+
+
+class LSTMPyTorch(nn.Module):
+    def __init__(self):
+        super(LSTMPyTorch, self).__init__()
+        
+        self.lstm1 = nn.LSTM(input_size=512, hidden_size=512, batch_first=True, bidirectional=False)
+        self.lstm2 = nn.LSTM(input_size=512, hidden_size=256, batch_first=True, bidirectional=False)
+        self.fc = nn.Linear(256, 7)
+        self.softmax = nn.Softmax(dim=1)
+
+    def forward(self, x):
+        x, _ = self.lstm1(x)
+        x, _ = self.lstm2(x)        
+        x = self.fc(x[:, -1, :])
+        x = self.softmax(x)
+        return x

config.toml

@@ -2,9 +2,9 @@ APP_VERSION = "0.2.0"
 FRAME_DOWNSAMPLING = 5
 
 [model_static]
-url = "https://huggingface.co/ElenaRyumina/face_emotion_recognition/resolve/main/FER_static_ResNet50_AffectNet.pth"
-path = "FER_static_ResNet50_AffectNet.pth"
+url = "https://huggingface.co/ElenaRyumina/face_emotion_recognition/resolve/main/FER_static_ResNet50_AffectNet.pt"
+path = "FER_static_ResNet50_AffectNet.pt"
 
 [model_dynamic]
-url = "https://huggingface.co/ElenaRyumina/face_emotion_recognition/resolve/main/FER_dinamic_LSTM_IEMOCAP.pth"
-path = "FER_dinamic_LSTM_IEMOCAP.pth"
+url = "https://huggingface.co/ElenaRyumina/face_emotion_recognition/resolve/main/FER_dinamic_LSTM_IEMOCAP.pt"
+path = "FER_dinamic_LSTM_IEMOCAP.pt"