EVobj-Detection / app.py
asrcoddeploy's picture
Update app.py
a42d3cf verified
import gradio as gr
import cv2
import numpy as np
from ultralytics import YOLO
# Load the model
model = YOLO('best.pt')
def get_intelligence(frame, box, label):
x1, y1, x2, y2 = map(int, box.xyxy[0])
w_px = x2 - x1
h_px = y2 - y1
# Distance estimation logic
dist_m = (2.5 * 700) / w_px
dist_cat = "Near" if dist_m < 15 else "Medium" if dist_m < 30 else "Far"
# Aspect ratio logic for direction
aspect_ratio = w_px / h_px
direction = "Front/Rear" if aspect_ratio < 1.2 else "Side View"
return dist_cat, direction
def predict(img):
results = model(img, conf=0.75)
annotated_img = img.copy()
status_reports = []
for r in results:
for box in r.boxes:
cls_id = int(box.cls[0])
label = r.names[cls_id]
dist, direction = get_intelligence(img, box, label)
# Create a label for the report
report = f"Detected {label} ({direction}) - Distance: {dist}"
status_reports.append(report)
# Draw on image
x1, y1, x2, y2 = map(int, box.xyxy[0])
cv2.rectangle(annotated_img, (x1, y1), (x2, y2), (0, 255, 0), 3)
cv2.putText(annotated_img, f"{label} {dist}", (x1, y1 - 10),
cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)
return annotated_img, "\n".join(status_reports) if status_reports else "No emergency vehicles detected."
# Build the Gradio UI
demo = gr.Interface(
fn=predict,
inputs=gr.Image(type="numpy"),
outputs=[gr.Image(type="numpy"), gr.Textbox(label="Intelligence Report")],
title="EVobj: Emergency Vehicle Detection System",
description="Upload an image of traffic to detect emergency vehicles and estimate proximity."
)
if __name__ == "__main__":
demo.launch()