Add source code

.gitattributes

@@ -32,3 +32,8 @@ 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
+egs/video1.mp4 filter=lfs diff=lfs merge=lfs -text
+egs/video2.mp4 filter=lfs diff=lfs merge=lfs -text
+egs/video3.mp4 filter=lfs diff=lfs merge=lfs -text
+egs/video4.mp4 filter=lfs diff=lfs merge=lfs -text
+src/model filter=lfs diff=lfs merge=lfs -text

app.py

@@ -0,0 +1,35 @@
+import gradio as gr
+import glob
+import os
+
+from src.core import load_model, predict_traits
+
+TRAIT_NAMES = [
+    "Extraversion",
+    "Agreeableness",
+    "Conscientiousness",
+    "Neurotisicm",
+    "Openness",
+]
+
+
+def get_traits(video):
+    model = load_model()
+    # if webcam_video:
+    #     trait_values = predict_traits(webcam_video, model)
+    # else:
+    trait_values = predict_traits(video, model)
+    return {k: float(v) for k, v in zip(TRAIT_NAMES, trait_values)}
+
+
+demo = gr.Interface(
+    get_traits,
+    inputs=gr.Video(label="Video", include_audio=True),
+    outputs=gr.Label(num_top_classes=5, label="Results"),
+    title="Personality Traits Prediction [Prototype]",
+    description="Predicts the 5 psychological traits using an introduction video",
+    thumbnail="https://cdn-icons-png.flaticon.com/512/3392/3392044.png",
+    examples="egs",
+    cache_examples=True,
+)
+demo.launch()

egs/video1.mp4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:956627f2dcb5559054718462157e436c5be0bc70ff32ae4194f334e32feafcee
+size 2969523

egs/video2.mp4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8e67ba771391717777022cf604dc1bef79e20ded23a7055f54f33325399f5296
+size 2311239

egs/video3.mp4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7408594f52bb5edf658ed32c559c6e05526e21ec3d6ff7e0b69be03f8221784a
+size 2579751

egs/video4.mp4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:410f03f6e0ecb6b170ed65b47e520f7e071c19bb7e0f11f1c57808d095a2298f
+size 2790813

src/audiovisual_stream.py

@@ -0,0 +1,39 @@
+from . import auditory_stream, visual_stream
+import chainer
+
+
+class ResNet18(chainer.Chain):
+    def __init__(self):
+        super(ResNet18, self).__init__(
+            aud=auditory_stream.ResNet18(),
+            vis=visual_stream.ResNet18(),
+            fc=chainer.links.Linear(512, 5, initialW=chainer.initializers.HeNormal()),
+        )
+
+    def __call__(self, x):
+        h = [
+            self.aud(chainer.Variable(chainer.cuda.to_cpu(x[0]))),
+            chainer.functions.expand_dims(
+                chainer.functions.sum(
+                    self.vis(chainer.Variable(chainer.cuda.to_cpu(x[1][:256]))), 0
+                ),
+                0,
+            ),
+        ]
+
+        for i in range(256, x[1].shape[0], 256):
+            h[1] += chainer.functions.expand_dims(
+                chainer.functions.sum(
+                    self.vis(chainer.Variable(chainer.cuda.to_cpu(x[1][i : i + 256]))),
+                    0,
+                ),
+                0,
+            )
+
+        h[1] /= x[1].shape[0]
+
+        return chainer.cuda.to_cpu(
+            (
+                (chainer.functions.tanh(self.fc(chainer.functions.concat(h))) + 1) / 2
+            ).data[0]
+        )

src/auditory_stream.py

