initial commit

app.py

@@ -1,82 +1,275 @@
-import gradio as gr
+from fastapi import FastAPI, File, UploadFile, HTTPException
+from fastapi.responses import FileResponse, JSONResponse
+from fastapi.middleware.cors import CORSMiddleware
+import uvicorn
 import cv2
 import numpy as np
 from ultralytics import YOLO
+import os
+import shutil
+from typing import Optional
+import uuid
+import base64
+from io import BytesIO
+from PIL import Image
+
+# Create FastAPI app
+app = FastAPI(
+    title="YOLO Intrusion Detection API",
+    description="API for detecting intrusions using YOLOv8 model",
+    version="1.0.0"
+)
+
+# Add CORS middleware
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
 
 # Load YOLO model
-model = YOLO("yolov8n.pt")  # Updated to correct model name
+model_name = 'yolov8n.pt'
+model = None  # Will be loaded on startup
 
-# Define trapezoidal restricted area (top-left, top-right, bottom-right, bottom-left)
+# Define trapezoidal restricted area
 trapezoid_pts = np.array([[250, 150], [400, 150], [450, 300], [200, 300]], np.int32)
 
+# Create temp directory for uploads if it doesn't exist
+os.makedirs("temp", exist_ok=True)
+
 def is_inside_trapezoid(box, trapezoid_pts):
     """Check if the center of a detected object is inside the trapezoidal area."""
     x1, y1, x2, y2 = box
-    cx, cy = int((x1 + x2) / 2), int((y1 + y2) / 2)
+    cx, cy = int((x1 + x2) / 2), int((y1 + y2) / 2)  # Calculate center of detected object
+    
+    # Use point-in-polygon check
     return cv2.pointPolygonTest(trapezoid_pts, (cx, cy), False) >= 0
 
-def detect_objects(frame):
-    if frame is None:
-        return np.zeros((480, 640, 3), dtype=np.uint8), "No input frame"
+def process_image(frame):
+    """Process a single image and return the annotated image and intrusion data."""
+    global model
     
+    # Perform object detection
     results = model.predict(frame, conf=0.5)
-    annotated_frame = results[0].plot()
-    
-    # Draw the trapezoid area
+    annotated_frame = results[0].plot()  # Draw bounding boxes
+
+    # Draw trapezoidal restricted area
     cv2.polylines(annotated_frame, [trapezoid_pts.reshape((-1, 1, 2))], isClosed=True, color=(0, 0, 255), thickness=2)
+
+    isAlert = {'alert': [False, ""], 'personCount': 0}
+    classInIntrusion = ['person', 'bicycle', 'car', 'motorcycle']
     
-    isAlert = {"alert": [False, ""], "personCount": 0}
-    classInIntrusion = ["person", "bicycle", "car", "motorcycle"]
-    
+    detections = []
+
+    # Loop through detected objects
     for r in results:
-        for box, cls in zip(r.boxes.xyxy, r.boxes.cls):
-            class_id = int(cls.item())
+        for box, cls, conf in zip(r.boxes.xyxy, r.boxes.cls, r.boxes.conf):
+            class_id = int(cls.item())  # Convert to integer
+            confidence = float(conf.item())
+            x1, y1, x2, y2 = map(int, box.tolist())
+            
+            class_name = classInIntrusion[class_id] if class_id < len(classInIntrusion) else f"class_{class_id}"
+            
+            # Add to detections list
+            detections.append({
+                "class": class_name,
+                "confidence": confidence,
+                "bbox": [x1, y1, x2, y2],
+                "in_restricted_area": is_inside_trapezoid(box.tolist(), trapezoid_pts)
+            })
+            
             if class_id == 0:  # Person
-                isAlert["personCount"] += 1
+                isAlert['personCount'] += 1
+            
             if class_id in [0, 1, 2, 3]:  # Person, bicycle, car, motorcycle
                 if is_inside_trapezoid(box.tolist(), trapezoid_pts):
-                    isAlert["alert"] = [True, classInIntrusion[class_id]]
+                    isAlert['alert'] = [True, classInIntrusion[class_id]]
                     # Mark the intrusion with a red box
