Update app.py

app.py

@@ -23,10 +23,25 @@ def extract_keypoints(frame):
             return flattened_keypoints
     return None  # Return None if no keypoints are detected
 
-def process_frame(frame):
+def process_input(input_media):
     """
-    Process each frame for suspicious activity detection
+    Process either a video or an image for suspicious activity detection
     """
+    # Determine if input is a video or image path
+    is_video = input_media.lower().endswith(('.mp4', '.avi', '.mov'))
+    
+    if is_video:
+        return process_video(input_media)
+    else:
+        return process_image(input_media)
+
+def process_image(image_path):
+    """
+    Process a single image for suspicious activity detection
+    """
+    # Read the image
+    frame = cv2.imread(image_path)
+    
     # Perform YOLO detection
     results = yolo_model(frame, verbose=False)
     for box in results[0].boxes:
@@ -57,15 +72,18 @@ def process_frame(frame):
                     cv2.rectangle(frame, (x1, y1), (x2, y2), color, 2)
                     cv2.putText(frame, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
                 else:
-                    print("No valid keypoints detected for ROI. Skipping frame.")
+                    print("No valid keypoints detected for ROI. Skipping.")
             else:
-                print("ROI size is zero. Skipping frame.")
+                print("ROI size is zero. Skipping.")
     
-    return frame
+    # Save the processed image
+    output_path = 'output_image.jpg'
+    cv2.imwrite(output_path, frame)
+    return output_path
 
-def detect_suspicious_activity(input_video):
+def process_video(input_video):
     """
-    Main function to process video for suspicious activity detection
+    Process video for suspicious activity detection
     """
     # Open video capture
     cap = cv2.VideoCapture(input_video)
@@ -85,11 +103,42 @@ def detect_suspicious_activity(input_video):
         if not ret:
             break
         
-        # Process and annotate frame
-        processed_frame = process_frame(frame)
+        # Perform YOLO detection
+        results = yolo_model(frame, verbose=False)
+        for box in results[0].boxes:
+            cls = int(box.cls[0])  # Class ID
+            confidence = float(box.conf[0])
+            
+            # Detect persons only (class_id 0 for 'person')
+            if cls == 0 and confidence > 0.5:
+                x1, y1, x2, y2 = map(int, box.xyxy[0])  # Bounding box coordinates
+                
+                # Extract ROI for classification
+                roi = frame[y1:y2, x1:x2]
+                if roi.size > 0:
+                    # Preprocess ROI to extract keypoints
+                    keypoints = extract_keypoints(roi)
+                    
+                    if keypoints is not None and len(keypoints) > 0:
+                        # Standardize and reshape keypoints for LSTM input
+                        keypoints_scaled = scaler.fit_transform([keypoints])  # Standardize features
+                        keypoints_reshaped = keypoints_scaled.reshape((1, 1, len(keypoints)))  # Reshape for LSTM
+                        
+                        # Predict with LSTM model
+                        prediction = (lstm_model.predict(keypoints_reshaped) > 0.5).astype(int)[0][0]
+                        
+                        # Draw bounding box and label
+                        color = (0, 0, 255) if prediction == 1 else (0, 255, 0)
+                        label = 'Suspicious' if prediction == 1 else 'Normal'
+                        cv2.rectangle(frame, (x1, y1), (x2, y2), color, 2)
+                        cv2.putText(frame, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
+                    else:
+                        print("No valid keypoints detected for ROI. Skipping frame.")
+                else:
+                    print("ROI size is zero. Skipping frame.")
         
         # Write processed frame to output video
-        out.write(processed_frame)
+        out.write(frame)
     
     # Release resources
     cap.release()
@@ -99,11 +148,17 @@ def detect_suspicious_activity(input_video):
 
 # Create Gradio interface
 iface = gr.Interface(
-    fn=detect_suspicious_activity,
-    inputs=gr.Video(label="Upload Video"),
-    outputs=gr.Video(label="Processed Video"),
+    fn=process_input,
+    inputs=[
+        gr.File(label="Upload Image or Video", 
+                file_types=['image', 'video'], 
+                type="filepath")
+    ],
+    outputs=[
+        gr.File(label="Processed Media")
+    ],
     title="Suspicious Activity Detection",
-    description="Upload a video to detect suspicious activities using YOLO and LSTM models"
+    description="Upload an image or video to detect suspicious activities using YOLO and LSTM models. Suspicious activities will be marked with red bounding boxes, normal activities with green."
 )
 
 # Launch the interface