@@ -0,0 +1,147 @@
+import chainer
+
+
+class ConvolutionBlock(chainer.Chain):
+    def __init__(self, in_channels, out_channels):
+        super(ConvolutionBlock, self).__init__(
+            conv=chainer.links.Convolution2D(
+                in_channels,
+                out_channels,
+                (1, 49),
+                (1, 4),
+                (0, 24),
+                initialW=chainer.initializers.HeNormal(),
+            ),
+            bn_conv=chainer.links.BatchNormalization(out_channels),
+        )
+
+    def __call__(self, x):
+        # Set Train to False.
+        chainer.config.train = False
+
+        h = self.conv(x)
+        h = self.bn_conv(h)
+        y = chainer.functions.relu(h)
+
+        return y
+
+
+class ResidualBlock(chainer.Chain):
+    def __init__(self, in_channels, out_channels):
+        super(ResidualBlock, self).__init__(
+            res_branch2a=chainer.links.Convolution2D(
+                in_channels,
+                out_channels,
+                (1, 9),
+                pad=(0, 4),
+                initialW=chainer.initializers.HeNormal(),
+            ),
+            bn_branch2a=chainer.links.BatchNormalization(out_channels),
+            res_branch2b=chainer.links.Convolution2D(
+                out_channels,
+                out_channels,
+                (1, 9),
+                pad=(0, 4),
+                initialW=chainer.initializers.HeNormal(),
+            ),
+            bn_branch2b=chainer.links.BatchNormalization(out_channels),
+        )
+
+    def __call__(self, x):
+        chainer.config.train = False
+
+        h = self.res_branch2a(x)
+        h = self.bn_branch2a(h)
+        h = chainer.functions.relu(h)
+        h = self.res_branch2b(h)
+        h = self.bn_branch2b(h)
+        h = x + h
+        y = chainer.functions.relu(h)
+
+        return y
+
+
+class ResidualBlockA:
+    def __init__(self):
+        pass
+
+    def __call__(self):
+        pass
+
+
+class ResidualBlockB(chainer.Chain):
+    def __init__(self, in_channels, out_channels):
+        super(ResidualBlockB, self).__init__(
+            res_branch1=chainer.links.Convolution2D(
+                in_channels,
+                out_channels,
+                (1, 1),
+                (1, 4),
+                initialW=chainer.initializers.HeNormal(),
+            ),
+            bn_branch1=chainer.links.BatchNormalization(out_channels),
+            res_branch2a=chainer.links.Convolution2D(
+                in_channels,
+                out_channels,
+                (1, 9),
+                (1, 4),
+                (0, 4),
+                initialW=chainer.initializers.HeNormal(),
+            ),
+            bn_branch2a=chainer.links.BatchNormalization(out_channels),
+            res_branch2b=chainer.links.Convolution2D(
+                out_channels,
+                out_channels,
+                (1, 9),
+                pad=(0, 4),
+                initialW=chainer.initializers.HeNormal(),
+            ),
+            bn_branch2b=chainer.links.BatchNormalization(out_channels),
+        )
+
+    def __call__(self, x):
+        chainer.config.train = False
+
+        temp = self.res_branch1(x)
+        temp = self.bn_branch1(temp)
+        h = self.res_branch2a(x)
+        h = self.bn_branch2a(h)
+        h = chainer.functions.relu(h)
+        h = self.res_branch2b(h)
+        h = self.bn_branch2b(h)
+        h = temp + h
+        y = chainer.functions.relu(h)
+
+        return y
+
+
+class ResNet18(chainer.Chain):
+    def __init__(self):
+        super(ResNet18, self).__init__(
+            conv1_relu=ConvolutionBlock(1, 32),
+            res2a_relu=ResidualBlock(32, 32),
+            res2b_relu=ResidualBlock(32, 32),
+            res3a_relu=ResidualBlockB(32, 64),
+            res3b_relu=ResidualBlock(64, 64),
+            res4a_relu=ResidualBlockB(64, 128),
+            res4b_relu=ResidualBlock(128, 128),
+            res5a_relu=ResidualBlockB(128, 256),
+            res5b_relu=ResidualBlock(256, 256),
+        )
+
+    def __call__(self, x):
+        chainer.config.train = False
+
+        h = self.conv1_relu(x)
+        h = chainer.functions.max_pooling_2d(h, (1, 9), (1, 4), (0, 4))
+        h = self.res2a_relu(h)
+        h = self.res2b_relu(h)
+        h = self.res3a_relu(h)
+        h = self.res3b_relu(h)
+        h = self.res4a_relu(h)
+        h = self.res4b_relu(h)
+        h = self.res5a_relu(h)
+        h = self.res5b_relu(h)
+        y = chainer.functions.average_pooling_2d(h, h.data.shape[2:])
+
+        return y

src/core.py

@@ -0,0 +1,43 @@
+from . import audiovisual_stream
+import chainer.serializers
+import librosa
+import numpy
+import skvideo.io
+import numpy as np
+
+FRAMES_LIMIT = 25
+
+
+def load_audio(data):
+    return librosa.load(data, 16000)[0][None, None, None, :]
+
+
+def load_model():
+    model = audiovisual_stream.ResNet18().to_cpu()
+    chainer.serializers.load_npz("src/model", model)
+    return model
+
+
+def predict_traits(data, model):
+    video_features = skvideo.io.vreader(data, num_frames=27)
+    # video_features = skvideo.io.vreader(data)
+
+    audio_features = load_audio(data)
+
+    x = []
+    predictions = []
+
+    frame_count = 0
+    for frame in video_features:
+        x.append(numpy.rollaxis(frame, 2))
+
+        frame_count += 1
+
+        if frame_count == FRAMES_LIMIT:
+            x = [audio_features, numpy.array(x, "float32")]
+            predictions.append(model(x))
+            
+            frame_count = 0
+            x = []
+
+    return np.mean(np.asarray(predictions), axis=0)

