Update app.py

app.py

@@ -1,21 +1,17 @@
 import gradio as gr
 import torch
+import numpy as np
+from transformers import AutoProcessor, AutoModel
+from PIL import Image
+from decord import VideoReader, cpu
+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
 
-import numpy as np
-from transformers import AutoProcessor, AutoModel
-from PIL import Image
-from decord import VideoReader, gpu
-
 def sample_uniform_frame_indices(clip_len, seg_len):
-    """
-    Samples `clip_len` uniformly spaced frame indices from a video of length `seg_len`.
-    Handles edge cases where `seg_len` might be less than `clip_len`.
-    """
     if seg_len < clip_len:
         repeat_factor = np.ceil(clip_len / seg_len).astype(int)
         indices = np.arange(seg_len).tolist() * repeat_factor
@@ -23,24 +19,32 @@ def sample_uniform_frame_indices(clip_len, seg_len):
     else:
         spacing = seg_len // clip_len
         indices = [i * spacing for i in range(clip_len)]
-    
     return np.array(indices).astype(np.int64)
 
 def read_video_decord(file_path, indices):
-    vr = VideoReader(file_path, num_threads=1, ctx=gpu(0))
+    vr = VideoReader(file_path, num_threads=1, ctx=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 np.array(frames)
+
+
 def concatenate_frames(frames, clip_len):
-    assert len(frames) == clip_len, f"The function expects {clip_len} frames as input."
-    
     layout = {
         32: (4, 8),
         16: (4, 4),
         8:  (2, 4)
     }
     rows, cols = layout[clip_len]
-    
     combined_image = Image.new('RGB', (frames[0].shape[1]*cols, frames[0].shape[0]*rows))
     frame_iter = iter(frames)
     y_offset = 0
@@ -51,26 +55,22 @@ def concatenate_frames(frames, clip_len):
             combined_image.paste(img, (x_offset, y_offset))
             x_offset += frames[0].shape[1]
         y_offset += frames[0].shape[0]
-        
     return combined_image
 
-
-def model_interface(uploaded_video, model_choice, activities):
+def model_interface(uploaded_video, model_choice, activity):
     clip_len = {
         "microsoft/xclip-base-patch16-zero-shot": 32,
         "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_decord(uploaded_video, indices)
-    concatenated_image = concatenate_frames(video, clip_len)  # Passed clip_len as argument
+    video = read_video_opencv(uploaded_video, indices)
+    concatenated_image = concatenate_frames(video, clip_len)
 
+    # Appending "other" to the list of activities
+    activities_list = [activity, "other"]
     processor = AutoProcessor.from_pretrained(model_choice)
     model = AutoModel.from_pretrained(model_choice)
-    model = model.to("cuda")
-    
-    activities_list = activities.split(",")
     inputs = processor(
         text=activities_list,
         videos=list(video),
@@ -86,19 +86,18 @@ def model_interface(uploaded_video, model_choice, activities):
 
     results_probs = []
     results_logits = []
+    max_prob_index = torch.argmax(probs[0]).item()
     for i in range(len(activities_list)):
-        activity = activities_list[i]
+        current_activity = activities_list[i]
         prob = float(probs[0][i])
         logit = float(logits_per_video[0][i])
-        results_probs.append((activity, f"Probability: {prob * 100:.2f}%"))
-        results_logits.append((activity, f"Raw Score: {logit:.2f}"))
+        results_probs.append((current_activity, f"Probability: {prob * 100:.2f}%"))
+        results_logits.append((current_activity, f"Raw Score: {logit:.2f}"))
 
-    # Retrieve most likely predicted label and its probability
-    max_prob_idx = probs[0].argmax().item()
-    most_likely_activity = activities_list[max_prob_idx]
-    most_likely_prob = float(probs[0][max_prob_idx])
+    likely_label = activities_list[max_prob_index]
+    likely_probability = float(probs[0][max_prob_index]) * 100
 
-    return concatenated_image, results_probs, results_logits, (most_likely_activity, f"Probability: {most_likely_prob * 100:.2f}%")
+    return concatenated_image, results_probs, results_logits, [ likely_label , likely_probability ]
 
 iface = gr.Interface(
     fn=model_interface,
@@ -109,15 +108,15 @@ iface = gr.Interface(
             "microsoft/xclip-base-patch32-16-frames",
             "microsoft/xclip-base-patch32"
         ], label="Model Choice"),
-        gr.components.Textbox(lines=4, label="Enter activities (comma-separated)"),
+        gr.components.Textbox(default="dancing", label="Desired Activity to Recognize"),
     ],
     outputs=[
-        gr.components.Image(type="pil", label="sampled frames"),
+        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="Most Likely Prediction")
+        gr.components.Textbox(type="text", label="Top Prediction")
     ],
     live=False
 )
 
-iface.launch()
+iface.launch()