Update app.py

app.py

@@ -1,41 +1,36 @@
-import sys, os
+import os
 
 import gradio as gr
 import plotly.express as px
 import numpy as np
-import random
-from ultralytics import YOLO
-#from sahi.models.yolov8 import *
-from src.sahi_onnx import *
+from sahi.models.ultralytics import *
 from sahi.predict import get_sliced_prediction
 from sahi.utils.cv import visualize_object_predictions
 import PIL
 
-#model_base = "https://huggingface.co/mayrajeo/marine-vessel-detection/resolve/main/"
-model_base = 'onnx_models'
+MODEL_PATH = 'yolo11s_tci.pt'
 
 def inference(
     im:gr.Image=None,
-    model_path:gr.Dropdown='YOLOv8n',
     conf_thr:gr.Slider=0.25
 ):
-    #model = Yolov8DetectionModel(model_path=f'{model_base}/{model_path}/{model_path}.pt',
-    model = Yolov8onnxDetectionModel(model_path=f'{model_base}/{model_path}/{model_path.lower()}.onnx',
-                                     config_path=f'{model_base}/{model_path}/args.yaml',
-                                 device='cpu',
-                                 confidence_threshold=conf_thr,
-                                 category_mapping={'0': 'Boat'},
-                                 image_size=640)
+    model = UltralyticsDetectionModel(model_path=MODEL_PATH,
+                                      device='cuda',
+                                      confidence_threshold=conf_thr,
+                                      category_mapping={'0': 'Boat'},
+                                      image_size=640)
 
     res = get_sliced_prediction(im, model, slice_width=320, 
                                 slice_height=320, overlap_height_ratio=0.2,
                                 overlap_width_ratio=0.2, verbose=0)
     img = PIL.Image.open(im)
-    visual_result = visualize_object_predictions(image=np.array(img),
+    img = PIL.ImageEnhance.Brightness(img)
+    img = img.enhance(2.0)
+    visual_result = visualize_object_predictions(image=np.array(img), color=(100, 193, 203),
                                                  object_prediction_list=res.object_prediction_list,
-                                                 text_size=0.4,
+                                                 text_size=0.4, hide_labels=True, hide_conf=True,
                                                  rect_th=1)
-    fig = px.imshow(visual_result['image'])
+    fig = px.imshow(visual_result['image'], width=900, height=900)
     fig.update_layout(showlegend=False, hovermode=False)
     fig.update_xaxes(visible=False)
     fig.update_yaxes(visible=False)
@@ -43,15 +38,7 @@ def inference(
 
 inputs = [
     gr.Image(type='filepath', label='Input'),
-    gr.Dropdown([
-        'YOLOv8n', 
-        'YOLOv8s', 
-        'YOLOv8m', 
-        'YOLOv8l', 
-        'YOLOv8x'
-        ],
-        value='YOLOv8n', label='Model'),
-    gr.Slider(minimum=0.0, maximum=1.0, value=0.25, step=0.05, label='Confidence Threshold'),
+    gr.components.Slider(minimum=0.0, maximum=1.0, value=0.25, step=0.05, label='Confidence Threshold'),
 ]
 
 outputs = [
@@ -61,7 +48,6 @@ outputs = [
 example_images = [[f'examples/{f}'] for f in os.listdir('examples')] 
 
 
-
 gr.Interface(
     fn=inference,
     inputs=inputs,
@@ -71,9 +57,9 @@ gr.Interface(
     cache_examples=False,
     examples_per_page=10,
     title='Marine vessel detection from Sentinel 2 images',
-    description="""Models detect potential marine vessels from Sentinel 2 imagery. 
-    Each example image covers 7.68x7.68 km (768x768 pixels). 
+    description="""Model detects potential marine vessels from Sentinel 2 imagery. 
+    Each example image covers 7.68x7.68 km (768x768 pixels). The result image has its brightness increased, 
+    but the predictions are made based on Sentinel-2 L1C-TCI data.
     As we don't clean the prediction with stationary targets that look like vessels in this resolution,
-    there will most likely be false positives from lighthouses, above-water rocks and on land.\n
-    Be patient with responses, as free tier only has 2vCPUs so app might be slow sometimes."""
-    ).launch()
+    there will most likely be false positives from lighthouses, above-water rocks and on land."""
+    ).launch(share=True, server_name='0.0.0.0')