Felipefe20's picture
Update app.py
65eca32 verified
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')
# object classes
classNames = ["book","lego","plush"]
def show_preds_image(image):
if isinstance(image, str): # If the input is a string, treat it as a file path
image = cv2.imread(image)
elif isinstance(image, np.ndarray): # If the input is a numpy array, treat it as an image
pass
else:
raise ValueError("Invalid input type for image")
frame_copy = image.copy()
#results=model.predict(source=image)
outputs = model.predict(source=image)
results = outputs[0].cpu().numpy()
#print(results.boxes.xyxy)
#print("end of results")
# coordinates
for r in results:
boxes = r.boxes
for box in boxes:
# bounding box
x1, y1, x2, y2 = box.xyxy[0]
x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2) # convert to int values
# confidence
confidence = math.ceil((box.conf[0]*100))/100
print("Confidence --->",confidence)
if confidence > 0.7:
# put box in cam
cv2.rectangle(image, (x1, y1), (x2, y2), (255, 0, 255), 3)
# class name
cls = int(box.cls[0])
print("Class name -->", classNames[cls])
# object details
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",
#examples=video_path,
cache_examples=False,
)
interface_video.queue().launch()