new version added

app.py

@@ -1,71 +1,138 @@
-from fastapi import FastAPI, File, UploadFile
-from fastapi.responses import JSONResponse
-from fastapi.middleware.cors import CORSMiddleware
+import gradio as gr
 from transformers import DetrImageProcessor, DetrForObjectDetection
-from PIL import Image, ImageDraw
-import io
+from PIL import Image
 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")
+# Load the DETR model and processor
+model_name = "hilmantm/detr-traffic-accident-detection"
+processor = DetrImageProcessor.from_pretrained(model_name)
+model = DetrForObjectDetection.from_pretrained(model_name)
 
 def detect_accident(image):
-    """Runs accident detection on the input image."""
+    """Process an image and detect traffic accidents using the DETR model."""
+    # Preprocess the input image
     inputs = processor(images=image, return_tensors="pt")
+    
+    # Get model predictions
     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
+    
+    # Post-process predictions to extract bounding boxes and labels
+    target_sizes = torch.tensor([image.size[::-1]])  # Image size in (height, width)
+    results = processor.post_process_object_detection(
+        outputs, target_sizes=target_sizes, threshold=0.9
+    )[0]
+    
+    # Draw bounding boxes and labels on the image
     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")
+    for box, score, label in zip(results["boxes"], results["scores"], results["labels"]):
+        box = [int(b) for b in box]
+        label_text = f"{model.config.id2label[label]}: {score:.2f}"
+        draw.rectangle(box, outline="red", width=3)
+        draw.text((box[0], box[1]), label_text, 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)
+# Define the Gradio interface
+iface = gr.Interface(
+    fn=detect_accident,
+    inputs=gr.inputs.Image(type="pil"),
+    outputs=gr.outputs.Image(type="pil"),
+    title="Traffic Accident Detection",
+    description="Upload an image to detect traffic accidents using the DETR model."
+)
+
+# Launch the app
+iface.launch()
+
+
+
+
+
+
+
+
+
+
+
+
+# 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)
+
+
+
+
+
+
+
+
+
+    
 
 # import gradio as gr
 # from transformers import DetrImageProcessor, DetrForObjectDetection