Update app.py

app.py

@@ -15,6 +15,14 @@ model.eval()
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model = model.to(device)
 
+# Map index to label
+id2label = model.config.id2label
+label2id = {v: k for k, v in id2label.items()}
+
+# Just pick highest index as "AI" for now (based on SMOGY default config)
+ai_class_index = max(id2label.keys())
+ai_label_name = id2label[ai_class_index]
+
 @GPU
 def analyze_video(video):
     cap = cv2.VideoCapture(video)
@@ -23,7 +31,6 @@ def analyze_video(video):
     if total_frames == 0:
         return "Could not read video frames."
 
-    # Evenly sample 10 frames
     num_frames_to_process = min(200, total_frames)
     frame_indices = [int(i * total_frames / num_frames_to_process) for i in range(num_frames_to_process)]
 
@@ -59,24 +66,13 @@ def analyze_video(video):
         probs = torch.nn.functional.softmax(outputs.logits, dim=1)
         predictions = torch.argmax(probs, dim=1).tolist()
 
-    # Identify index for AI/fake class
-    ai_class_index = None
-    for k, v in model.config.id2label.items():
-        if v.strip().lower() in ["ai", "fake", "ai_generated"]:
-            ai_class_index = k
-            break
-
-    if ai_class_index is None:
-        return "Error: Could not identify the AI-generated class label from model config."
-
     ai_scores = probs[:, ai_class_index].tolist()
 
     avg_score = sum(ai_scores) / len(ai_scores)
     total_frames = len(ai_scores)
-    ai_frame_count = sum(1 for score in ai_scores if score >= 0.5)  # Soft threshold for "suspicious" frame
+    ai_frame_count = sum(1 for score in ai_scores if score >= 0.5)
     ai_frame_ratio = ai_frame_count / total_frames
-    
-    # Refined verdict based on AI frame ratio and confidence
+
     if ai_frame_ratio < 0.25:
         verdict = "✅ Very likely real — no strong signs of AI generation."
     elif 0.25 <= ai_frame_ratio < 0.5:
@@ -88,10 +84,10 @@ def analyze_video(video):
             verdict = "⚠️ Likely AI-generated — multiple signs suggest manipulation."
         else:
             verdict = "🤔 Unclear — the video shows mixed signals."
-    # Output results
+
     results = []
     for frame_idx, pred, ai_score in zip(frame_numbers, predictions, ai_scores):
-        label = model.config.id2label[pred]
+        label = id2label[pred]
         confidence = ai_score * 100
         results.append(f"Frame {frame_idx}: {label} (AI Confidence: {confidence:.1f}%)")