SakuraD's picture
init
469404a
raw
history blame
4.11 kB
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 uniformerv2 import uniformerv2_b16
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
class Uniformerv2(nn.Module):
def __init__(self, model):
super().__init__()
self.backbone = model
def forward(self, x):
return self.backbone(x)
# Device on which to run the model
# Set to cuda to load on GPU
device = "cpu"
model_path = hf_hub_download(repo_id="Andy1621/uniformerv2", filename="k400+k710_uniformerv2_b16_8x224.pyth")
# Pick a pretrained model
model = Uniformerv2(uniformerv2_b16(pretrained=False, t_size=8, no_lmhra=True, temporal_downsample=False))
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=8):
seg_size = float(num_frames - 1) / num_segments
start = int(seg_size / 2)
offsets = np.array([
start + int(np.round(seg_size * idx)) for idx in range(num_segments)
])
return offsets
def load_video(video_path):
vr = VideoReader(video_path, ctx=cpu(0))
num_frames = len(vr)
frame_indices = get_index(num_frames, 8)
# 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(
"""
# UniFormerV2-B
Gradio demo for <a href='https://github.com/OpenGVLab/UniFormerV2' target='_blank'>UniFormerV2</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/2211.09552' target='_blank'>[Arxiv] UniFormerV2: Spatiotemporal Learning by Arming Image ViTs with Video UniFormer</a> | <a href='https://github.com/OpenGVLab/UniFormerV2' 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)