Update app.py

app.py

@@ -3,13 +3,11 @@ import torch
 import numpy as np
 from transformers import AutoProcessor, AutoModel
 from PIL import Image
-from decord import VideoReader, cpu
+from decord import VideoReader, cpu, gpu
 import cv2
 
 print(f"Is CUDA available: {torch.cuda.is_available()}")
-# True
 print(f"CUDA device: {torch.cuda.get_device_name(torch.cuda.current_device())}")
-# Tesla T4
 
 def sample_uniform_frame_indices(clip_len, seg_len):
     if seg_len < clip_len:
@@ -22,22 +20,10 @@ def sample_uniform_frame_indices(clip_len, seg_len):
     return np.array(indices).astype(np.int64)
 
 def read_video_decord(file_path, indices):
-    vr = VideoReader(file_path, num_threads=1, ctx=cpu(0))
+    vr = VideoReader(file_path, num_threads=1, ctx=gpu(0) if torch.cuda.is_available() else cpu(0))
     video = vr.get_batch(indices).asnumpy()
     return video
 
-def read_video_opencv(file_path, indices):
-    vidcap = cv2.VideoCapture(file_path)
-    frames = []
-    for idx in indices:
-        vidcap.set(cv2.CAP_PROP_POS_FRAMES, idx)
-        success, image = vidcap.read()
-        if success:
-            # Convert BGR to RGB
-            frames.append(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
-    return frames
-
-
 def concatenate_frames(frames, clip_len):
     layout = {
         32: (4, 8),
@@ -63,26 +49,30 @@ def model_interface(uploaded_video, model_choice, activity):
         "microsoft/xclip-base-patch32-16-frames": 16,
         "microsoft/xclip-base-patch32": 8
     }.get(model_choice, 32)
+
     indices = sample_uniform_frame_indices(clip_len, seg_len=len(VideoReader(uploaded_video)))
-    video = read_video_opencv(uploaded_video, indices)
+    video = read_video_decord(uploaded_video, indices)
     concatenated_image = concatenate_frames(video, clip_len)
 
-    # Appending "other" to the list of activities
+    # Convert list of numpy arrays to a single numpy ndarray
+    video_np = np.array(video)
+
     activities_list = [activity, "other"]
     processor = AutoProcessor.from_pretrained(model_choice)
-    model = AutoModel.from_pretrained(model_choice)
-    
+    model = AutoModel.from_pretrained(model_choice).to('cuda')
+
     inputs = processor(
         text=activities_list,
-        videos=list(video),
+        videos=video_np,  # Use the ndarray instead of the list
         return_tensors="pt",
         padding=True,
     )
+    inputs = {name: tensor.to('cuda') for name, tensor in inputs.items()}
 
     with torch.no_grad():
         outputs = model(**inputs)
 
-    logits_per_video = outputs.logits_per_video
+    logits_per_video = outputs.logits_per_video.cpu()
     probs = logits_per_video.softmax(dim=1)
 
     results_probs = []
@@ -98,28 +88,4 @@ def model_interface(uploaded_video, model_choice, activity):
     likely_label = activities_list[max_prob_index]
     likely_probability = float(probs[0][max_prob_index]) * 100
 
-    return concatenated_image, results_probs, results_logits, [ likely_label , likely_probability ]
-
-iface = gr.Interface(
-    fn=model_interface,
-    inputs=[
-        gr.components.Video(label="Upload a video file"),
-        gr.components.Dropdown(choices=[
-            "microsoft/xclip-base-patch16-zero-shot",
-            "microsoft/xclip-base-patch32-16-frames",
-            "microsoft/xclip-base-patch32"
-        ], label="Model Choice"),
-        gr.components.Textbox(default="dancing", label="Desired Activity to Recognize"),
-    ],
-    outputs=[
-        gr.components.Image(type="pil", label="Sampled Frames"),
-        gr.components.Textbox(type="text", label="Probabilities"),
-        gr.components.Textbox(type="text", label="Raw Scores"),
-        gr.components.Textbox(type="text", label="Top Prediction")
-    ],
-    live=False
-)
-
-
-
-iface.launch()
+    return concatenated_image, results_probs, results_logits, [likely_label, likely_probability]