Hugging Face's logo

Spaces:
Gradio-Blocks
/
uniformer_video_demo
Runtime error

App Files Community
uniformer_video_demo
File size: 4,242 Bytes 
0b70dc1
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
7312740
 
 
 
 
 
 
 
 
d4ed5f1
7312740
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
 
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
from huggingface_hub import hf_hub_download

# 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")
model_path = hf_hub_download(repo_id="Andy1621/uniformer", filename="uniformer_small_k400_16x8.pth")
# Pick a pretrained model 
model = uniformer_small()
# state_dict = torch.load('d5fd7b0c49ee6a5422ef5d0c884d962c742003bfbd900747485eb99fa269d0db', map_location='cpu')
state_dict = torch.load(model_path, map_location='cpu')
model.load_state_dict(state_dict)

# 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(400)}
    

def set_example_video(example: list) -> dict:
    return gr.Video.update(value=example[0])


demo = gr.Blocks()
with demo:
    gr.Markdown(
        """
        # UniFormer-S
        Gradio demo for <a href='https://github.com/Sense-X/UniFormer' target='_blank'>UniFormer</a>: To use it, simply upload your video, or click one of the examples to load them. Read more at the links below.
        """
    )

    with gr.Box():
        with gr.Row():
                with gr.Column():
                    with gr.Row():
                        input_video = gr.Video(label='Input Video')
                    with gr.Row():
                        submit_button = gr.Button('Submit')
                with gr.Column():
                    label = gr.Label(num_top_classes=5)
        with gr.Row():
            example_videos = gr.Dataset(components=[input_video], samples=[['hitting_baseball.mp4'], ['hoverboarding.mp4'], ['yoga.mp4']])

    gr.Markdown(
        """
        <p style='text-align: center'><a href='https://arxiv.org/abs/2201.04676' target='_blank'>[ICLR2022] UniFormer: Unified Transformer for Efficient Spatiotemporal Representation Learning</a> | <a href='https://github.com/Sense-X/UniFormer' target='_blank'>Github Repo</a></p>
        """
    )

    submit_button.click(fn=inference, inputs=input_video, outputs=label)
    example_videos.click(fn=set_example_video, inputs=example_videos, outputs=example_videos.components)

demo.launch(enable_queue=True)