re-added UI, test

app.py

@@ -1,152 +1,152 @@
-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 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
-# import cv2
-# import numpy as np
 
-# # Load model and processor
+# # 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")
 
-# # Function to detect accidents in an image
 # def detect_accident(image):
+#     """Runs accident detection on the input image."""
 #     inputs = processor(images=image, return_tensors="pt")
 #     outputs = model(**inputs)
 
-#     # Post-process the results
+#     # 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 boxes and labels on the image
+#     # 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
 
-# # Function to detect accidents frame-by-frame in a video
-# def detect_accident_in_video(video_path):
-#     cap = cv2.VideoCapture(video_path)
-#     frames = []
-#     while True:
-#         ret, frame = cap.read()
-#         if not ret:
-#             break
+# @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
+from PIL import Image, ImageDraw
+import torch
+import cv2
+import numpy as np
+
+# Load model and processor
+model = DetrForObjectDetection.from_pretrained("hilmantm/detr-traffic-accident-detection")
+processor = DetrImageProcessor.from_pretrained("hilmantm/detr-traffic-accident-detection")
+
+# Function to detect accidents in an image
+def detect_accident(image):
+    inputs = processor(images=image, return_tensors="pt")
+    outputs = model(**inputs)
+
+    # Post-process the results
+    target_sizes = torch.tensor([image.size[::-1]])
+    results = processor.post_process_object_detection(outputs, target_sizes=target_sizes, threshold=0.9)[0]
+
+    # Draw 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")
+
+    return image
+
+# Function to detect accidents frame-by-frame in a video
+def detect_accident_in_video(video_path):
+    cap = cv2.VideoCapture(video_path)
+    frames = []
+    while True:
+        ret, frame = cap.read()
+        if not ret:
+            break
         
-#         # Convert frame to PIL Image
-#         frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
-#         pil_frame = Image.fromarray(frame_rgb)
-
-#         # Run accident detection on the frame
-#         processed_frame = detect_accident(pil_frame)
-
-#         # Convert PIL image back to numpy array for video
-#         frames.append(np.array(processed_frame))
-
-#     cap.release()
-
-#     # Save processed frames as output video
-#     height, width, _ = frames[0].shape
-#     out = cv2.VideoWriter("output.mp4", cv2.VideoWriter_fourcc(*"mp4v"), 10, (width, height))
-#     for frame in frames:
-#         out.write(cv2.cvtColor(frame, cv2.COLOR_RGB2BGR))
-#     out.release()
-
-#     return "output.mp4"
-
-# # Gradio app interface
-# with gr.Blocks() as interface:
-#     gr.Markdown("# Traffic Accident Detection")
-#     gr.Markdown(
-#         "Upload an image or video to detect traffic accidents using the DETR model. "
-#         "For videos, the system processes frame by frame and outputs a new video with accident detection."
-#     )
-
-#     # Input components
-#     with gr.Tab("Image Input"):
-#         image_input = gr.Image(type="pil", label="Upload Image")
-#         image_output = gr.Image(type="pil", label="Detection Output")
-#         image_button = gr.Button("Detect Accidents in Image")
+        # Convert frame to PIL Image
+        frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+        pil_frame = Image.fromarray(frame_rgb)
+
+        # Run accident detection on the frame
+        processed_frame = detect_accident(pil_frame)
+
+        # Convert PIL image back to numpy array for video
+        frames.append(np.array(processed_frame))
+
+    cap.release()
+
+    # Save processed frames as output video
+    height, width, _ = frames[0].shape
+    out = cv2.VideoWriter("output.mp4", cv2.VideoWriter_fourcc(*"mp4v"), 10, (width, height))
+    for frame in frames:
+        out.write(cv2.cvtColor(frame, cv2.COLOR_RGB2BGR))
+    out.release()
+
+    return "output.mp4"
+
+# Gradio app interface
+with gr.Blocks() as interface:
+    gr.Markdown("# Traffic Accident Detection")
+    gr.Markdown(
+        "Upload an image or video to detect traffic accidents using the DETR model. "
+        "For videos, the system processes frame by frame and outputs a new video with accident detection."
+    )
+
+    # Input components
+    with gr.Tab("Image Input"):
+        image_input = gr.Image(type="pil", label="Upload Image")
+        image_output = gr.Image(type="pil", label="Detection Output")
+        image_button = gr.Button("Detect Accidents in Image")
     
-#     with gr.Tab("Video Input"):
-#         video_input = gr.Video(label="Upload Video")
-#         video_output = gr.Video(label="Processed Video")
-#         video_button = gr.Button("Detect Accidents in Video")
+    with gr.Tab("Video Input"):
+        video_input = gr.Video(label="Upload Video")
+        video_output = gr.Video(label="Processed Video")
+        video_button = gr.Button("Detect Accidents in Video")
 
-#     # Define behaviors
-#     image_button.click(fn=detect_accident, inputs=image_input, outputs=image_output)
-#     video_button.click(fn=detect_accident_in_video, inputs=video_input, outputs=video_output)
+    # Define behaviors
+    image_button.click(fn=detect_accident, inputs=image_input, outputs=image_output)
+    video_button.click(fn=detect_accident_in_video, inputs=video_input, outputs=video_output)
 
-# interface.launch()
+interface.launch()