Update app.py

app.py

@@ -1,4 +1,5 @@
-from fastapi import FastAPI, File, UploadFile
+from fastapi import FastAPI, File, UploadFile, HTTPException
+from fastapi.responses import FileResponse
 import cv2
 import numpy as np
 from ultralytics import YOLO
@@ -11,12 +12,12 @@ import tempfile
 app = FastAPI()
 
 # Load YOLO model safely
-model_path = "12x.pt"
-if not os.path.exists(model_path):
-    print(f"Warning: Model file '{model_path}' not found. API will not work properly.")
-    model = None  # Handle model loading failure
+MODEL_PATH = "12x.pt"
+if not os.path.exists(MODEL_PATH):
+    print(f"⚠ Warning: Model file '{MODEL_PATH}' not found. API will not work properly.")
+    model = None
 else:
-    model = YOLO(model_path)
+    model = YOLO(MODEL_PATH)
 
 
 def process_frame(frame):
@@ -49,82 +50,91 @@ def process_frame(frame):
 async def upload_image(file: UploadFile = File(...)):
     """Upload an image and get object detection results."""
     if model is None:
-        return {"error": "Model not loaded. Please upload '12x.pt' to run detection."}
+        raise HTTPException(status_code=500, detail="Model not loaded. Please upload '12x.pt' to run detection.")
 
-    contents = await file.read()
-    nparr = np.frombuffer(contents, np.uint8)
-    img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
+    try:
+        contents = await file.read()
+        nparr = np.frombuffer(contents, np.uint8)
+        img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
 
-    predictions, object_count = process_frame(img)
+        predictions, object_count = process_frame(img)
 
-    # Draw bounding boxes on the image
-    for pred in predictions:
-        x1, y1, x2, y2 = map(int, pred["bbox"])
-        label = f"{pred['class']} ({pred['confidence']:.2f})"
-        cv2.rectangle(img, (x1, y1), (x2, y2), (0, 255, 0), 2)
-        cv2.putText(img, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
+        # Draw bounding boxes on the image
+        for pred in predictions:
+            x1, y1, x2, y2 = map(int, pred["bbox"])
+            label = f"{pred['class']} ({pred['confidence']:.2f})"
+            cv2.rectangle(img, (x1, y1), (x2, y2), (0, 255, 0), 2)
+            cv2.putText(img, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
 
-    _, buffer = cv2.imencode('.jpg', img)
-    img_base64 = base64.b64encode(buffer).decode('utf-8')
+        _, buffer = cv2.imencode('.jpg', img)
+        img_base64 = base64.b64encode(buffer).decode('utf-8')
 
-    return {
-        "image": f"data:image/jpeg;base64,{img_base64}",
-        "object_count": object_count
-    }
+        return {
+            "image": f"data:image/jpeg;base64,{img_base64}",
+            "object_count": object_count
+        }
+
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Error processing image: {str(e)}")
 
 
 @app.post("/upload-video/")
 async def upload_video(file: UploadFile = File(...)):
     """Upload a video, process it frame by frame, and return the processed video."""
     if model is None:
-        return {"error": "Model not loaded. Please upload '12x.pt' to run detection."}
+        raise HTTPException(status_code=500, detail="Model not loaded. Please upload '12x.pt' to run detection.")
+
+    try:
+        with tempfile.NamedTemporaryFile(delete=False, suffix=".mp4") as temp_video:
+            shutil.copyfileobj(file.file, temp_video)
+            temp_video_path = temp_video.name
+
+        cap = cv2.VideoCapture(temp_video_path)
+        if not cap.isOpened():
+            os.remove(temp_video_path)
+            raise HTTPException(status_code=400, detail="Could not open video file.")
 
-    with tempfile.NamedTemporaryFile(delete=False, suffix=".mp4") as temp_video:
-        shutil.copyfileobj(file.file, temp_video)
-        temp_video_path = temp_video.name
+        # Get video properties
+        fps = int(cap.get(cv2.CAP_PROP_FPS))
+        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
 
-    cap = cv2.VideoCapture(temp_video_path)
-    if not cap.isOpened():
-        os.remove(temp_video_path)
-        return {"error": "Could not open video file"}
+        # Output video
+        output_video_path = temp_video_path.replace(".mp4", "_processed.mp4")
+        out = cv2.VideoWriter(output_video_path, fourcc, fps, (width, height))
 
-    # Get video properties
-    fps = int(cap.get(cv2.CAP_PROP_FPS))
-    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
-    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
-    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+        frame_interval = 5  # Process every 5th frame for efficiency
+        frame_index = 0
 
-    # Output video
-    output_video_path = temp_video_path.replace(".mp4", "_processed.mp4")
-    out = cv2.VideoWriter(output_video_path, fourcc, fps, (width, height))
+        while True:
+            ret, frame = cap.read()
+            if not ret:
+                break
 
-    frame_interval = 5  # Process every 5th frame for efficiency
-    frame_index = 0
+            if frame_index % frame_interval == 0:
+                predictions, _ = process_frame(frame)
 
-    while True:
-        ret, frame = cap.read()
-        if not ret:
-            break
+                # Draw bounding boxes on the frame
+                for pred in predictions:
+                    x1, y1, x2, y2 = map(int, pred["bbox"])
+                    label = f"{pred['class']} ({pred['confidence']:.2f})"
+                    cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
+                    cv2.putText(frame, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
 
-        if frame_index % frame_interval == 0:
-            predictions, _ = process_frame(frame)
+            out.write(frame)
+            frame_index += 1
 
-            # Draw bounding boxes on the frame
-            for pred in predictions:
-                x1, y1, x2, y2 = map(int, pred["bbox"])
-                label = f"{pred['class']} ({pred['confidence']:.2f})"
-                cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
-                cv2.putText(frame, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
+        cap.release()
+        out.release()
+        os.remove(temp_video_path)  # Clean up temp file
 
-        out.write(frame)
-        frame_index += 1
+        return FileResponse(output_video_path, media_type="video/mp4", filename="processed_video.mp4")
 
-    cap.release()
-    out.release()
-    os.remove(temp_video_path)  # Clean up temp file
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Error processing video: {str(e)}")
 
-    return FileResponse(output_video_path, media_type="video/mp4", filename="processed_video.mp4")
 
 @app.get("/")
 def home():
-    return {"message": "Object Detection API for Images and Videos using 12x.pt"}
+    return {"message": "🎯 Object Detection API for Images and Videos using 12x.pt"}