Yolo / app.py
GiaoH's picture
Upload 3 files
34b002e verified
raw
history blame
No virus
4.75 kB
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)
cv2.rectangle(
image,
(int(x1), int(y1), int(x2), int(y2)),
color=(0, 0, 255),
thickness=2,
lineType=cv2.LINE_AA
)
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)
def show_preds_webcam(pil_image):
image= np.array(pil_image)
outputs = model.predict(image)
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)
cv2.rectangle(
image,
(int(x1), int(y1), int(x2), int(y2)),
color=(0, 0, 255),
thickness=2,
lineType=cv2.LINE_AA
)
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 image
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,
)
outputs_video = [
gr.components.Image(type="numpy", label="Output Image"),
]
interface_webcam = gr.Interface(
fn=show_preds_webcam,
live=True,
inputs=gr.Image(source="webcam", streaming=True, type="pil"),
outputs=outputs_video,
)
gr.TabbedInterface(
[ interface_webcam, interface_image],
tab_names=[ 'Webcam', "Image"]
).queue().launch()