Create app.py

app.py

@@ -0,0 +1,101 @@
+import os
+
+import torch
+import numpy as np
+import torch.nn.functional as F
+import torchvision.transforms as T
+from PIL import Image
+from decord import VideoReader
+from decord import cpu
+from uniformer import uniformer_small
+from kinetics_class_index import kinetics_classnames
+from transforms import (
+    GroupNormalize, GroupScale, GroupCenterCrop, 
+    Stack, ToTorchFormatTensor
+)
+
+import gradio as gr
+
+# Device on which to run the model
+# Set to cuda to load on GPU
+device = "cpu"
+os.system("wget https://cdn-lfs.huggingface.co/Andy1621/uniformer/d5fd7b0c49ee6a5422ef5d0c884d962c742003bfbd900747485eb99fa269d0db")
+# Pick a pretrained model 
+model = uniformer_small()
+state_dict = torch.load('d5fd7b0c49ee6a5422ef5d0c884d962c742003bfbd900747485eb99fa269d0db', map_location='cpu')
+model.load_state_dict(state_dict['model'])
+
+# Set to eval mode and move to desired device
+model = model.to(device)
+model = model.eval()
+
+# Create an id to label name mapping
+kinetics_id_to_classname = {}
+for k, v in kinetics_classnames.items():
+    kinetics_id_to_classname[k] = v
+
+
+def get_index(num_frames, num_segments=16, dense_sample_rate=8):
+    sample_range = num_segments * dense_sample_rate
+    sample_pos = max(1, 1 + num_frames - sample_range)
+    t_stride = dense_sample_rate
+    start_idx = 0 if sample_pos == 1 else sample_pos // 2
+    offsets = np.array([
+        (idx * t_stride + start_idx) %
+        num_frames for idx in range(num_segments)
+    ])
+    return offsets + 1
+
+
+def load_video(video_path):
+    vr = VideoReader(video_path, ctx=cpu(0))
+    num_frames = len(vr)
+    frame_indices = get_index(num_frames, 16, 16)
+
+    # transform
+    crop_size = 224
+    scale_size = 256
+    input_mean = [0.485, 0.456, 0.406]
+    input_std = [0.229, 0.224, 0.225]
+
+    transform = T.Compose([
+        GroupScale(int(scale_size)),
+        GroupCenterCrop(crop_size),
+        Stack(),
+        ToTorchFormatTensor(),
+        GroupNormalize(input_mean, input_std) 
+    ])
+
+    images_group = list()
+    for frame_index in frame_indices:
+        img = Image.fromarray(vr[frame_index].asnumpy())
+        images_group.append(img)
+    torch_imgs = transform(images_group)
+    return torch_imgs
+    
+
+def inference(video):
+    vid = load_video(video)
+    
+    # The model expects inputs of shape: B x C x H x W
+    TC, H, W = vid.shape
+    inputs = vid.reshape(1, TC//3, 3, H, W).permute(0, 2, 1, 3, 4)
+    
+    prediction = model(inputs)
+    prediction = F.softmax(prediction, dim=1).flatten()
+
+    return {kinetics_id_to_classname[str(i)]: float(prediction[i]) for i in range(1000)}
+    
+
+inputs = gr.inputs.Video()
+label = gr.outputs.Label(num_top_classes=5)
+
+title = "UniFormer-S"
+description = "Gradio demo for UniFormer: To use it, simply upload your video, or click one of the examples to load them. Read more at the links below."
+article = "<p style='text-align: center'><a href='https://arxiv.org/abs/2201.09450' target='_blank'>UniFormer: Unifying Convolution and Self-attention for Visual Recognition</a> | <a href='https://github.com/Sense-X/UniFormer' target='_blank'>Github Repo</a></p>"
+
+gr.Interface(
+    inference, inputs, outputs=label, 
+    title=title, description=description, article=article, 
+    examples=[['hitting_baseball.mp4'], ['hoverboarding.mp4'], ['yoga.mp4']]
+    ).launch(enable_queue=True, cache_examples=True)