Update app.py

app.py

@@ -64,19 +64,23 @@ class ViolenceDetector3DCNN(nn.Module):
 # ============================================
 @st.cache_resource
 def load_model():
-    # Download model from Hugging Face
-    model_path = hf_hub_download(
-        repo_id="santa47/violence-detection-3dcnn",
-        filename="violence_detector.pth"
-    )
-    
-    # Load model
-    model = ViolenceDetector3DCNN(num_classes=2, dropout=0.5)
-    checkpoint = torch.load(model_path, map_location=torch.device('cpu'))
-    model.load_state_dict(checkpoint['model_state_dict'])
-    model.eval()
-    
-    return model
+    try:
+        # Download model from Hugging Face
+        model_path = hf_hub_download(
+            repo_id="santa47/violence-detection-3dcnn",
+            filename="violence_detector.pth"
+        )
+        
+        # Load model
+        model = ViolenceDetector3DCNN(num_classes=2, dropout=0.5)
+        checkpoint = torch.load(model_path, map_location=torch.device('cpu'), weights_only=False)
+        model.load_state_dict(checkpoint['model_state_dict'])
+        model.eval()
+        
+        return model
+    except Exception as e:
+        st.error(f"Failed to load model: {e}")
+        return None
 
 
 # ============================================
@@ -84,45 +88,56 @@ def load_model():
 # ============================================
 def process_video(video_path, num_frames=16, frame_size=(112, 112)):
     """Extract and preprocess frames from video"""
-    cap = cv2.VideoCapture(video_path)
-    frames = []
-    
-    while True:
-        ret, frame = cap.read()
-        if not ret:
-            break
-        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
-        frame = cv2.resize(frame, frame_size)
-        frames.append(frame)
-    
-    cap.release()
+    try:
+        cap = cv2.VideoCapture(video_path)
+        
+        if not cap.isOpened():
+            st.error(f"❌ Cannot open video file: {video_path}")
+            return None
+        
+        frames = []
+        
+        while True:
+            ret, frame = cap.read()
+            if not ret:
+                break
+            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+            frame = cv2.resize(frame, frame_size)
+            frames.append(frame)
+        
+        cap.release()
+        
+        if len(frames) == 0:
+            st.error("❌ No frames extracted from video")
+            return None
+        
+        # Sample frames uniformly
+        total_frames = len(frames)
+        if total_frames >= num_frames:
+            indices = np.linspace(0, total_frames - 1, num_frames, dtype=int)
+        else:
+            indices = list(range(total_frames)) + [total_frames - 1] * (num_frames - total_frames)
+        
+        sampled_frames = [frames[i] for i in indices]
+        
+        # Convert to tensor: (T, H, W, C) -> (C, T, H, W)
+        video_tensor = np.stack(sampled_frames, axis=0)
+        video_tensor = video_tensor.transpose(3, 0, 1, 2)
+        video_tensor = video_tensor.astype(np.float32) / 255.0
+        
+        # Normalize
+        mean = np.array([0.485, 0.456, 0.406]).reshape(3, 1, 1, 1)
+        std = np.array([0.229, 0.224, 0.225]).reshape(3, 1, 1, 1)
+        video_tensor = (video_tensor - mean) / std
+        
+        # Add batch dimension
+        video_tensor = torch.from_numpy(video_tensor).unsqueeze(0).float()
+        
+        return video_tensor
     
-    if len(frames) == 0:
+    except Exception as e:
+        st.error(f"❌ Error processing video: {e}")
         return None
-    
-    # Sample frames uniformly
-    total_frames = len(frames)
-    if total_frames >= num_frames:
-        indices = np.linspace(0, total_frames - 1, num_frames, dtype=int)
-    else:
-        indices = list(range(total_frames)) + [total_frames - 1] * (num_frames - total_frames)
-    
-    sampled_frames = [frames[i] for i in indices]
-    
-    # Convert to tensor: (T, H, W, C) -> (C, T, H, W)
-    video_tensor = np.stack(sampled_frames, axis=0)
-    video_tensor = video_tensor.transpose(3, 0, 1, 2)
-    video_tensor = video_tensor.astype(np.float32) / 255.0
-    
-    # Normalize
-    mean = np.array([0.485, 0.456, 0.406]).reshape(3, 1, 1, 1)
-    std = np.array([0.229, 0.224, 0.225]).reshape(3, 1, 1, 1)
-    video_tensor = (video_tensor - mean) / std
-    
-    # Add batch dimension
-    video_tensor = torch.from_numpy(video_tensor).unsqueeze(0).float()
-    
-    return video_tensor
 
 
 # ============================================
@@ -156,67 +171,90 @@ def main():
     # Load model
     with st.spinner("Loading model..."):
         model = load_model()
