File size: 6,412 Bytes
7dfcc1a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 |
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--oil-storage-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__}")
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/588fc1fb-b86a-4fb4-8161-d9bd3a1556ca.jpg", 0.50],
["./demo/605ffac0-69d5-4748-92c2-48d43f51afc1.jpg", 0.50],
["./demo/67f7c7ad-11a1-4c7f-9f2a-da7ef50bfdd8.jpg", 0.50],
["./demo/b8c0e212-3669-4ff8-81a5-32191d456f86.jpg", 0.50],
["./demo/df5ec618-c1f3-4cfe-88b1-86799d23c22d.jpg", 0.50]]
# Define CSS for some elements
css = """
.image-preview {
height: 820px !important;
width: 800px !important;
}
"""
TITLE = "Oil storage detection on SPOT images (1.5 m) 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 objects", 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! They are not included in the training dataset.'
)
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.kaggle.com/datasets/airbusgeo/airbus-oil-storage-detection-dataset'>Airbus Oil Storage Dataset</a>.
The associated license is <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>""")
demo.launch(
inline=False,
show_api=False,
debug=False
)
|