Upload 4 files

inference.py

@@ -0,0 +1,145 @@
+import cv2
+import torch
+import streamlit as st
+
+from PIL import Image
+from torch.nn import functional as nnf
+
+# @st.cache_data
+def generate2(
+        model,
+        tokenizer,
+        tokens=None,
+        prompt='',
+        embed=None,
+        entry_count=1,
+        entry_length=67,
+        top_p=0.98,
+        temperature=1,
+        stop_token='.',
+):
+
+    # model.eval()
+
+    generated_list = []
+    stop_token_index = tokenizer.encode(stop_token)[0]
+    filter_value = -float("Inf")
+    device = next(model.parameters()).device
+
+    with torch.no_grad():
+        for entry_idx in range(entry_count):
+            if not tokens:
+                tokens = torch.tensor(tokenizer.encode(prompt))
+                tokens = tokens.unsqueeze(0).to(device)
+
+            emb_tokens = model.gpt.transformer.wte(tokens)
+
+            if embed is not None:
+                generated = torch.cat((embed, emb_tokens), dim=1)
+            else:
+                generated = emb_tokens
+
+            for i in range(entry_length):
+                outputs = model.gpt(inputs_embeds=generated)
+                logits = outputs.logits
+                logits = logits[:, -1, :] / (temperature if temperature > 0 else 1.0)
+                sorted_logits, sorted_indices = torch.sort(logits, descending=True)
+                cumulative_probs = torch.cumsum(nnf.softmax(sorted_logits, dim=-1), dim=-1)
+                sorted_indices_to_remove = cumulative_probs > top_p
+                sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[
+                                                    ..., :-1
+                                                    ].clone()
+                sorted_indices_to_remove[..., 0] = 0
+
+                indices_to_remove = sorted_indices[sorted_indices_to_remove]
+                logits[:, indices_to_remove] = filter_value
+
+                top_k = 2000
+                top_p = 0.98
+
+                next_token = torch.argmax(logits, -1).unsqueeze(0)
+                next_token_embed = model.gpt.transformer.wte(next_token)
+
+                if tokens is None:
+                    tokens = next_token
+                else:
+                    tokens = torch.cat((tokens, next_token), dim=1)
+
+                generated = torch.cat((generated, next_token_embed), dim=1)
+
+                if stop_token_index == next_token.item():
+                    break
+
+            output_list = list(tokens.squeeze().cpu().numpy())
+
+            output_text = tokenizer.decode(output_list)
+            output_text = filter_ngrams(output_text)
+            generated_list.append(output_text)
+
+    return generated_list[0]
+
+def filter_ngrams(output_text):
+    a_pos = output_text.find(' A:')
+    sec_a_pos = output_text.find(' A:', a_pos + 1)
+
+    return output_text[:sec_a_pos]
+
+def image_grid(imgs, rows, cols):
+    assert len(imgs) == rows * cols
+
+    w, h = imgs[0].size
+    grid = Image.new('RGB', size=(cols * w, rows * h))
+    grid_w, grid_h = grid.size
+
+    for i, img in enumerate(imgs):
+        grid.paste(img, box=(i % cols * w, i // cols * h))
+
+    return grid
+
+@st.cache_data
+def read_video(path, transform=None, frames_num=9, window=30):
+    frames = []
+
+    cap = cv2.VideoCapture(path)
+
+    fps = int(cap.get(cv2.CAP_PROP_FPS))
+    length = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+    N = length // (frames_num)
+    current_frame = 1
+
+    for i in range(length):
+        ret, frame = cap.read(current_frame)
+
+        if ret and i == current_frame and len(frames) < frames_num:
+            size = 193, 193
+            frame = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
+            frame.thumbnail(size, Image.ANTIALIAS)
+
+            frames.append(frame)
+            current_frame += N
+
+    cap.release()
+
+    return frames
+
+# @st.cache_data
+def get_caption(model, tokenizer, prefix, prefix_length, prompt=''):
+    device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
+    prefix = prefix.to(device)
+
+    with torch.no_grad():
+        prefix_embed = model.clip_project(prefix).reshape(1, prefix_length, -1)
+
+        if prompt:
+            generated_text_prefix = generate2(model, tokenizer, prompt=prompt, embed=prefix_embed)
+        else:
+            generated_text_prefix = generate2(model, tokenizer, embed=prefix_embed)
+
+    return generated_text_prefix.replace('\n', ' ')
+
+# @st.cache_data
+def get_ans(model, tokenizer, clip_emb, prefix_length, prompt):
+    output = get_caption(model, tokenizer, clip_emb, prefix_length, prompt=prompt)
+    ans = output[len(prompt):].strip()
+
+    return {'answer': ans}

main.py

@@ -0,0 +1,115 @@
+import streamlit as st
+import pandas as pd
+import numpy as np
+import torch
+import clip
+import tempfile
+
+from tqdm import tqdm
+from transformers import GPT2Tokenizer
+from model import *
+from inference import *
+
+st.set_page_config(
+    page_title="Video Analysis AI",
+    page_icon="🕶️",
+)
+
+@st.cache_resource
+def load_model():
+    device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
+    clip_model, preprocess = clip.load("ViT-L/14@336px", device=device, jit=False)
+    tokenizer = GPT2Tokenizer.from_pretrained('sberbank-ai/rugpt3large_based_on_gpt2')
+
+    prefix_length = 50
+    model_path = 'transformer_clip_gpt-007.pt'
+    model = ClipCaptionModel('sberbank-ai/rugpt3small_based_on_gpt2', prefix_length=prefix_length)
+
+    model.load_state_dict(torch.load(model_path, map_location='cpu'))
+    model.to(device)
+    model.eval()
+
+    return model, clip_model, preprocess, tokenizer
+
+def _max_width_():
+    max_width_str = f"max-width: 1400px;"
+    st.markdown(
+        f"""
+    <style>
+    .reportview-container .main .block-container{{
+        {max_width_str}
+    }}
+    </style>
+    """,
+        unsafe_allow_html=True,
+    )
+
+_max_width_()
+
+
+def main():
+    model, clip_model, preprocess, tokenizer = load_model()
+    prefix_length = 50
+
+    st.title("🦾 Video Analysis for Education")
+    st.header("")
+
+    with st.sidebar.expander("ℹ️ - About application", expanded=True):
+        st.write(
+            """
+            - Upload the video
+            - Make a question about the content of the video
+            - Recieve answer according your question prompt
+    	    """
+        )
+
+
+    uploaded_file = st.file_uploader("📌 Upload video: ", ['.mp4'])
+
+    # if play_video:
+    #     video_bytes = uploaded_file.read()
+    #     st.video(video_bytes)
+
+    st.write("---")
+
+    question = st.text_input("❔ Enter question prompt: ", "")
+
+    
+    tfile = tempfile.NamedTemporaryFile(delete=False)
+    tfile.write(uploaded_file.read())
+
+    device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
+    val_embeddings = []
+    val_captions = []
+    result = ''
+    text = f'Question: {question}? Answer:'
+
+    #read video -> get_ans
+    video = read_video(tfile.name, transform=None, frames_num=4)
+
+    if len(video) > 0:
+        i = image_grid(video, 2, 2)
+        image = preprocess(i).unsqueeze(0).to(device)
+
+        with torch.no_grad():
+            prefix = clip_model.encode_image(image).to(device, dtype=torch.float32)
+
+        val_embeddings.append(prefix)
+        val_captions.append(text)
+
+    answers = []
+
+    for i in tqdm(range(len(val_embeddings))):
+        emb = val_embeddings[i]
+        caption = val_captions[i]
+
+        ans = get_ans(model, tokenizer, emb, prefix_length, caption)
+        answers.append(ans['answer'])
+
+    result = answers[0].split(' A: ')[0]
+    
+    res = st.text_input('✅ Answer to the question', result, disabled=False)
+
+
+if __name__ == '__main__':
+    main()

model.py

@@ -0,0 +1,222 @@
+import os
+import clip
+import numpy as np
+import pandas as pd
+import torch
+import transformers
+import torch
+import torch.nn as nn
+
+from enum import Enum
+from torch.nn import functional as nnf
+from typing import Tuple, Optional, Union
+from transformers import GPT2Tokenizer, GPT2LMHeadModel
+
+class MappingType(Enum):
+    MLP = 'mlp'
+    Transformer = 'transformer'
+
+class MlpTransformer(nn.Module):
+     def __init__(self, in_dim, h_dim, out_d: Optional[int] = None, act=nnf.relu, dropout=0.):
+         super().__init__()
+         out_d = out_d if out_d is not None else in_dim
+         self.fc1 = nn.Linear(in_dim, h_dim)
+         self.act = act
+         self.fc2 = nn.Linear(h_dim, out_d)
+         self.dropout = nn.Dropout(dropout)
+
+     def forward(self, x):
+         x = self.fc1(x)
+         x = self.act(x)
+         x = self.dropout(x)
+         x = self.fc2(x)
+         x = self.dropout(x)
+
+         return x
+
+class MLP(nn.Module):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.model(x)
+
+    def __init__(self, sizes: Tuple[int, ...], bias=True, act=nn.Tanh):
+        super(MLP, self).__init__()
+        layers = []
+        for i in range(len(sizes) - 1):
+            layers.append(nn.Linear(sizes[i], sizes[i + 1], bias=bias))
+            if i < len(sizes) - 2:
+                layers.append(act())
+
+        self.model = nn.Sequential(*layers)
+
+
+class MultiHeadAttention(nn.Module):
+    def __init__(self, dim_self, dim_ref, num_heads, bias=True, dropout=0.):
+        super().__init__()
+        self.num_heads = num_heads
+        head_dim = dim_self // num_heads
+        self.scale = head_dim ** -0.5
+        self.to_queries = nn.Linear(dim_self, dim_self, bias=bias)
+        self.to_keys_values = nn.Linear(dim_ref, dim_self * 2, bias=bias)
+        self.project = nn.Linear(dim_self, dim_self)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x, y=None, mask=None):
+        y = y if y is not None else x
+        b, n, c = x.shape
+        _, m, d = y.shape
+
+        queries = self.to_queries(x).reshape(b, n, self.num_heads, c // self.num_heads)
+        keys_values = self.to_keys_values(y).reshape(b, m, 2, self.num_heads, c // self.num_heads)
+        keys, values = keys_values[:, :, 0], keys_values[:, :, 1]
+        attention = torch.einsum('bnhd,bmhd->bnmh', queries, keys) * self.scale
+
+        if mask is not None:
+            if mask.dim() == 2:
+                mask = mask.unsqueeze(1)
+            attention = attention.masked_fill(mask.unsqueeze(3), float("-inf"))
+
+        attention = attention.softmax(dim=2)
+
+        out = torch.einsum('bnmh,bmhd->bnhd', attention, values).reshape(b, n, c)
+        out = self.project(out)
+
+        return out, attention
+
+
+class TransformerLayer(nn.Module):
+    def forward_with_attention(self, x, y=None, mask=None):
+        x_, attention = self.attn(self.norm1(x), y, mask)
+        x = x + x_
+        x = x + self.mlp(self.norm2(x))
+
+        return x, attention
+
+    def forward(self, x, y=None, mask=None):
+        x = x + self.attn(self.norm1(x), y, mask)[0]
+        x = x + self.mlp(self.norm2(x))
+
+        return x
+
+    def __init__(self, dim_self, dim_ref, num_heads, mlp_ratio=4., bias=False, dropout=0., act=nnf.relu,
+                 norm_layer: nn.Module = nn.LayerNorm):
+        super().__init__()
+        self.norm1 = norm_layer(dim_self)
+        self.attn = MultiHeadAttention(dim_self, dim_ref, num_heads, bias=bias, dropout=dropout)
+        self.norm2 = norm_layer(dim_self)
+        self.mlp = MlpTransformer(dim_self, int(dim_self * mlp_ratio), act=act, dropout=dropout)
+
+
+class Transformer(nn.Module):
+    def forward_with_attention(self, x, y=None, mask=None):
+        attentions = []
+        for layer in self.layers:
+            x, att = layer.forward_with_attention(x, y, mask)
+            attentions.append(att)
+
+        return x, attentions
+
+    def forward(self, x, y=None, mask=None):
+        for i, layer in enumerate(self.layers):
+            if i % 2 == 0 and self.enc_dec:
+                x = layer(x, y)
+            elif self.enc_dec:
+                x = layer(x, x, mask)
+            else:
+                x = layer(x, y, mask)
+        return x
+
+    def __init__(self, dim_self: int, num_heads: int, num_layers: int, dim_ref: Optional[int] = None,
+                 mlp_ratio: float = 2., act=nnf.relu, norm_layer: nn.Module = nn.LayerNorm, enc_dec: bool = False):
+        super(Transformer, self).__init__()
+        dim_ref = dim_ref if dim_ref is not None else dim_self
+        self.enc_dec = enc_dec
+
+        if enc_dec:
+            num_layers = num_layers * 2
+
+        layers = []
+
+        for i in range(num_layers):
+            if i % 2 == 0 and enc_dec:
+                layers.append(TransformerLayer(dim_self, dim_ref, num_heads, mlp_ratio, act=act, norm_layer=norm_layer))
+            elif enc_dec:
+                layers.append(TransformerLayer(dim_self, dim_self, num_heads, mlp_ratio, act=act, norm_layer=norm_layer))
+            else:
+                layers.append(TransformerLayer(dim_self, dim_ref, num_heads, mlp_ratio, act=act, norm_layer=norm_layer))
+
+        self.layers = nn.ModuleList(layers)
+
+
+class TransformerMapper(nn.Module):
+    def forward(self, x):
+        x = self.linear(x).view(x.shape[0], self.clip_length, -1)
+        prefix = self.prefix_const.unsqueeze(0).expand(x.shape[0], *self.prefix_const.shape)
+        prefix = torch.cat((x, prefix), dim=1)
+        out = self.transformer(prefix)[:, self.clip_length:]
+
+        return out
+
+    def __init__(self, dim_clip: int, dim_embedding: int, prefix_length: int, clip_length: int, num_layers: int = 8):
+        super(TransformerMapper, self).__init__()
+        self.clip_length = clip_length
+        self.transformer = Transformer(dim_embedding, 8, num_layers)
+        self.linear = nn.Linear(dim_clip, clip_length * dim_embedding)
+        self.prefix_const = nn.Parameter(torch.randn(prefix_length, dim_embedding), requires_grad=True)
+
+class MLP(nn.Module):
+    def __init__(self, sizes: Tuple[int, ...], bias=True, act=nn.Tanh):
+        super(MLP, self).__init__()
+        layers = []
+        for i in range(len(sizes) - 1):
+            layers.append(nn.Linear(sizes[i], sizes[i + 1], bias=bias))
+            if i < len(sizes) - 2:
+                layers.append(act())
+        self.model = nn.Sequential(*layers)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.model(x)
+
+
+class ClipCaptionModel(nn.Module):
+    def __init__(self, gpt, prefix_length: int, prefix_size: int = 768):
+        super(ClipCaptionModel, self).__init__()
+
+        self.prefix_length = prefix_length
+        clip_length = prefix_length
+        num_layers = 8
+
+        self.gpt = GPT2LMHeadModel.from_pretrained(gpt)
+        # self.gpt = freeze(self.gpt)
+        self.gpt_embedding_size = self.gpt.transformer.wte.weight.shape[1]
+        self.clip_project = TransformerMapper(prefix_size, self.gpt_embedding_size, prefix_length,
+                                              clip_length, num_layers)
+
+    def get_dummy_token(self, batch_size: int, device: torch.device) -> torch.Tensor:
+        return torch.zeros(batch_size, self.prefix_length, dtype=torch.int64, device=device)
+
+    def forward(self, tokens: torch.Tensor, prefix: torch.Tensor,
+                mask: Optional[torch.Tensor] = None,
+                labels: Optional[torch.Tensor] = None):
+        embedding_text = self.gpt.transformer.wte(tokens)
+        prefix_projections = self.clip_project(prefix).view(-1, self.prefix_length, self.gpt_embedding_size)
+
+        embedding_cat = torch.cat((prefix_projections, embedding_text), dim=1)
+
+        if labels is not None:
+            dummy_token = self.get_dummy_token(tokens.shape[0], tokens.device)
+            labels = torch.cat((dummy_token, tokens), dim=1)
+
+        out = self.gpt(inputs_embeds=embedding_cat, labels=labels, attention_mask=mask)
+
+        return out
+
+
+class ClipCaptionPrefix(ClipCaptionModel):
+    def parameters(self, recurse: bool = True):
+        return self.clip_project.parameters()
+
+    def train(self, mode: bool = True):
+        super(ClipCaptionPrefix, self).train(mode)
+        self.gpt.eval()
+
+        return self

transformer_clip_gpt-007.pt

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