thejasbh/I3D_Sign_Language_Classification

1import torch2import cv23import videotransforms4import numpy as np5import gradio as gr6from einops import rearrange7from torchvision import transforms8from pytorch_i3d import InceptionI3d9 10 11def preprocess(vidpath):12    # Fetch video13    cap = cv2.VideoCapture(vidpath)14 15    frames = []16    cap.set(cv2.CAP_PROP_POS_FRAMES, 0)17    num = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))18    19    # Extract frames from video20    for _ in range(num):21        _, img = cap.read()22        23        # Skip NoneType frames24        if img is None:25            continue26 27        # Resize if (w,h) < (226,226)28        w, h, c = img.shape29        if w < 226 or h < 226:30            d = 226. - min(w, h)31            sc = 1 + d / min(w, h)32            img = cv2.resize(img, dsize=(0, 0), fx=sc, fy=sc)33 34        # Normalize35        img = (img / 255.) * 2 - 136 37        frames.append(img)38    39    frames = torch.Tensor(np.asarray(frames, dtype=np.float32))40    41    # Transform tensor and reshape to (1, c, t ,w, h)42    transform = transforms.Compose([videotransforms.CenterCrop(224)])43    frames = transform(frames)44    frames = rearrange(frames, 't w h c-> 1 c t w h')45 46    return frames47 48def classify(video,dataset='WLASL100'):49    to_load = {50        'WLASL100':{'logits':100,'path':'weights/asl100/FINAL_nslt_100_iters=896_top1=65.89_top5=84.11_top10=89.92.pt'},51        'WLASL2000':{'logits':2000,'path':'weights/asl2000/FINAL_nslt_2000_iters=5104_top1=32.48_top5=57.31_top10=66.31.pt'}52        }53 54    # Preprocess video55    input = preprocess(video)56 57    # Load model58    model = InceptionI3d()59    model.load_state_dict(torch.load('weights/rgb_imagenet.pt',map_location=torch.device('cpu')))60    model.replace_logits(to_load[dataset]['logits'])61    model.load_state_dict(torch.load(to_load[dataset]['path'],map_location=torch.device('cpu')))62 63    # Run on cpu. Spaces environment is limited to CPU for free users. 64    model.cpu()65 66    # Evaluation mode67    model.eval()68 69    with torch.no_grad(): # Disable gradient computation70        per_frame_logits = model(input) # Inference71    72    per_frame_logits.cpu()73    model.cpu()74 75    # Load predictions76    predictions = rearrange(per_frame_logits,'1 j k -> j k')77    predictions = torch.mean(predictions, dim = 1)78 79    # Fetch top 10 predictions80    _, index = torch.topk(predictions,10)81    index = index.cpu().numpy()82 83    # Load labels 84    with open('wlasl_class_list.txt') as f:85        idx2label = dict()86        for line in f:87            idx2label[int(line.split()[0])]=line.split()[1]88    89    # Get probabilities90    predictions = torch.nn.functional.softmax(predictions, dim=0).cpu().numpy()91 92    # Return dict {label:pred}93    return {idx2label[i]:float(predictions[i]) for i in index}94 95# Gradio App config96title = "I3D Sign Language Recognition"97description =   "Gradio demo of word-level sign language classification using I3D model pretrained on the WLASL video dataset. " \98                "WLASL is a large-scale dataset containing more than 2000 words in American Sign Language. " \99                "Examples used in the demo are videos from the the test subset. "  \100                "Note that WLASL100 contains 100 words while WLASL2000 contains 2000."101examples = [102        ['videos/no.mp4','WLASL100'],103        ['videos/all.mp4','WLASL100'],104        ['videos/before.mp4','WLASL100'],105        ['videos/blue.mp4','WLASL2000'],106        ['videos/white.mp4','WLASL2000'],107        ['videos/accident2.mp4','WLASL2000']108    ]109 110article =   "NOTE: This is not the official demonstration of the I3D sign language classification on the WLASL dataset. "\111            "More information about the WLASL dataset and pretrained I3D models can be found <a href=https://github.com/dxli94/WLASL>here</a>."112 113# Gradio App interface114gr.Interface(   fn=classify,115                inputs=[gr.inputs.Video(label="Video (*.mp4)"),gr.inputs.Radio(choices=['WLASL100','WLASL2000'], default='WLASL100', label='Trained on:')], 116                outputs=[gr.outputs.Label(num_top_classes=5, label='Top 5 Predictions')],117                allow_flagging="never",118                title=title, 119                description=description, 120                examples=examples,121                article=article).launch()122