test

app.py

@@ -1,4 +1,23 @@
-import streamlit as st 
+import streamlit as st
 
+from model_processing import *
+from video_input import video_input
 
-st.title('Hello')
+st.set_page_config(
+    page_title=r"Analysis TikTok",
+    page_icon=":dizzy:",
+    layout="wide"
+)
+
+model = TikTokAnalytics()
+video_path, video_container, message_container = video_input()
+st.session_state['video_path'] = video_path
+
+if 'video_path' in st.session_state:
+    with st.spinner('Wait. We analyze your video'):
+        final_metric = [0, 0.1]  # model(video_path)[0]
+        message_container.write(final_metric)
+
+    message_container.write(
+        'People will like your video' if final_metric[1] > final_metric[0] else "People won't like your video"
+    )

config.py

@@ -0,0 +1,6 @@
+FPS_DIV = 3  # Во сколько раз уменьшаем fps. Все видео в тиктоке 30 fps. По сути анализируем с fps = 10
+MAX_LENGTH = 90  # Фиксированное количество кадров. То есть обрезаем каждое видео до max_length / 30 * fps_mult секунд
+
+BATCH_SIZE = 4  # Количество кадров для обработки батчем
+
+MODEL_PATH = './src/cash/best-model.th'

model_processing.py

@@ -0,0 +1,158 @@
+import ffmpegio
+import gc
+import torch
+from transformers import MobileViTImageProcessor, MobileViTForSemanticSegmentation
+from config import FPS_DIV, MAX_LENGTH, BATCH_SIZE, MODEL_PATH
+
+
+class PreprocessModel(torch.nn.Module):
+    device = 'cpu'
+
+    def __init__(self):
+        super().__init__()
+        self.feature_extractor = MobileViTImageProcessor.from_pretrained("apple/deeplabv3-mobilevit-small")
+        self.mobile_vit = MobileViTForSemanticSegmentation.from_pretrained("apple/deeplabv3-mobilevit-small")
+        self.convs = torch.nn.Sequential(
+            torch.nn.MaxPool2d(2, 2)
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.mobile_vit(x).logits
+        x = self.convs(x)
+        return x
+
+    def read_video(self, path: str) -> torch.Tensor:
+        """
+        Читает видео и возвращает тензор с фичами
+        """
+
+        _, video = ffmpegio.video.read(path, t=1.0)
+        video = video[::FPS_DIV][:MAX_LENGTH]
+
+        out_seg_video = []
+
+        for i in range(0, video.shape[0], BATCH_SIZE):
+            frames = [video[j] for j in range(i, min(i + BATCH_SIZE, video.shape[0]))]
+            frames = self.feature_extractor(images=frames, return_tensors='pt')['pixel_values']
+
+            out = self.forward(frames.to(self.device)).detach().to('cpu')
+            out_seg_video.append(out)
+
+            del frames, out
+            gc.collect()
+            if self.device == 'cuda':
+                torch.cuda.empty_cache()
+
+        return torch.cat(out_seg_video)
+
+
+class VideoModel(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        p = 0.5
+        self.pic_cnn = torch.nn.Sequential(
+            torch.nn.Conv2d(21, 128, (2, 2), stride=2),
+            torch.nn.BatchNorm2d(128),
+            torch.nn.LeakyReLU(),
+            torch.nn.Conv2d(128, 256, (2, 2), stride=2),
+            torch.nn.BatchNorm2d(256),
+            torch.nn.Dropout2d(p),
+            torch.nn.LeakyReLU(),
+            torch.nn.Conv2d(256, 256, (4, 4), stride=2),
+            torch.nn.BatchNorm2d(256),
+            torch.nn.Dropout2d(p),
+            torch.nn.Flatten()
+        )
+
+        self.vid_cnn = torch.nn.Sequential(
+            torch.nn.Conv2d(21, 128, (2, 2), stride=2),
+            torch.nn.BatchNorm2d(128),
+            torch.nn.Tanh(),
+            torch.nn.Conv2d(128, 256, (2, 2), stride=2),
+            torch.nn.BatchNorm2d(256),
+            torch.nn.Dropout2d(p),
+            torch.nn.LeakyReLU(),
+            torch.nn.Conv2d(256, 512, (2, 2), stride=2),
+            torch.nn.BatchNorm2d(512),
+            torch.nn.Dropout2d(p),
+            torch.nn.Flatten()
+        )
+
+        self.lstm = torch.nn.LSTM(2048, 256, 1, batch_first=True, bidirectional=True)
+        self.fc1 = torch.nn.Linear(256 * 2, 1024)
+        self.fc_norm = torch.nn.BatchNorm1d(256 * 2)
+        self.tanh = torch.nn.Tanh()
+        self.fc2 = torch.nn.Linear(1024, 2)
+        self.sigmoid = torch.nn.Sigmoid()
+        self.dropout = torch.nn.Dropout(p)
+
+        # xaiver init
+        for m in self.modules():
+            if isinstance(m, torch.nn.Conv2d) or isinstance(m, torch.nn.Conv3d):
+                torch.nn.init.xavier_uniform_(m.weight)
+                if m.bias is not None:
+                    torch.nn.init.zeros_(m.bias)
+
+            elif isinstance(m, torch.nn.Linear):
+                torch.nn.init.xavier_uniform_(m.weight)
+                if m.bias is not None:
+                    torch.nn.init.zeros_(m.bias)
+
+    def forward(self, video: torch.Tensor) -> torch.Tensor:
+        """
+        Использует превью как начальное скрытое состояние, а кадры видео как последовательность.
+        video[0] - превью, video[1] - видео
+
+        :param video: torch.Tensor, shape = (batch_size, frames + 1, 1344)
+        """
+        frames = video.shape[0]
+        video = torch.nn.functional.pad(video, (0, 0, 0, 0, 0, 0, MAX_LENGTH + 1 - frames, 0))
+        video = video.unsqueeze(0)
+        _batch_size = video.shape[0]
+
+        _preview = video[:, 0, :, :]
+        _video = video[:, 1:, :, :]
+
+        h0 = self.pic_cnn(_preview).unsqueeze(0)
+        h0 = torch.nn.functional.pad(h0, (0, 0, 0, 0, 0, 1))
+        c0 = torch.zeros_like(h0)
+
+        _video = self.vid_cnn(_video.reshape(-1, 21, 16, 16))
+        _video = _video.reshape(_batch_size, 90, -1)
+
+        context, _ = self.lstm(_video, (h0, c0))
+        out = self.fc_norm(context[:, -1])
+        out = self.tanh(self.fc1(out))
+        out = self.dropout(out)
+        out = self.sigmoid(self.fc2(out))
+        return out
+
+
+# @st.cache_resource
+class TikTokAnalytics(torch.nn.Module):
+
+    def __init__(self):
+        super().__init__()
+        self.preprocessing_model = PreprocessModel()
+        self.predict_model = torch.load(MODEL_PATH, map_location=self.preprocessing_model.device)
+
+        self.preprocessing_model.eval()
+        self.predict_model.eval()
+
+    def forward(self, path: str) -> torch.Tensor:
+        """
+        Вызываем препроцесс, потом предикт
+        :param path:
+        :return:
+        """
+        tensor = self.preprocessing_model.read_video(path)
+        predict = self.predict_model(tensor)
+
+        return predict
+
+
+# if __name__ == '__main__':
+#     model = TikTokAnalytics()
+#     model = model(
+#         '/Users/victorbarbarich/PycharmProjects/nueramic/vktrbr-video-tiktok/data/videos/video-6930454291186502917.mp4')
+#     print(model)

requirements.txt

@@ -0,0 +1,256 @@
+# This file may be used to create an environment using:
+# $ conda create --name <env> --file <this file>
+# platform: osx-arm64
+aiohttp=3.8.3=pypi_0
+aiomoex=2.0.0=pypi_0
+aiosignal=1.2.0=pypi_0
+alembic=1.8.1=pypi_0
+altair=4.2.0=pypi_0
+ansi2html=1.8.0=pypi_0
+anyio=3.6.1=pypi_0
+appnope=0.1.3=pyhd8ed1ab_0
+apscheduler=3.9.1=pypi_0
+arch=5.2.0=pypi_0
+argon2-cffi=21.3.0=pyhd8ed1ab_0
+argon2-cffi-bindings=21.2.0=py39h02fc5c5_3
+asttokens=2.1.0=pyhd8ed1ab_0
+async-generator=1.10=pypi_0
+async-timeout=4.0.2=pypi_0
+attrs=22.1.0=pyh71513ae_1
+backcall=0.2.0=pyh9f0ad1d_0
+backports=1.0=py_2
+backports.functools_lru_cache=1.6.4=pyhd8ed1ab_0
+beautifulsoup4=4.11.1=pyha770c72_0
+bleach=5.0.1=pyhd8ed1ab_0
+blinker=1.5=pypi_0
+brotli=1.0.9=pypi_0
+brotlipy=0.7.0=py39h02fc5c5_1005
+build=0.8.0=pypi_0
+bzip2=1.0.8=h3422bc3_4
+ca-certificates=2022.12.7=h4653dfc_0
+cachetools=5.2.0=pypi_0
+catboost=1.0.6=pypi_0
+certifi=2022.12.7=pyhd8ed1ab_0
+certipy=0.1.3=pypi_0
+cffi=1.15.1=py39h7e6b969_1
+charset-normalizer=2.1.1=pyhd8ed1ab_0
+click=8.1.3=pypi_0
+colorama=0.4.6=pyhd8ed1ab_0
+commonmark=0.9.1=pypi_0
+conda=23.1.0=py39h2804cbe_0
+conda-package-handling=1.9.0=py39h02fc5c5_1
+cryptography=38.0.3=py39he2a39a8_0
+cycler=0.11.0=pypi_0
+dash=2.6.1=pypi_0
+dash-core-components=2.0.0=pypi_0
+dash-html-components=2.0.0=pypi_0
+dash-table=5.0.0=pypi_0
+dataclasses=0.6=pypi_0
+debugpy=1.6.3=py39h23fbdae_1
+decorator=5.1.1=pyhd8ed1ab_0
+defusedxml=0.7.1=pyhd8ed1ab_0
+docutils=0.18.1=pypi_0
+entrypoints=0.4=pyhd8ed1ab_0
+esgscraper=1.0.0=pypi_0
+et-xmlfile=1.1.0=pypi_0
+exceptiongroup=1.0.4=pypi_0
+executing=1.2.0=pyhd8ed1ab_0
+ffmpegio=0.7.0=pypi_0
+ffmpegio-core=0.8.0=pypi_0
+flask=2.2.2=pypi_0
+flask-compress=1.12=pypi_0
+flit-core=3.8.0=pyhd8ed1ab_0
+fonttools=4.33.3=pypi_0
+frozenlist=1.3.1=pypi_0
+gh=2.22.1=h75b854d_0
+gitdb=4.0.9=pypi_0
+gitpython=3.1.27=pypi_0
+greenlet=2.0.1=pypi_0
+h11=0.12.0=pypi_0
+httpcore=0.15.0=pypi_0
+httpx=0.23.0=pypi_0
+huggingface-hub=0.13.1=pypi_0
+idna=3.4=pyhd8ed1ab_0
+importlib-metadata=5.0.0=pyha770c72_1
+importlib_resources=5.10.0=pyhd8ed1ab_0
+iniconfig=1.1.1=pypi_0
+ipykernel=6.13.0=py39h32adebf_0
+ipython=8.6.0=pyhd1c38e8_1
+ipython_genutils=0.2.0=py_1
+ipywidgets=7.7.1=pypi_0
+itsdangerous=2.1.2=pypi_0
+jedi=0.18.1=pyhd8ed1ab_2
+jinja2=3.1.2=pyhd8ed1ab_1
+joblib=1.1.0=pypi_0
+jsonschema=4.17.0=pyhd8ed1ab_0
+jupyter=1.0.0=py39h2804cbe_7
+jupyter-dash=0.4.2=pypi_0
+jupyter-telemetry=0.1.0=pypi_0
+jupyter_client=7.3.1=pyhd8ed1ab_0
+jupyter_console=6.4.3=pyhd8ed1ab_0
+jupyter_core=5.0.0=py39h2804cbe_0
+jupyterhub=3.1.0=pypi_0
+jupyterlab_pygments=0.2.2=pyhd8ed1ab_0
+jupyterlab_widgets=1.1.1=pyhd8ed1ab_0
+kaggle=1.5.12=pypi_0
+kaleido=0.2.1=pypi_0
+keyring=23.5.1=pypi_0
+kiwisolver=1.4.2=pypi_0
+libblas=3.9.0=16_osxarm64_openblas
+libcblas=3.9.0=16_osxarm64_openblas
+libcxx=14.0.6=h2692d47_0
+libffi=3.4.2=h3422bc3_5
+libgfortran=5.0.0=11_3_0_hd922786_26
+libgfortran5=11.3.0=hdaf2cc0_26
+liblapack=3.9.0=16_osxarm64_openblas
+libopenblas=0.3.21=openmp_hc731615_3
+libsodium=1.0.18=h27ca646_1
+libsqlite=3.40.0=h76d750c_0
+libzlib=1.2.13=h03a7124_4
+lightgbm=3.3.2=py39h8685b90_0
+llvm-openmp=15.0.5=h7cfbb63_0
+llvmlite=0.38.1=pypi_0
+lxml=4.9.0=pypi_0
+mako=1.2.4=pypi_0
+markupsafe=2.1.1=py39h02fc5c5_2
+matplotlib=3.5.2=pypi_0
+matplotlib-inline=0.1.6=pyhd8ed1ab_0
+mistune=0.8.4=py39h5161555_1005
+mpmath=1.2.1=pypi_0
+multidict=6.0.2=pypi_0
+multitasking=0.0.10=pypi_0
+nbclient=0.5.13=pyhd8ed1ab_0
+nbconvert=6.4.5=py39h2804cbe_0
+nbformat=5.7.0=pyhd8ed1ab_0
+ncurses=6.3=h07bb92c_1
+nest-asyncio=1.5.6=pyhd8ed1ab_0
+networkx=2.8.4=pypi_0
+notebook=6.4.11=pyha770c72_0
+nueramic-mathml=0.75.2=pypi_0
+numpy=1.22.4=pypi_0
+oauthlib=3.2.2=pypi_0
+openpyxl=3.0.10=pypi_0
+openssl=3.0.8=h03a7124_0
+outcome=1.2.0=pypi_0
+packaging=21.3=pyhd8ed1ab_0
+pamela=1.0.0=pypi_0
+pandas=1.4.2=pypi_0
+pandoc=2.19.2=hce30654_1
+pandocfilters=1.5.0=pyhd8ed1ab_0
+parso=0.8.3=pyhd8ed1ab_0
+patsy=0.5.2=pypi_0
+pep517=0.12.0=pypi_0
+pexpect=4.8.0=pyh1a96a4e_2
+pickleshare=0.7.5=py_1003
+pillow=9.2.0=pypi_0
+pip=22.3.1=pypi_0
+pkginfo=1.8.2=pypi_0
+pkgutil-resolve-name=1.3.10=pyhd8ed1ab_0
+platformdirs=2.5.2=pyhd8ed1ab_1
+playwright=1.28.0=pypi_0
+plotly=5.8.2=pypi_0
+pluggy=1.0.0=pyhd8ed1ab_5
+prometheus_client=0.15.0=pyhd8ed1ab_0
+prompt-toolkit=3.0.30=pypi_0
+prompt_toolkit=3.0.32=hd8ed1ab_0
+property-cached=1.6.4=pypi_0
+protobuf=3.20.3=pypi_0
+psutil=5.9.4=py39h02fc5c5_0
+ptyprocess=0.7.0=pyhd3deb0d_0
+pure_eval=0.2.2=pyhd8ed1ab_0
+pyarrow=9.0.0=pypi_0
+pycosat=0.6.4=py39h02fc5c5_1
+pycparser=2.21=pyhd8ed1ab_0
+pydeck=0.8.0b3=pypi_0
+pyee=9.0.4=pypi_0
+pygments=2.13.0=pyhd8ed1ab_0
+pympler=1.0.1=pypi_0
+pyopenssl=22.1.0=pyhd8ed1ab_0
+pyparsing=3.0.9=pyhd8ed1ab_0
+pyrsistent=0.19.2=py39h02fc5c5_0
+pysocks=1.7.1=pyha2e5f31_6
+pytest=7.2.0=pypi_0
+python=3.9.13=h96fcbfb_0_cpython
+python-dateutil=2.8.2=pyhd8ed1ab_0
+python-fastjsonschema=2.16.2=pyhd8ed1ab_0
+python-graphviz=0.20=pypi_0
+python-json-logger=2.0.4=pypi_0
+python-slugify=6.1.2=pypi_0
+python-telegram-bot=20.0a4=pypi_0
+python_abi=3.9=2_cp39
+pytz=2022.1=pypi_0
+pytz-deprecation-shim=0.1.0.post0=pypi_0
+pyyaml=6.0=pypi_0
+pyzmq=24.0.1=py39h0553236_1
+readline=8.1.2=h46ed386_0
+readme-renderer=35.0=pypi_0
+regex=2022.10.31=pypi_0
+requests=2.28.1=pyhd8ed1ab_1
+requests-toolbelt=0.9.1=pypi_0
+retrying=1.3.3=pypi_0
+rfc3986=1.5.0=pypi_0
+rich=12.4.4=pypi_0
+ruamel.yaml=0.17.21=py39h02fc5c5_2
+ruamel.yaml.clib=0.2.7=py39h02fc5c5_1
+ruamel_yaml=0.15.80=py39h02fc5c5_1008
+scikit-learn=1.1.1=pypi_0
+scipy=1.8.1=pypi_0
+selenium=4.6.1=pypi_0
+semver=2.13.0=pypi_0
+send2trash=1.8.0=pyhd8ed1ab_0
+setuptools=62.6.0=pypi_0
+six=1.16.0=pyh6c4a22f_0
+slicer=0.0.7=pypi_0
+smmap=5.0.0=pypi_0
+sniffio=1.3.0=pypi_0
+sortedcontainers=2.4.0=pypi_0
+soupsieve=2.3.2.post1=pyhd8ed1ab_0
+split-folders=0.5.1=pypi_0
+sqlalchemy=1.4.41=pypi_0
+sqldf=0.4.2=pypi_0
+sqlite=3.40.0=h2229b38_0
+stack_data=0.6.1=pyhd8ed1ab_0
+statsmodels=0.13.2=pypi_0
+streamlit=1.20.0=pypi_0
+sympy=1.10.1=pypi_0
+telegram=0.0.1=pypi_0
+tenacity=8.0.1=pypi_0
+terminado=0.15.0=py39h2804cbe_0
+testpath=0.6.0=pyhd8ed1ab_0
+text-unidecode=1.3=pypi_0
+threadpoolctl=3.1.0=pyh8a188c0_0
+tiktokapi=5.2.2=pypi_0
+tk=8.6.12=he1e0b03_0
+tokenizers=0.13.2=pypi_0
+toml=0.10.2=pypi_0
+tomli=2.0.1=pypi_0
+toolz=0.12.0=pyhd8ed1ab_0
+torch=1.13.0.dev20220529=pypi_0
+torchaudio=0.11.0=pypi_0
+torchvision=0.12.0=pypi_0
+tornado=6.2=py39h02fc5c5_1
+tqdm=4.64.0=pypi_0
+traitlets=5.3.0=pypi_0
+transformers=4.26.1=pypi_0
+trio=0.21.0=pypi_0
+trio-websocket=0.9.2=pypi_0
+twine=4.0.1=pypi_0
+typing-extensions=4.2.0=pypi_0
+typing_extensions=4.4.0=pyha770c72_0
+tzdata=2022.4=pypi_0
+tzlocal=4.2=pypi_0
+urllib3=1.26.11=pyhd8ed1ab_0
+validators=0.20.0=pypi_0
+wcwidth=0.2.5=pyh9f0ad1d_2
+webencodings=0.5.1=py_1
+werkzeug=2.2.2=pypi_0
+wheel=0.38.4=pyhd8ed1ab_0
+widgetsnbextension=3.6.0=py39h2804cbe_0
+wsproto=1.2.0=pypi_0
+xz=5.2.6=h57fd34a_0
+yaml=0.2.5=h3422bc3_2
+yarl=1.8.1=pypi_0
+yfinance=0.1.70=pypi_0
+zeromq=4.3.4=hbdafb3b_1
+zipp=3.10.0=pyhd8ed1ab_0
+zlib=1.2.13=h03a7124_4

video_input.py

@@ -0,0 +1,27 @@
+import streamlit as st
+
+
+def video_input() -> [str, st.container, st.container]:
+    """
+
+    :return: путь к файлу
+    """
+
+    st.title("Video Input")
+    video_file = st.file_uploader("Upload a video", type=["mp4"])
+
+    if video_file is None:
+        st.stop()
+
+    _, video_container, message_container, _ = st.columns([1, 2, 2, 1])
+    video_container: st.empty()
+
+    if video_file is not None:
+        video_container.video(video_file)
+
+    file_path = f'./{abs(hash(str(video_file)))}.mp4'
+
+    with open(file_path, 'wb') as file:
+        file.write(video_file.read())
+
+    return file_path, video_container, message_container