v0

app.py

@@ -0,0 +1,68 @@
+from fastapi import FastAPI, File, UploadFile
+from fastapi.responses import JSONResponse
+from fastapi.middleware.cors import CORSMiddleware
+from transformers import DetrImageProcessor, DetrForObjectDetection
+from PIL import Image, ImageDraw
+import io
+import torch
+
+# Initialize FastAPI app
+app = FastAPI()
+
+# Add CORS middleware to allow communication with external clients
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],  # Change this to the specific domain in production
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+# Load the model and processor
+model = DetrForObjectDetection.from_pretrained("hilmantm/detr-traffic-accident-detection")
+processor = DetrImageProcessor.from_pretrained("hilmantm/detr-traffic-accident-detection")
+
+def detect_accident(image):
+    """Runs accident detection on the input image."""
+    inputs = processor(images=image, return_tensors="pt")
+    outputs = model(**inputs)
+
+    # Post-process results
+    target_sizes = torch.tensor([image.size[::-1]])
+    results = processor.post_process_object_detection(outputs, target_sizes=target_sizes, threshold=0.9)[0]
+
+    # Draw bounding boxes and labels
+    draw = ImageDraw.Draw(image)
+    for box, label, score in zip(results["boxes"], results["labels"], results["scores"]):
+        x_min, y_min, x_max, y_max = box
+        draw.rectangle((x_min, y_min, x_max, y_max), outline="red", width=3)
+        draw.text((x_min, y_min), f"{label}: {score:.2f}", fill="red")
+    
+    return image
+
+@app.post("/detect_accident")
+async def process_frame(file: UploadFile = File(...)):
+    """API endpoint to process an uploaded frame."""
+    try:
+        # Read and preprocess image
+        image = Image.open(io.BytesIO(await file.read()))
+        image = image.resize((256, int(image.height * 256 / image.width)))  # Resize while maintaining aspect ratio
+
+        # Detect accidents
+        processed_image = detect_accident(image)
+
+        # Save the processed image into bytes to send back
+        img_byte_arr = io.BytesIO()
+        processed_image.save(img_byte_arr, format="JPEG")
+        img_byte_arr.seek(0)
+
+        return JSONResponse(
+            content={"status": "success", "message": "Frame processed successfully"},
+            media_type="image/jpeg"
+        )
+    except Exception as e:
+        return JSONResponse(content={"status": "error", "message": str(e)}, status_code=500)
+
+# Run the app
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=8000)