Update app.py

app.py

@@ -1,5 +1,4 @@
 from fastapi import FastAPI, File, UploadFile
-from fastapi.responses import StreamingResponse
 import cv2
 import numpy as np
 from ultralytics import YOLO
@@ -7,27 +6,33 @@ import base64
 import os
 import shutil
 import tempfile
-import asyncio
 
 # Initialize FastAPI app
 app = FastAPI()
 
-# Load YOLO model (Ensure 12x.pt exists)
+# Load YOLO model safely
 model_path = "12x.pt"
 if not os.path.exists(model_path):
-    raise FileNotFoundError(f"Model file '{model_path}' not found. Please place it in the project directory.")
+    print(f"Warning: Model file '{model_path}' not found. API will not work properly.")
+    model = None  # Handle model loading failure
+else:
+    model = YOLO(model_path)
 
-model = YOLO(model_path)
 
 def process_frame(frame):
     """Process a single frame with YOLO and return predictions."""
+    if model is None:
+        return [], {}
+
     results = model(frame)
     predictions = []
     object_count = {}
 
     for result in results:
         for box in result.boxes:
-            class_name = result.names[int(box.cls)]
+            class_id = int(box.cls)
+            class_name = model.names.get(class_id, f"Unknown_{class_id}")  # Handle missing class names
+
             predictions.append({
                 "class": class_name,
                 "confidence": float(box.conf),
@@ -39,48 +44,20 @@ def process_frame(frame):
 
     return predictions, object_count
 
-def encode_frame(frame):
-    """Encode a frame as JPEG and return base64-encoded string."""
-    _, buffer = cv2.imencode('.jpg', frame)
-    return base64.b64encode(buffer).decode('utf-8')
-
-@app.get("/video-stream/")
-async def video_stream():
-    """Endpoint to stream video frames with real-time object detection."""
-    cap = cv2.VideoCapture(0)
-    if not cap.isOpened():
-        return {"error": "Could not open webcam"}
-
-    async def generate():
-        while True:
-            ret, frame = cap.read()
-            if not ret:
-                break
-
-            predictions, _ = process_frame(frame)
-            
-            # Draw bounding boxes
-            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)
-            
-            _, buffer = cv2.imencode('.jpg', frame)
-            yield (b'--frame\r\n' b'Content-Type: image/jpeg\r\n\r\n' + buffer.tobytes() + b'\r\n')
-            await asyncio.sleep(0.1)  # Adjust frame rate
-
-    return StreamingResponse(generate(), media_type="multipart/x-mixed-replace; boundary=frame")
 
 @app.post("/upload-image/")
 async def upload_image(file: UploadFile = File(...)):
-    """Endpoint to upload an image and get object detection results."""
+    """Upload an image and get object detection results."""
+    if model is None:
+        return {"error": "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)
 
     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})"
@@ -90,17 +67,25 @@ async def upload_image(file: UploadFile = File(...)):
     _, buffer = cv2.imencode('.jpg', img)
     img_base64 = base64.b64encode(buffer).decode('utf-8')
 
-    return {"image": img_base64, "object_count": object_count}
+    return {
+        "image": f"data:image/jpeg;base64,{img_base64}",
+        "object_count": object_count
+    }
+
 
 @app.post("/upload-video/")
 async def upload_video(file: UploadFile = File(...)):
-    """Endpoint to upload a video, process it frame by frame, and return detection results."""
+    """Upload a video, process it frame by frame, and return detection results."""
+    if model is None:
+        return {"error": "Model not loaded. Please upload '12x.pt' to run detection."}
+
     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)
         return {"error": "Could not open video file"}
 
     frame_results = []
@@ -116,16 +101,18 @@ async def upload_video(file: UploadFile = File(...)):
             predictions, object_count = process_frame(frame)
             frame_results.append({
                 "frame_index": frame_index,
-                "object_count": object_count
+                "object_count": object_count,
+                "detections": predictions
             })
 
         frame_index += 1
 
     cap.release()
-    os.remove(temp_video_path)
+    os.remove(temp_video_path)  # Clean up temporary file
 
     return {"video_results": frame_results}
 
+
 @app.get("/")
 def home():
-    return {"message": "Real-Time Object Detection API with Image, Video, and Streaming Support"}
+    return {"message": "Object Detection API for Images and Videos using 12x.pt"}