+    
+    if model is None:
+        st.error("❌ Failed to load model. Please refresh the page.")
+        return
+    
     st.success("✅ Model loaded!")
     
     # File uploader
     st.markdown("### Upload a Video")
     uploaded_file = st.file_uploader(
-        "Choose a video file (AVI, MP4, MKV)",
-        type=['avi', 'mp4', 'mkv']
+        "Choose a video file (AVI, MP4, MKV, MOV, WEBM)",
+        type=['avi', 'mp4', 'mkv', 'mov', 'webm']
     )
     
     if uploaded_file is not None:
-        # Save uploaded file temporarily
-        with tempfile.NamedTemporaryFile(delete=False, suffix='.avi') as tmp_file:
-            tmp_file.write(uploaded_file.read())
-            tmp_path = tmp_file.name
-        
-        # Display video
-        st.video(uploaded_file)
-        
-        # Process and predict
-        if st.button("🔍 Analyze Video", type="primary"):
-            with st.spinner("Processing video..."):
-                # Process video
-                video_tensor = process_video(tmp_path)
-                
-                if video_tensor is None:
-                    st.error("❌ Could not process video. Please try another file.")
-                else:
-                    # Predict
-                    pred_class, confidence, probs = predict(model, video_tensor)
-                    
-                    # Display results
-                    st.markdown("---")
-                    st.markdown("### 📊 Results")
-                    
-                    col1, col2 = st.columns(2)
-                    
-                    with col1:
-                        if pred_class == 1:
-                            st.error("⚠️ **VIOLENCE DETECTED**")
-                        else:
-                            st.success("✅ **NO VIOLENCE**")
-                    
-                    with col2:
-                        st.metric("Confidence", f"{confidence * 100:.1f}%")
-                    
-                    # Probability bars
-                    st.markdown("### Probability Distribution")
-                    
-                    col1, col2 = st.columns(2)
-                    with col1:
-                        st.markdown("**Non-Violence**")
-                        st.progress(float(probs[0]))
-                        st.write(f"{probs[0] * 100:.1f}%")
-                    
-                    with col2:
-                        st.markdown("**Violence**")
-                        st.progress(float(probs[1]))
-                        st.write(f"{probs[1] * 100:.1f}%")
+        # FIX 1: Get proper file extension from uploaded file
+        file_extension = os.path.splitext(uploaded_file.name)[1].lower()
+        if not file_extension:
+            file_extension = '.mp4'
+        
+        # FIX 2: Read file bytes ONCE and store
+        video_bytes = uploaded_file.read()
+        
+        # FIX 3: Save with correct extension
+        tmp_path = None
+        try:
+            with tempfile.NamedTemporaryFile(delete=False, suffix=file_extension) as tmp_file:
+                tmp_file.write(video_bytes)
+                tmp_path = tmp_file.name
+            
+            # FIX 4: Display video using bytes (not the file object after read)
+            st.video(video_bytes)
             
-            # Cleanup
-            os.unlink(tmp_path)
+            # Process and predict
+            if st.button("🔍 Analyze Video", type="primary"):
+                with st.spinner("Processing video..."):
+                    # Process video
+                    video_tensor = process_video(tmp_path)
+                    
+                    if video_tensor is None:
+                        st.error("❌ Could not process video. Please try another file.")
+                    else:
+                        # Predict
+                        pred_class, confidence, probs = predict(model, video_tensor)
+                        
+                        # Display results
+                        st.markdown("---")
+                        st.markdown("### 📊 Results")
+                        
+                        col1, col2 = st.columns(2)
+                        
+                        with col1:
+                            if pred_class == 1:
+                                st.error("⚠️ **VIOLENCE DETECTED**")
+                            else:
+                                st.success("✅ **NO VIOLENCE**")
+                        
+                        with col2:
+                            st.metric("Confidence", f"{confidence * 100:.1f}%")
+                        
+                        # Probability bars
+                        st.markdown("### Probability Distribution")
+                        
+                        col1, col2 = st.columns(2)
+                        with col1:
+                            st.markdown("**Non-Violence**")
+                            st.progress(float(probs[0]))
+                            st.write(f"{probs[0] * 100:.1f}%")
+                        
+                        with col2:
+                            st.markdown("**Violence**")
+                            st.progress(float(probs[1]))
+                            st.write(f"{probs[1] * 100:.1f}%")
+        
+        except Exception as e:
+            st.error(f"❌ Error: {e}")
+        
+        finally:
+            # FIX 5: Cleanup in finally block (always runs)
+            if tmp_path and os.path.exists(tmp_path):
+                try:
+                    os.unlink(tmp_path)
+                except:
+                    pass
     
     # Footer
     st.markdown("---")
@@ -232,4 +270,4 @@ def main():
 
 
 if __name__ == "__main__":
-    main()
+    main()