hocheewai's picture
Update app.py
d79cbd2
import cv2
import gradio as gr
import imutils
import numpy as np
import torch
from pytorchvideo.transforms import (
ApplyTransformToKey,
Normalize,
RandomShortSideScale,
RemoveKey,
ShortSideScale,
UniformTemporalSubsample,
)
from torchvision.transforms import (
Compose,
Lambda,
RandomCrop,
RandomHorizontalFlip,
Resize,
)
from transformers import VideoMAEFeatureExtractor, VideoMAEForVideoClassification
MODEL_CKPT = "hocheewai/videomae-base-finetuned-ucf101-subset"
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
MODEL = VideoMAEForVideoClassification.from_pretrained(MODEL_CKPT).to(DEVICE)
PROCESSOR = VideoMAEFeatureExtractor.from_pretrained(MODEL_CKPT)
RESIZE_TO = PROCESSOR.size["shortest_edge"]
NUM_FRAMES_TO_SAMPLE = MODEL.config.num_frames
IMAGE_STATS = {"image_mean": [0.485, 0.456, 0.406], "image_std": [0.229, 0.224, 0.225]}
VAL_TRANSFORMS = Compose(
[
UniformTemporalSubsample(NUM_FRAMES_TO_SAMPLE),
Lambda(lambda x: x / 255.0),
Normalize(IMAGE_STATS["image_mean"], IMAGE_STATS["image_std"]),
Resize((RESIZE_TO, RESIZE_TO)),
]
)
LABELS = list(MODEL.config.label2id.keys())
def parse_video(video_file):
"""A utility to parse the input videos.
Reference: https://pyimagesearch.com/2018/11/12/yolo-object-detection-with-opencv/
"""
vs = cv2.VideoCapture(video_file)
# try to determine the total number of frames in the video file
try:
prop = (
cv2.cv.CV_CAP_PROP_FRAME_COUNT
if imutils.is_cv2()
else cv2.CAP_PROP_FRAME_COUNT
)
total = int(vs.get(prop))
print("[INFO] {} total frames in video".format(total))
# an error occurred while trying to determine the total
# number of frames in the video file
except:
print("[INFO] could not determine # of frames in video")
print("[INFO] no approx. completion time can be provided")
total = -1
frames = []
# loop over frames from the video file stream
while True:
# read the next frame from the file
(grabbed, frame) = vs.read()
if frame is not None:
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
frames.append(frame)
# if the frame was not grabbed, then we have reached the end
# of the stream
if not grabbed:
break
return frames
def preprocess_video(frames: list):
"""Utility to apply preprocessing transformations to a video tensor."""
# Each frame in the `frames` list has the shape: (height, width, num_channels).
# Collated together the `frames` has the the shape: (num_frames, height, width, num_channels).
# So, after converting the `frames` list to a torch tensor, we permute the shape
# such that it becomes (num_channels, num_frames, height, width) to make
# the shape compatible with the preprocessing transformations. After applying the
# preprocessing chain, we permute the shape to (num_frames, num_channels, height, width)
# to make it compatible with the model. Finally, we add a batch dimension so that our video
# classification model can operate on it.
video_tensor = torch.tensor(np.array(frames).astype(frames[0].dtype))
video_tensor = video_tensor.permute(
3, 0, 1, 2
) # (num_channels, num_frames, height, width)
video_tensor_pp = VAL_TRANSFORMS(video_tensor)
video_tensor_pp = video_tensor_pp.permute(
1, 0, 2, 3
) # (num_frames, num_channels, height, width)
video_tensor_pp = video_tensor_pp.unsqueeze(0)
return video_tensor_pp.to(DEVICE)
def infer(video_file):
frames = parse_video(video_file)
video_tensor = preprocess_video(frames)
inputs = {"pixel_values": video_tensor}
# forward pass
with torch.no_grad():
outputs = MODEL(**inputs)
logits = outputs.logits
softmax_scores = torch.nn.functional.softmax(logits, dim=-1).squeeze(0)
confidences = {LABELS[i]: float(softmax_scores[i]) for i in range(len(LABELS))}
return confidences
gr.Interface(
fn=infer,
inputs=gr.Video(type="file"),
outputs=gr.Label(num_top_classes=3),
examples=[
["examples/babycrawling.mp4"],
["examples/baseball.mp4"],
["examples/balancebeam.mp4"],
],
title="VideoMAE fine-tuned on a subset of UCF-101",
description=(
"Gradio demo for VideoMAE for video classification. To use it, simply upload your video or click one of the"
" examples to load them. Read more at the links below."
),
article=(
"<div style='text-align: center;'><a href='https://huggingface.co/docs/transformers/model_doc/videomae' target='_blank'>VideoMAE</a>"
" <center><a href='https://huggingface.co/sayakpaul/videomae-base-finetuned-kinetics-finetuned-ucf101-subset' target='_blank'>Fine-tuned Model</a></center></div>"
),
allow_flagging=False,
allow_screenshot=False,
).launch()