mussie1212's picture
Update app.py
35f1caf verified
from PIL import Image
import gradio as gr
# import gradio as gr
import PIL.Image as Image
from ultralytics import ASSETS, YOLO
model_path = 'new_data_improved_object_detector.pt'
model = YOLO(model_path)
from PIL import Image
import gradio as gr
object_detector_model_path = 'new_data_improved_object_detector.pt'
# logo_detector_model_path ='logo_detector_grayscale_v2.pt'
logo_detector_model_path = 'logo_detector_june_1.pt'
object_model = YOLO(object_detector_model_path)
logo_model = YOLO(logo_detector_model_path)
def Get_logo_xywh(model_result_input):
model_result = model_result_input[0]
xywh = model_result.boxes.xywh.cpu().tolist()
clss = model_result.boxes.cls.cpu().tolist()
# names = model_result_input[0].names
confidence = model_result.boxes.conf.cpu().tolist()
xyxy = model_result.boxes.cpu().xyxy.tolist()
return xywh, clss, confidence, xyxy
def predict_image(img, conf_threshold, iou_threshold):
"""Predicts and plots labeled objects in an image using YOLOv8 model with adjustable confidence and IOU thresholds."""
resized_image = img.resize((640, 640))
# Convert the input image to grayscale
img = resized_image.convert('L')
logo_results = logo_model.predict(
source=img,
conf=conf_threshold,
iou=iou_threshold,
show_labels=True,
show_conf=True,
imgsz=640,
)
logo_im_arrays = []
for r in logo_results:
im_array = r.plot()
im = Image.fromarray(im_array[..., ::-1])
logo_im_arrays.append(im)
object_result = object_model.predict(
source=img,
conf=conf_threshold,
iou=iou_threshold,
show_labels=True,
show_conf=True,
imgsz=640,
)
im_arrays = []
for r in object_result:
im_array = r.plot()
im = Image.fromarray(im_array[..., ::-1])
im_arrays.append(im)
logo_xywh, logo_clss, logo_confidence, logo_xyxy = Get_logo_xywh(logo_results)
object_xywh, object_clss, object_confidence, object_xyxy = Get_logo_xywh(object_result)
return logo_im_arrays,im_arrays, logo_xywh, logo_clss, logo_confidence, logo_xyxy,object_xywh, object_clss, object_confidence, object_xyxy
iface = gr.Interface(
fn=predict_image,
inputs=[
gr.Image(type="pil", label="Upload Image"),
gr.Slider(minimum=0, maximum=1, value=0.25, label="Confidence threshold"),
gr.Slider(minimum=0, maximum=1, value=0.45, label="IoU threshold"),
],
outputs=[
gr.Gallery(label="logo Images"),
gr.Gallery(label="object Images"),
gr.JSON(label="Detection Bounding Boxes (l_xywh)"),
gr.JSON(label="Detection Class Indices"),
gr.JSON(label="Detection Confidence Scores"),
gr.JSON(label="Detection Bounding Boxes (l_xyxy)"),
gr.JSON(label="Detection Bounding Boxes (l_xywh)"),
gr.JSON(label="Detection Class Indices"),
gr.JSON(label="Detection Confidence Scores"),
gr.JSON(label="Detection Bounding Boxes (l_xyxy)"),
],
title="Ultralytics Gradio",
description="Upload images for inference. The Ultralytics YOLOv8n model is used by default.",
)
if __name__ == "__main__":
iface.launch(show_error=True)