src/model

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a4794bfeae59d64c7ae9f78404c11d4517cf6da78b313301297dca4bc148deac
+size 19218612

src/visual_stream.py

@@ -0,0 +1,145 @@
+import chainer
+
+class ConvolutionBlock(chainer.Chain):
+    def __init__(self, in_channels, out_channels):
+        super(ConvolutionBlock, self).__init__(
+            conv=chainer.links.Convolution2D(
+                in_channels,
+                out_channels,
+                7,
+                2,
+                3,
+                initialW=chainer.initializers.HeNormal(),
+            ),
+            bn_conv=chainer.links.BatchNormalization(out_channels),
+        )
+
+    def __call__(self, x):
+        chainer.config.train = False
+
+        h = self.conv(x)
+        h = self.bn_conv(h)
+        y = chainer.functions.relu(h)
+
+        return y
+
+
+class ResidualBlock(chainer.Chain):
+    def __init__(self, in_channels, out_channels):
+        super(ResidualBlock, self).__init__(
+            res_branch2a=chainer.links.Convolution2D(
+                in_channels,
+                out_channels,
+                3,
+                pad=1,
+                initialW=chainer.initializers.HeNormal(),
+            ),
+            bn_branch2a=chainer.links.BatchNormalization(out_channels),
+            res_branch2b=chainer.links.Convolution2D(
+                out_channels,
+                out_channels,
+                3,
+                pad=1,
+                initialW=chainer.initializers.HeNormal(),
+            ),
+            bn_branch2b=chainer.links.BatchNormalization(out_channels),
+        )
+
+    def __call__(self, x):
+        chainer.config.train = False
+
+        h = self.res_branch2a(x)
+        h = self.bn_branch2a(h)
+        h = chainer.functions.relu(h)
+        h = self.res_branch2b(h)
+        h = self.bn_branch2b(h)
+        h = x + h
+        y = chainer.functions.relu(h)
+
+        return y
+
+
+class ResidualBlockA:
+    def __init__(self):
+        pass
+
+    def __call__(self):
+        pass
+
+
+class ResidualBlockB(chainer.Chain):
+    def __init__(self, in_channels, out_channels):
+        super(ResidualBlockB, self).__init__(
+            res_branch1=chainer.links.Convolution2D(
+                in_channels,
+                out_channels,
+                1,
+                2,
+                initialW=chainer.initializers.HeNormal(),
+            ),
+            bn_branch1=chainer.links.BatchNormalization(out_channels),
+            res_branch2a=chainer.links.Convolution2D(
+                in_channels,
+                out_channels,
+                3,
+                2,
+                1,
+                initialW=chainer.initializers.HeNormal(),
+            ),
+            bn_branch2a=chainer.links.BatchNormalization(out_channels),
+            res_branch2b=chainer.links.Convolution2D(
+                out_channels,
+                out_channels,
+                3,
+                pad=1,
+                initialW=chainer.initializers.HeNormal(),
+            ),
+            bn_branch2b=chainer.links.BatchNormalization(out_channels),
+        )
+
+    def __call__(self, x):
+        chainer.config.train = False
+
+        temp = self.res_branch1(x)
+        temp = self.bn_branch1(temp)
+        h = self.res_branch2a(x)
+        h = self.bn_branch2a(h)
+        h = chainer.functions.relu(h)
+        h = self.res_branch2b(h)
+        h = self.bn_branch2b(h)
+        h = temp + h
+        y = chainer.functions.relu(h)
+
+        return y
+
+
+class ResNet18(chainer.Chain):
+    def __init__(self):
+        super(ResNet18, self).__init__(
+            conv1_relu=ConvolutionBlock(3, 32),
+            res2a_relu=ResidualBlock(32, 32),
+            res2b_relu=ResidualBlock(32, 32),
+            res3a_relu=ResidualBlockB(32, 64),
+            res3b_relu=ResidualBlock(64, 64),
+            res4a_relu=ResidualBlockB(64, 128),
+            res4b_relu=ResidualBlock(128, 128),
+            res5a_relu=ResidualBlockB(128, 256),
+            res5b_relu=ResidualBlock(256, 256),
+        )
+
+    def __call__(self, x):
+        chainer.config.train = False
+
+        h = self.conv1_relu(x)
+        h = chainer.functions.max_pooling_2d(h, 3, 2, 1)
+        h = self.res2a_relu(h)
+        h = self.res2b_relu(h)
+        h = self.res3a_relu(h)
+        h = self.res3b_relu(h)
+        h = self.res4a_relu(h)
+        h = self.res4b_relu(h)
+        h = self.res5a_relu(h)
+        h = self.res5b_relu(h)
+        y = chainer.functions.average_pooling_2d(h, h.data.shape[2:])
+
+        return y