Initial commit

.gitignore

@@ -0,0 +1,3 @@
+**/.DS_Store
+**/__pycache__/
+weights/

README.md

@@ -1,9 +1,11 @@
 ---
-title: Aircraft Detection Optical Satellite
-emoji: 🏆
+title: Aircraft Detection in Optical Satellite images
+emoji: ✈️
 colorFrom: indigo
-colorTo: red
-sdk: docker
+colorTo: blue
+sdk: gradio
+sdk_version: 4.19.2
+app_file: app.py
 pinned: false
 license: cc-by-nc-sa-4.0
 ---

app.py

@@ -0,0 +1,176 @@
+import os
+import socket
+import gradio as gr
+import numpy as np
+from PIL import Image, ImageDraw
+from pathlib import Path
+from loguru import logger
+import cv2
+import torch
+import ultralytics
+from ultralytics import YOLO
+import time
+import base64
+import requests
+import json
+ 
+# API for inferences
+DL4EO_API_URL = "https://dl4eo--predict.modal.run"
+
+# Auth Token to access API
+DL4EO_API_KEY = os.environ['DL4EO_API_KEY']
+
+# width of the boxes on image
+LINE_WIDTH = 2
+
+# Load a model if weights are present
+WEIGHTS_FILE = './weights/best.pt'
+model = None
+if os.path.exists(WEIGHTS_FILE):
+    model = YOLO(WEIGHTS_FILE)  # previously trained YOLOv8n model
+    logger.info(f"Setup for local inference")
+
+    # check if GPU if available
+    device = torch.device("cuda:0" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu")
+    logger.info(f"Using device: {device}")
+
+    # Check Ultralytics modules version
+    logger.info(f"Ultralytics version: {ultralytics.__version__}")
+
+# Check Gradio modules version
+logger.info(f"Gradio version: {gr.__version__}")
+
+# Define the inference function
+def predict_image(image, threshold):
+    
+    # Resize the image to the new size
+    #image = image.resize((image.size[0] * 2, image.size[1] * 2))
+
+    if isinstance(image, Image.Image):
+        img = np.array(image)
+    
+    if not isinstance(img, np.ndarray) or len(img.shape) != 3 or img.shape[2] != 3:
+        raise BaseException("predit_image(): input 'img' shoud be single RGB image in PIL or Numpy array format.")
+
+    width, height = img.shape[0], img.shape[1]
+
+    if model is None:
+        # Encode the image data as base64
+        image_base64 = base64.b64encode(np.ascontiguousarray(img)).decode()
+        
+        # Create a dictionary representing the JSON payload
+        payload = {
+            'image': image_base64,
+            'shape': img.shape,
+            'threshold': threshold,
+        }
+
+        headers = {
+            'Authorization': 'Bearer ' + DL4EO_API_KEY,
+            'Content-Type': 'application/json'  # Adjust the content type as needed
+        }
+
+        # Send the POST request to the API endpoint with the image file as binary payload
+        response = requests.post(DL4EO_API_URL, json=payload, headers=headers)
+        
+        # Check the response status
+        if response.status_code != 200:
+            raise Exception(
+                f"Received status code={response.status_code} in inference API"
+            )
+                
+        json_data = json.loads(response.content)
+        duration = json_data['duration']
+        boxes = json_data['boxes']
+    else:
+        start_time = time.time()
+        results = model.predict([img], imgsz=(width, height), conf=threshold)
+        end_time = time.time()
+        boxes = [box.xyxy.cpu().squeeze().int().tolist() for box in boxes]
+        duration = end_time - start_time
+        boxes = results[0].boxes
+
+    # drow boxes on image
+    draw = ImageDraw.Draw(image)
+
+    for box in boxes:
+        left, top, right, bottom = box
+
+        if left <= 0: left = -LINE_WIDTH 
+        if top <= 0: top = top - LINE_WIDTH 
+        if right >= img.shape[0] - 1: right = img.shape[0] - 1 + LINE_WIDTH
+        if bottom >= img.shape[1] - 1: bottom = img.shape[1] - 1 + LINE_WIDTH
+
+        draw.rectangle([left, top, right, bottom], outline="red", width=LINE_WIDTH)
+    
+    return image, str(image.size), len(boxes), duration
+
+
+# Define example images and their true labels for users to choose from
+example_data = [
+    ["./demo/airport01.jpg", 0.50],
+    ["./demo/airport02.jpg", 0.50],
+    ["./demo/airport03.jpg", 0.50],
+    ["./demo/airport04.jpg", 0.50],
+    ["./demo/Pleiades_Neo_Tucson_USA.jpg", 0.50],
+]
+
+# Define CSS for some elements
+css = """
+  .image-preview {
+    height: 820px !important; 
+    width: 800px !important;
+  } 
+"""
+TITLE = "Aircraft detection with YOLOv8"
+
+# Define the Gradio Interface
+demo = gr.Blocks(title=TITLE, css=css).queue()
+with demo:
+    gr.Markdown(f"<h1><center>{TITLE}<center><h1>")
+    #gr.Markdown("<p>This demo is provided by <a href='https://www.linkedin.com/in/faudi/'>Jeff Faudi</a> and <a href='https://www.dl4eo.com/'>DL4EO</a></p>")
+
+    with gr.Row():
+        with gr.Column(scale=0):
+            input_image = gr.Image(type="pil", interactive=True, scale=1)
+            run_button = gr.Button(value="Run", scale=0)
+            with gr.Accordion("Advanced options", open=True):
+                threshold = gr.Slider(label="Confidence threshold", minimum=0.0, maximum=1.0, value=0.25, step=0.01)
+                dimensions = gr.Textbox(label="Image size", interactive=False)
+                detections = gr.Number(label="Predicted aircrafts", interactive=False)
+                stopwatch = gr.Number(label="Execution time (sec.)", interactive=False, precision=3)
+
+        with gr.Column(scale=2):
+            output_image = gr.Image(type="pil", elem_classes='image-preview', interactive=False, width=800, height=800)
+
+    run_button.click(fn=predict_image, inputs=[input_image, threshold], outputs=[output_image, dimensions, detections, stopwatch])
+    gr.Examples(
+        examples=example_data,
+        inputs = [input_image, threshold],
+        outputs = [output_image, dimensions, detections, stopwatch],
+        fn=predict_image,
+        #cache_examples=True,
+        label='Try these images!'
+    )
+
+    gr.Markdown("<p>This demo is provided by <a href='https://www.linkedin.com/in/faudi/'>Jeff Faudi</a> and <a href='https://www.dl4eo.com/'>DL4EO</a>. The model has been trained with the <a href='https://www.ultralytics.com/yolo'>Ultralytics YOLOv8</a> framework on the <a href='https://www.cosmiqworks.org/rareplanes/'>RarePlanes</a>, <a href='https://www.kaggle.com/datasets/airbusgeo/airbus-aircrafts-sample-dataset'>Airbus</a> and <a href='https://github.com/dilsadunsal/HRPlanesv2-Data-Set'>HRPlanesv2</a> datasets. The associated licenses are <a href='https://about.google/brand-resource-center/products-and-services/geo-guidelines/#google-earth-web-and-apps'>GoogleEarth fair use</a> and <a href='https://creativecommons.org/licenses/by-nc-sa/4.0/deed.en'>CC-BY-SA-NC</a>. This demonstration CANNOT be used for commercial puposes. Please contact <a href='mailto:jeff@dl4eo.com'>me</a> for more information on how you could get access to a commercial grade model or API. </p>")
+
+if os.path.exists('/.dockerenv'):
+    print('Running inside a Docker container')
+
+    # Launch the interface on MacOS
+    hostname = socket.gethostname()
+
+    demo.launch(
+        server_name=hostname, 
+        inline=False, 
+        server_port=7860, 
+        debug=True
+    )
+else:
+    print('Not running inside a Docker container')
+    demo.launch(
+        inline=False, 
+        #server_port=7860, 
+        debug=False
+    )

requirements.txt

@@ -0,0 +1,3 @@
+loguru
+ultralytics==8.1.18
+gradio==3.35.2