Samrumi67/Ai_Enemy_Detector

1import streamlit as st2import torch3import cv24import numpy as np5import tempfile6from PIL import Image7import base648import colorsys9 10# Page settings11st.set_page_config(page_title="AI Enemy Detector", layout="wide")12 13# Radar-style background and Calibri font for all text14st.markdown(15    f"""16    <style>17    .stApp {{18        background-image: url("https://i.imghippo.com/files/TXyG7770rkk.jpg");19        background-size: cover;20        background-attachment: fixed;21        background-repeat: no-repeat;22        background-position: center;23    }}24    h1, h2, h3, .stMarkdown {{25        color: white;26        text-shadow: 1px 1px 2px black;27        font-family: Calibri; /* Apply Calibri font */28    }}29    h1 {{30        font-size: 24px; /* Adjust the size of the main heading */31    }}32    </style>33    """,34    unsafe_allow_html=True,35)36 37# Title38st.markdown(39    "<h1 style='font-size: 24px; font-family: Calibri;'>AI Enemy Detector🛡</h1>",40    unsafe_allow_html=True,41)  # Smaller title in Calibri42st.markdown("**Detect enemy drones or airplanes using AI-powered vision.**")43 44# Load model45@st.cache_resource46def load_model():47    model = torch.hub.load(48        "ultralytics/yolov5", "yolov5n", pretrained=True49    )  # Smaller model50    model.cpu()51    return model52 53 54model = load_model()55 56 57# Function to calculate the dominant color of an image58def get_dominant_color(image):59    """60    Calculates the dominant color (in RGB) of a given image.  Handles edge cases.61    """62    if image is None or image.size == 0:63        return (0, 0, 0)  # Return black for empty or invalid images64 65    # Resize the image for faster processing, while preserving aspect ratio66    height, width = image.shape[:2]67    max_size = 100  # Reduced size for faster processing68    if height > width:69        new_height = max_size70        new_width = int(width * (max_size / height))71    else:72        new_width = max_size73        new_height = int(height * (max_size / width))74    resized_image = cv2.resize(75        image, (new_width, new_height), interpolation=cv2.INTER_AREA76    )77 78    pixels = np.float32(resized_image.reshape(-1, 3))79 80    n_colors = 5  # Consider the top 5 dominant colors81    criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 10, 0.1)82    flags = cv2.KMEANS_RANDOM_CENTERS83 84    try:85        _, labels, palette = cv2.kmeans(pixels, n_colors, None, criteria, 10, flags)86    except:87        return (0, 0, 0)  # error case88 89    _, counts = np.unique(labels, return_counts=True)90 91    dominant_color = tuple(map(int, palette[np.argmax(counts)]))92    return dominant_color93 94 95# Function to detect Pakistan flag colors in cropped image96def contains_pakistan_flag_color(crop_img):97    """98    Check if cropped image contains predominant green + white colors similar to99    Pakistan flag.100    """101    if crop_img is None or crop_img.size == 0:102        return False103 104    hsv = cv2.cvtColor(crop_img, cv2.COLOR_BGR2HSV)105 106    # Define a wider range for green to be more inclusive107    lower_green = np.array([30, 20, 20])  # Adjusted lower bound108    upper_green = np.array([90, 255, 255])  # Adjusted upper bound109    green_mask = cv2.inRange(hsv, lower_green, upper_green)110    green_area = np.sum(green_mask > 0)111 112    # Define a range for white113    lower_white = np.array([0, 0, 200])114    upper_white = np.array([180, 30, 255])  # Allow some slightly off-white colors115    white_mask = cv2.inRange(hsv, lower_white, upper_white)116    white_area = np.sum(white_mask > 0)117    total_area = crop_img.shape[0] * crop_img.shape[1]118    # Check if green is dominant and white is present119    if green_area > total_area * 0.2 and white_area > total_area * 0.05:120        return True121    return False122 123 124def is_friendly_aircraft(label, crop_img, detected_classes):125    """126    Determine if an aircraft is friendly.127    """128    if label == "airplane":129        dominant_color = get_dominant_color(crop_img)130        r, g, b = dominant_color131        h, s, v = colorsys.rgb_to_hsv(r / 255.0, g / 255.0, b / 255.0)132        is_white = v > 0.8 and s < 0.2133        return is_white or contains_pakistan_flag_color(crop_img)134    return False  # Drones are not friendly.135 136 137# Detect objects in image/frame138def detect_objects(image):139    results = model(image)140    return results141 142 143# Draw bounding boxes and label friendly/enemy jets144def draw_results(img, results):145    detected_enemy = False146    detected_classes = (147        []148    )  # To store detected class names.  Important for Airbus detection.149    for *box, conf, cls in results.xyxy[0]:150        label = model.names[int(cls)]151        detected_classes.append(label)152        if label in ["drone", "airplane"]:153            x1, y1, x2, y2 = map(int, box)154            crop_img = img[y1:y2, x1:x2]155            if is_friendly_aircraft(label, crop_img, detected_classes):156                color = (0, 255, 0)  # Green for friendly157                text = f"FRIENDLY {label.upper()} ({conf:.2f})"158            else:159                color = (0, 0, 255)  # Red for enemy160                text = f"POTENTIAL ENEMY {label.upper()} ({conf:.2f})"161                detected_enemy = True162 163            cv2.rectangle(img, (x1, y1), (x2, y2), color, 2)164            cv2.putText(165                img, text, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 2166            )167    return img, detected_enemy, detected_classes168 169 170# Load the siren file as base64171@st.cache_data172def load_siren_audio():173    try:174        with open("siren.mp3", "rb") as f:175            audio_bytes = f.read()176        return base64.b64encode(audio_bytes).decode("utf-8")177    except FileNotFoundError:178        st.error(179            "Siren file not found! Please ensure 'siren.mp3' is in the same directory."180        )181        return None  # Important: Return None in case of error182 183 184siren_audio_base64 = load_siren_audio()185 186 187# Log detections and play siren if enemy detected188def log_detections(detected_enemy):189    if detected_enemy:190        st.error("Potential Enemy Aircraft Detected - Siren Activated!")191        if siren_audio_base64:  # Check if the siren loaded successfully192            # Use HTML to embed and autoplay the audio193            st.markdown(194                f"""195                <audio autoplay style="display:none;">196                    <source src="data:audio/mp3;base64,{siren_audio_base64}" type="audio/mp3">197                    Your browser does not support the audio element.198                </audio>199                """,200                unsafe_allow_html=True,201            )202    else:203        st.success("No potential enemy aircraft or drone detected.")204 205 206# Upload section207uploaded_file = st.file_uploader(208    "Upload an aerial image or video", type=["jpg", "jpeg", "png", "mp4"]209)210 211# Main logic212if uploaded_file:213    if uploaded_file.type.startswith("image"):214        img = Image.open(uploaded_file).convert("RGB")215        img_np = np.array(img)216        results = detect_objects(img_np)217        output_img, enemy_detected, detected_classes = draw_results(218            img_np.copy(), results219        )220        st.image(output_img, channels="RGB")221        log_detections(enemy_detected)222 223    elif uploaded_file.type == "video/mp4":224        tfile = tempfile.NamedTemporaryFile(delete=False)225        tfile.write(uploaded_file.read())226        cap = cv2.VideoCapture(tfile.name)227        stframe = st.empty()228 229        while cap.isOpened():230            ret, frame = cap.read()231            if not ret:232                break233            results = detect_objects(frame)234            frame, enemy_detected, detected_classes = draw_results(frame, results)235            stframe.image(frame, channels="BGR")236            log_detections(enemy_detected)237        cap.release()