File size: 1,599 Bytes
cea4224 5f0bc09 d943e69 7997fdd 015372b cea4224 02f91fb cea4224 |
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 |
import gradio as gr
import sys
sys.path.append('./utils')
from yolo_utils import preprocess_image_pil, run_model, process_results, plot_results_gradio
import matplotlib.pyplot as plt
import io
try:
from ultralytics import YOLO
except ImportError:
import os
os.system('pip install ./yolov8-to')
from ultralytics import YOLO
def process_image(image,conf,iou):
model = YOLO('./trained_models/nano.pt')
# Preprocess the image
preprocessed_image = preprocess_image_pil(image, threshold_value=0.9, upscale=False)
# Run the model
results = run_model(model, preprocessed_image, conf=conf, iou=iou, imgsz=640)
# Process the results
input_image_array_tensor, seg_result, pred_Phi, sum_pred_H, final_H, dice_loss, tversky_loss = process_results(results, preprocessed_image)
# Plot the results
fig = plot_results_gradio(input_image_array_tensor, seg_result, pred_Phi, sum_pred_H, final_H, dice_loss, tversky_loss)
# Convert the plot to an image
return fig
# Create the Gradio interface
title = "YOLOV8-TO Demo App"
description = "Upload an image and see the processed results. Adjust the confidence and IOU thresholds as needed. Runs the YOLOv8-TO Nano model size. Runs on 2 CPU cores so please be patient!"
iface = gr.Interface(
fn=process_image,
inputs=[
gr.Image(type='pil'),
gr.Slider(minimum=0, maximum=1, value=0.1, label="Confidence Threshold"),
gr.Slider(minimum=0, maximum=1, value=0.5, label="IOU Threshold")
],
outputs="image",
title=title,
description=description
)
iface.launch()
|