BoraBasaran/Drone-Detection-and-Classification

1import gradio as gr2import numpy as np3import os4import cv25from PIL import Image6from tensorflow.keras.models import load_model7from tensorflow.keras.applications.imagenet_utils import preprocess_input8import torch9import seaborn10 11# ========== Load Drone Classifier Models ==========12model_files = {13    "CNN": "CNN.h5",14    "ResNet": "ResNet50.h5",15    "EfficientNet": "EfficientNet.h5",16    "MobileNet": "MobileNetV2.h5",17    "VGG16": "VGG16.h5",18    "VGG19": "VGG19.h5",19    "DenseNet": "DenseNet.h5"20}21 22models = {}23for name, path in model_files.items():24    try:25        models[name] = load_model(path)26    except Exception as e:27        print(f"Error loading model {name}: {e}")28        models[name] = None29 30def preprocess_image(image_np, model_name):31    image_resized = cv2.resize(image_np, (224, 224))32 33    if model_name in ["ResNet", "MobileNet", "VGG16", "VGG19", "EfficientNet", "DenseNet"]:34        image_array = preprocess_input(image_resized.astype(np.float32))35    else:36        image_array = image_resized.astype(np.float32) / 255.037 38    return np.expand_dims(image_array, axis=0)39 40def get_examples(base_path="Drone-NonDrone/data/test"):41    examples = []42    for class_folder in os.listdir(base_path):43        class_folder_path = os.path.join(base_path, class_folder)44        if os.path.isdir(class_folder_path):45            image_files = [46                f for f in os.listdir(class_folder_path)47                if f.lower().endswith(('.jpg', '.jpeg', '.png'))48            ]49            for image_file in image_files:50                example_path = os.path.join(class_folder_path, image_file)51                examples.append([example_path])52    return examples53 54examples = get_examples()55 56model_thresholds = {57    "CNN": 0.1,58    "ResNet": 0.15,59    "EfficientNet": 0.3,60    "MobileNet": 0.2,61    "VGG16": 0.25,62    "VGG19": 0.25,63    "DenseNet": 0.2,64}65 66def classify_drone(image, selected_models):67    predictions = {}68 69    for name in selected_models:70        threshold = model_thresholds.get(name, 0.3)  # default to 0.3 if missing71        model = models.get(name)72        if model is None:73            predictions[name] = "Model Error"74            continue75        try:76            input_image = preprocess_image(image, name)77            prob = model.predict(input_image)[0][0]78            print(f"{name} prediction probability: {prob:.4f}")79            label = "Drone" if prob > threshold else "Not Drone"80            predictions[name] = f"{label} ({prob:.2f})"81        except Exception as e:82            predictions[name] = "Prediction Error"83 84    return predictions85 86 87# Load the custom YOLOv5 model88yolo_model = torch.hub.load('ultralytics/yolov5', 'custom', path='yolov5s.pt', force_reload=True)89 90## Define aerial object labels and corresponding colors91aerial_objects = {92    "helicopter": (255, 0, 0),  # Blue93    "drone": (0, 0, 255),       # Red94    "kite": (255, 255, 0),      # Cyan95    "bird": (0, 255, 255),      # Yellow96    "airplane": (255, 0, 255)   # Magenta97}98 99def yolo_detect(image):100    img_bgr = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)101    results = yolo_model(img_bgr)102 103    labels, cords = results.xyxyn[0][:, -1], results.xyxyn[0][:, :-1]104    n = len(labels)105    if n == 0:106        return image107 108    h, w, _ = image.shape109    for i in range(n):110        row = cords[i]111        conf = row[4]112        if conf < 0.3:113            continue114 115        x1, y1, x2, y2 = int(row[0]*w), int(row[1]*h), int(row[2]*w), int(row[3]*h)116        cls = int(labels[i])117        label = yolo_model.names[cls].lower() if hasattr(yolo_model, 'names') else 'object'118 119        # Heuristic: Re-label small 'airplane' as 'drone'120        box_width = x2 - x1121        box_height = y2 - y1122        area = box_width * box_height123        image_area = w * h124 125        if label == "airplane" and area < 0.01 * image_area:126            label = "drone"127 128        # Choose color129        color = aerial_objects.get(label, (0, 255, 0))  # Default green130 131        cv2.rectangle(img_bgr, (x1, y1), (x2, y2), color, 2)132        cv2.putText(img_bgr, f"{label} {conf:.2f}", (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 2)133 134    img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)135    return img_rgb136 137# ========== Gradio Interface ==========138interface = gr.Interface(139    fn=classify_drone,140    inputs=[141        gr.Image(type="numpy", label="Upload Image"),142        gr.CheckboxGroup(choices=list(model_files.keys()), label="Select Models", value=["CNN"])143    ],144    outputs=gr.JSON(label="Model Predictions"),145    title="Drone Classification (Binary)",146    examples=[[ex[0], ["CNN", "ResNet", "EfficientNet","MobileNet","VGG16","VGG19","DenseNet"]] for ex in examples]147)148 149# Your YOLOv5 detection interface150yolo_interface = gr.Interface(151    fn=yolo_detect,152    inputs=gr.Image(type="numpy", label="Upload Image for Object Detection"),153    outputs=gr.Image(type="numpy", label="Detected Objects"),154    title="YOLOv5 Drone & Object Detection",155    description="Detects drones, helicopters, kites, and birds with bounding boxes.",156    examples = examples157)158 159# ========== Launch ==========160app = gr.TabbedInterface(161    interface_list=[interface, yolo_interface],162    tab_names=["Drone Binary Classification", "YOLOv5 Object Detection"]163)164 165if __name__ == "__main__":166    app.launch()167