File size: 1,380 Bytes
b962de2 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 |
# Importamos librerias
from ultralytics import YOLO
import cv2
import math
# Modelo
model = YOLO('Modelos/best.pt')
# Cap
cap = cv2.VideoCapture(0)
cap.set(3, 1280)
cap.set(4, 720)
# Clases
clsName = ['Metal', 'Glass', 'Plastic', 'Carton', 'Medical']
# Inference
while True:
# Frames
ret, frame = cap.read()
# Yolo | AntiSpoof
results = model(frame, stream=True, verbose=False)
for res in results:
# Box
boxes = res.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)
# Error < 0
if x1 < 0: x1 = 0
if y1 < 0: y1 = 0
if x2 < 0: x2 = 0
if y2 < 0: y2 = 0
# Class
cls = int(box.cls[0])
# Confidence
conf = math.ceil(box.conf[0])
print(f"Clase: {cls} Confidence: {conf}")
if conf > 0:
# Draw
cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 0, 255), 2)
cv2.putText(frame, f'{clsName[cls]} {int(conf * 100)}%', (x1, y1 - 20),
cv2.FONT_HERSHEY_COMPLEX, 1, (0, 0, 255), 2)
# Show
cv2.imshow("Waste Detect", frame)
# Close
t = cv2.waitKey(5)
if t == 27:
break
cap.release()
cv2.destroyAllWindows() |