-                    x1, y1, x2, y2 = map(int, box.tolist())
                     cv2.rectangle(annotated_frame, (x1, y1), (x2, y2), (0, 0, 255), 3)
     
     # Add alert text on the frame
     alert_text = f"Intrusion Alert: {isAlert['alert'][0]}, Object: {isAlert['alert'][1]}, Persons: {isAlert['personCount']}"
     cv2.putText(annotated_frame, alert_text, (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2)
     
-    return annotated_frame, alert_text
+    # Convert the response
+    response = {
+        "intrusion_detected": isAlert['alert'][0],
+        "intruding_object": isAlert['alert'][1],
+        "person_count": isAlert['personCount'],
+        "detections": detections
+    }
+    
+    return annotated_frame, response
+
+def encode_image_to_base64(image):
+    """Convert an OpenCV image to base64 encoded string."""
+    _, buffer = cv2.imencode('.jpg', image)
+    return base64.b64encode(buffer).decode('utf-8')
+
+@app.on_event("startup")
+async def startup_event():
+    """Load the YOLO model when the app starts."""
+    global model
+    model = YOLO(model_name)
+    print(f"Model {model_name} loaded successfully")
 
-def webcam_feed():
-    cap = cv2.VideoCapture(0)
+@app.get("/")
+async def root():
+    """Root endpoint."""
+    return {
+        "message": "YOLO Intrusion Detection API is running",
+        "documentation": "/docs",
+        "endpoints": {
+            "process_image": "/process_image/",
+            "process_video": "/process_video/",
+            "health": "/health/"
+        }
+    }
+
+@app.get("/health/")
+async def health_check():
+    """Health check endpoint."""
+    return {"status": "healthy", "model": model_name}
+
+@app.post("/process_image/")
+async def api_process_image(file: UploadFile = File(...), return_image: bool = True):
+    """
+    Process an image file and detect intrusions.
+    
+    Args:
+        file: The image file to process
+        return_image: If True, returns the annotated image as base64
+    
+    Returns:
+        JSON with detection results and optionally the annotated image
+    """
+    # Check file extension
+    if not file.filename.lower().endswith(('.png', '.jpg', '.jpeg')):
+        raise HTTPException(status_code=400, detail="Only PNG and JPG images are supported")
+    
+    # Read and process image
+    contents = await file.read()
+    nparr = np.frombuffer(contents, np.uint8)
+    img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
+    
+    if img is None:
+        raise HTTPException(status_code=400, detail="Could not decode image")
+    
+    # Process the image
+    annotated_img, results = process_image(img)
+    
+    # Optionally include the annotated image
+    if return_image:
+        results["image"] = encode_image_to_base64(annotated_img)
     
-    # Check if the webcam opened successfully
+    return results
+
+@app.post("/process_video/")
+async def api_process_video(file: UploadFile = File(...)):
+    """
+    Process a video file and detect intrusions.
+    
+    Args:
+        file: The video file to process
+    
+    Returns:
+        JSON with detection results and path to processed video
+    """
+    # Check file extension
+    if not file.filename.lower().endswith(('.mp4', '.avi', '.mov')):
+        raise HTTPException(status_code=400, detail="Only MP4, AVI, and MOV videos are supported")
+    
+    # Create a unique temporary file name
+    temp_input = f"temp/input_{uuid.uuid4()}.mp4"
+    temp_output = f"temp/output_{uuid.uuid4()}.mp4"
+    
+    # Save uploaded file
+    with open(temp_input, "wb") as buffer:
+        shutil.copyfileobj(file.file, buffer)
+    
+    # Process the video
+    cap = cv2.VideoCapture(temp_input)
     if not cap.isOpened():
-        return np.zeros((480, 640, 3), dtype=np.uint8), "Failed to open webcam"
+        os.remove(temp_input)
+        raise HTTPException(status_code=400, detail="Could not open video file")
+    
+    # Get video properties
+    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+    fps = cap.get(cv2.CAP_PROP_FPS)
+    
+    # Create output video file
+    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+    out = cv2.VideoWriter(temp_output, fourcc, fps, (width, height))
+    
+    # Process frames
+    final_results = {
+        "intrusion_detected": False,
+        "intruding_objects": set(),
+        "max_person_count": 0,
+        "frames_processed": 0,
+        "total_detections": 0
+    }
     
     while True:
         ret, frame = cap.read()
         if not ret:
             break
         
-        # Process frame
-        result_frame, alert_message = detect_objects(frame)
+        # Process the frame
+        annotated_frame, frame_results = process_image(frame)
         
-        # Return the processed frame and alert message
-        yield result_frame, alert_message
-
-# Create the Gradio interface with webcam
-demo = gr.Interface(
-    fn=webcam_feed,
-    inputs=[],
-    outputs=[
-        gr.Image(label="Detection Output"),
-        gr.Textbox(label="Alert Status")
-    ],
-    live=True,
-    title="YOLO Intrusion Detection",
-    description="Real-time detection of persons and vehicles inside a restricted trapezoidal area.",
-    allow_flagging="never"
-)
+        # Update final results
+        final_results["frames_processed"] += 1
+        final_results["total_detections"] += len(frame_results["detections"])
+        
+        if frame_results["intrusion_detected"]:
+            final_results["intrusion_detected"] = True
+            if frame_results["intruding_object"]:
+                final_results["intruding_objects"].add(frame_results["intruding_object"])
+        
+        final_results["max_person_count"] = max(
+            final_results["max_person_count"], 
+            frame_results["person_count"]
+        )
+        
+        # Write the frame
+        out.write(annotated_frame)
+    
+    # Release resources
+    cap.release()
+    out.release()
+    
+    # Convert set to list for JSON serialization
+    final_results["intruding_objects"] = list(final_results["intruding_objects"])
+    
+    # Clean up input file
+    os.remove(temp_input)
+    
+    return {
+        "results": final_results,
+        "video_path": f"/download_video/{os.path.basename(temp_output)}"
+    }
+
+@app.get("/download_video/{filename}")
+async def download_video(filename: str):
+    """
+    Download the processed video file.
+    
+    Args:
+        filename: The name of the processed video file
+    
+    Returns:
+        The video file
+    """
+    file_path = f"temp/{filename}"
+    if not os.path.exists(file_path):
+        raise HTTPException(status_code=404, detail="Video not found")
+    
+    return FileResponse(file_path, media_type="video/mp4", filename="processed_video.mp4")
 
-if __name__ == "__main__":
-    demo.queue(max_size=1).launch()
+# # For local development
+# if __name__ == "__main__":
+#     uvicorn.run("app:app", host="0.0.0.0", port=7860, reload=True)

requirements.txt

@@ -1,4 +1,7 @@
 ultralytics
-opencv-python
 numpy
-gradio
+fastapi
+uvicorn
+opencv-python-headless
+pillow
+python-multipart