Update app.py

app.py

@@ -1,24 +1,20 @@
-import gradio as gr
-import torch
-import numpy as np
 import cv2
+import numpy as np
+import torch
+from transformers import AutoImageProcessor, AutoModelForVideoClassification
+import gradio as gr
 from PIL import Image
-from transformers import AutoImageProcessor, AutoModelForVideoClassification, ViTForImageClassification
 
-# Load video model and processor
-video_processor = AutoImageProcessor.from_pretrained("facebook/timesformer-base-finetuned-k400")
-video_model = AutoModelForVideoClassification.from_pretrained("facebook/timesformer-base-finetuned-k400")
-
-# Load image model and processor
-image_processor = AutoImageProcessor.from_pretrained("google/vit-base-patch16-224")
-image_model = ViTForImageClassification.from_pretrained("google/vit-base-patch16-224")
+# Use AutoModelForVideoClassification which is generic and supports TimeSformer
+processor = AutoImageProcessor.from_pretrained("facebook/timesformer-base-finetuned-k400")
+model = AutoModelForVideoClassification.from_pretrained("facebook/timesformer-base-finetuned-k400")
 
 def extract_frames(video_path, num_frames=8):
     cap = cv2.VideoCapture(video_path)
     total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
     frames = []
-    frame_indices = np.linspace(0, total_frames - 1, num_frames).astype(int)
-    for idx in frame_indices:
+    frame_idxs = np.linspace(0, total_frames - 1, num_frames).astype(int)
+    for idx in frame_idxs:
         cap.set(cv2.CAP_PROP_POS_FRAMES, idx)
         ret, frame = cap.read()
         if not ret:
@@ -29,43 +25,35 @@ def extract_frames(video_path, num_frames=8):
     cap.release()
     return frames
 
+# Image classification for fallback
+image_processor = AutoImageProcessor.from_pretrained("google/vit-base-patch16-224")
+image_model = torch.hub.load('pytorch/vision:v0.14.0', 'vit_b_16', pretrained=True)
+image_model.eval()
+
 def predict(file):
     if file is None:
-        return "Please upload a file.", None
+        return "Please upload a file."
 
-    path = file.name
+    filepath = file.name
     video_exts = [".mp4", ".avi", ".mov", ".mkv"]
-
-    if any(path.lower().endswith(ext) for ext in video_exts):
-        # Video
-        frames = extract_frames(path)
-        inputs = video_processor(frames, return_tensors="pt")
+    if any(filepath.lower().endswith(ext) for ext in video_exts):
+        frames = extract_frames(filepath)
+        inputs = processor(frames, return_tensors="pt")
         with torch.no_grad():
-            outputs = video_model(**inputs)
+            outputs = model(**inputs)
         logits = outputs.logits
-        pred_id = torch.argmax(logits, dim=-1).item()
-        label = video_model.config.id2label[pred_id]
-        confidence = torch.nn.functional.softmax(logits, dim=-1)[0, pred_id].item()
-        return f"Video Prediction: {label} (Confidence: {confidence:.2f})", path
+        pred_idx = torch.argmax(logits, dim=-1).item()
+        label = model.config.id2label[pred_idx]
+        prob = torch.nn.functional.softmax(logits, dim=-1)[0, pred_idx].item()
+        return f"Video Prediction: {label} (Confidence: {prob:.2f})"
     else:
-        # Image
-        image = Image.open(path).convert("RGB")
-        inputs = image_processor(images=image, return_tensors="pt")
+        image = Image.open(filepath).convert("RGB")
+        inputs = image_processor(image, return_tensors="pt")
         with torch.no_grad():
             outputs = image_model(**inputs)
         logits = outputs.logits
-        pred_id = torch.argmax(logits, dim=-1).item()
-        label = image_model.config.id2label[pred_id]
-        confidence = torch.nn.functional.softmax(logits, dim=-1)[0, pred_id].item()
-        return f"Image Prediction: {label} (Confidence: {confidence:.2f})", None
-
-iface = gr.Interface(
-    fn=predict,
-    inputs=gr.File(file_types=[".mp4", ".avi", ".mov", ".mkv", ".jpg", ".jpeg", ".png"]),
-    outputs=[gr.Textbox(label="Prediction"), gr.Video(label="Uploaded Video")],
-    title="Image and Video Classification",
-    description="Upload an image or video. Video predictions use manual frame extraction to avoid decord dependency."
-)
+        pred_idx = torch.argmax(logits, dim=-1).item()
+        return f"Image Prediction index: {pred_idx} (Raw model, no labels)"
 
-if __name__ == "__main__":
-    iface.launch()
+iface = gr.Interface(fn=predict, inputs=gr.File(file_types=[".mp4", ".avi", ".mov", ".mkv", ".jpg", ".png"]), outputs="text")
+iface.launch()