Yolo / app.py
Huu076's picture
Update app.py
17d191e verified
import gradio as gr
import cv2
import requests
import os
import cvzone
import math
import random
import numpy as np
from ultralytics import YOLO
model = YOLO('yolov8x-seg.pt')
path = []
video_path = []
listClasses = ['person', 'bicycle', 'car']
def show_preds_image(image_path):
image = cv2.imread(image_path)
outputs = model.predict(source=image_path)
results = outputs[0].cpu().numpy()
yolo_classes = list(model.names.values())
classes_ids = [yolo_classes.index(clas) for clas in yolo_classes]
colors = [random.choices(range(256), k=3) for _ in classes_ids]
for result in outputs:
for mask, box in zip(result.masks.xy, result.boxes):
#for r in results:
#boxes = r.boxes
#for box in boxes:
cls = box.cls[0]
conf = math.ceil((box.conf[0]*100))/100
if (int(cls)<3) and (conf > 0.70):
points = np.int32([mask])
cv2.polylines(img, points, True, (255, 0, 0), 1)
color_number = classes_ids.index(int(box.cls[0]))
color = colors[color_number]
cv2.fillPoly(image, points, color)
x1, y1, x2, y2 = box.xyxy[0]
x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
name = yolo_classes[int(cls)]
# fontScale
fontScale = 0.5
color_number = classes_ids.index(int(box.cls[0]))
color = colors[color_number]
# Line thickness of 2 px
thickness = 1
font = cv2.FONT_HERSHEY_SIMPLEX
cv2.putText(image, str(name) + " " + str(conf), (max(0,x1), max(35,y1)), font,
fontScale, color, thickness, cv2.LINE_AA)
return cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
inputs_image = [
gr.components.Image(type="filepath", label="Input Image"),
]
outputs_image = [
gr.components.Image(type="numpy", label="Output Image"),
]
interface_image = gr.Interface(
fn=show_preds_image,
inputs=inputs_image,
outputs=outputs_image,
title="Object segmentation",
examples=path,
cache_examples=False,
).launch()