|
import gradio as gr |
|
import cv2 |
|
import requests |
|
import os |
|
import numpy as np |
|
from ultralytics import YOLO |
|
import math |
|
|
|
model = YOLO('yolo-Weights/best.pt') |
|
|
|
|
|
classNames = ["book","lego","plush"] |
|
|
|
def show_preds_image(image): |
|
if isinstance(image, str): |
|
image = cv2.imread(image) |
|
elif isinstance(image, np.ndarray): |
|
pass |
|
else: |
|
raise ValueError("Invalid input type for image") |
|
|
|
frame_copy = image.copy() |
|
|
|
outputs = model.predict(source=image) |
|
results = outputs[0].cpu().numpy() |
|
|
|
|
|
|
|
for r in results: |
|
boxes = r.boxes |
|
|
|
for box in boxes: |
|
|
|
x1, y1, x2, y2 = box.xyxy[0] |
|
x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2) |
|
|
|
|
|
|
|
confidence = math.ceil((box.conf[0]*100))/100 |
|
print("Confidence --->",confidence) |
|
|
|
if confidence > 0.7: |
|
|
|
cv2.rectangle(image, (x1, y1), (x2, y2), (255, 0, 255), 3) |
|
|
|
|
|
|
|
cls = int(box.cls[0]) |
|
print("Class name -->", classNames[cls]) |
|
|
|
|
|
org = [x1, y1] |
|
font = cv2.FONT_HERSHEY_SIMPLEX |
|
fontScale = 1 |
|
color = (255, 0, 0) |
|
thickness = 2 |
|
|
|
cv2.putText(image, classNames[cls], org, font, fontScale, color, thickness) |
|
|
|
yield image |
|
|
|
inputs_video = [ |
|
gr.Image(label="Input Video", sources="webcam"), |
|
|
|
] |
|
outputs_video = [ |
|
gr.components.Image(type="numpy", label="Output Image"), |
|
] |
|
interface_video = gr.Interface( |
|
fn=show_preds_image, |
|
inputs=inputs_video, |
|
outputs=outputs_video, |
|
title="Plush-lego-book detector", |
|
|
|
cache_examples=False, |
|
) |
|
|
|
interface_video.queue().launch() |