AriasTech/Airplane-Detection

1import cv22import numpy as np3import gradio as gr4from ultralytics import YOLO5from tile_inference import tile_inference6 7# --------------------------8# Thresholds (shared logic)9# --------------------------10CONF_THRES = 0.3511IOU_THRES  = 0.4512 13# --------------------------14# Load model15# --------------------------16model = YOLO("best.pt")17 18CLASS_NAMES = ["civil", "military"]19COLORS = [20    (0, 255, 0),   # civil21    (0, 0, 255)    # military22]23 24# --------------------------25# Normal YOLO Inference26# --------------------------27def normal_inference(model, img):28    result = model(29        img,30        conf=CONF_THRES,31        iou=IOU_THRES,32        agnostic_nms=False33    )[0]34 35    detections = []36    if result.boxes is None:37        return detections38 39    for box in result.boxes:40        detections.append({41            "cls": int(box.cls),42            "conf": float(box.conf),43            "bbox": box.xyxy[0].cpu().numpy().tolist()44        })45 46    return detections47 48# --------------------------49# Unified Prediction50# --------------------------51def predict(image, mode):52    if image is None:53        return None, 0, 054 55    img = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)56 57    if mode == "Tiled Inference":58        detections = tile_inference(59            model,60            img,61            conf_threshold=CONF_THRES,62            iou_threshold=IOU_THRES63        )64    else:65        detections = normal_inference(model, img)66 67    civil_count = 068    military_count = 069 70    for det in detections:71        cls = det["cls"]72        conf = det["conf"]73        x1, y1, x2, y2 = map(int, det["bbox"])74 75        if cls == 0:76            civil_count += 177        else:78            military_count += 179 80        color = COLORS[cls]81        cv2.rectangle(img, (x1, y1), (x2, y2), color, 3)82 83        # Responsive label84        box_w = x2 - x185        box_h = y2 - y186 87        font_scale = max(0.6, min(2.0, box_w / 180))88        font_scale = min(font_scale, box_h / 80)89        thickness = max(1, int(font_scale * 2))90 91        label = f"{CLASS_NAMES[cls]} {conf:.2f}"92 93        (tw, th), base = cv2.getTextSize(94            label,95            cv2.FONT_HERSHEY_SIMPLEX,96            font_scale,97            thickness98        )99 100        tx = x1101        ty = y1 - 10102        if ty - th < 0:103            ty = y1 + th + 10104 105        cv2.rectangle(106            img,107            (tx, ty - th - base),108            (tx + tw, ty + base),109            color,110            cv2.FILLED111        )112 113        cv2.putText(114            img,115            label,116            (tx, ty),117            cv2.FONT_HERSHEY_SIMPLEX,118            font_scale,119            (255, 255, 255),120            thickness,121            cv2.LINE_AA122        )123 124    img_out = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)125    return img_out, civil_count, military_count126 127# --------------------------128# Gradio UI129# --------------------------130with gr.Blocks() as demo:131    gr.Markdown("# 🛩️ Plane Detection (Civil vs Military)")132 133    with gr.Row():134        with gr.Column():135            img_in = gr.Image(label="Upload Image", type="numpy")136            mode = gr.Radio(137                ["Normal YOLO Inference", "Tiled Inference"],138                value="Normal YOLO Inference"139            )140            btn = gr.Button("▶ Run Detection", variant="primary")141 142        with gr.Column():143            img_out = gr.Image(label="Prediction Output")144            civil_out = gr.Number(label="Civil Planes Detected")145            military_out = gr.Number(label="Military Planes Detected")146 147    btn.click(148        predict,149        inputs=[img_in, mode],150        outputs=[img_out, civil_out, military_out]151    )152 153demo.launch()154