Spaces:
Runtime error
Runtime error
File size: 3,303 Bytes
b757ffe 84a6f17 |
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 |
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
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)
# The model expects inputs of shape: B x C x T x H x W
TC, H, W = torch_imgs.shape
torch_imgs = torch_imgs.reshape(1, TC//3, 3, H, W).permute(0, 2, 1, 3, 4)
return torch_imgs
def inference(video):
vid = load_video(video)
prediction = model(vid)
prediction = F.softmax(prediction, dim=1).flatten()
return {kinetics_id_to_classname[str(i)]: float(prediction[i]) for i in range(400)}
# Device on which to run the model
# Set to cuda to load on GPU
device = "cpu"
model_path = hf_hub_download(repo_id="Sense-X/uniformer_video", filename="uniformer_small_k400_16x8.pth")
# Pick a pretrained model
model = uniformer_small()
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
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.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>"
gr.Interface(
inference, inputs, outputs=label,
title=title, description=description, article=article,
examples=[['hitting_baseball.mp4'], ['hoverboarding.mp4'], ['yoga.mp4']]
).launch(enable_queue=True)
|