45 lines
1.2 KiB
Python
45 lines
1.2 KiB
Python
import cv2
|
|
from ultralytics import YOLO
|
|
|
|
model = YOLO('yolov5s.pt') # Replace with your trained logo detection model path
|
|
cap = cv2.VideoCapture(0) # Webcam input
|
|
|
|
if not cap.isOpened():
|
|
print("Error: Could not open video stream.")
|
|
exit()
|
|
|
|
while True:
|
|
ret, frame = cap.read()
|
|
if not ret:
|
|
break
|
|
|
|
results = model(frame)
|
|
|
|
for result in results:
|
|
boxes = result.boxes.xyxy.cpu().numpy()
|
|
scores = result.boxes.conf.cpu().numpy()
|
|
class_ids = result.boxes.cls.cpu().numpy()
|
|
|
|
for box, score, cls_id in zip(boxes, scores, class_ids):
|
|
if score < 0.3:
|
|
continue
|
|
|
|
x1, y1, x2, y2 = map(int, box)
|
|
label = model.names[int(cls_id)]
|
|
|
|
# Draw bounding box and label
|
|
cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
|
|
cv2.putText(frame, f'{label} {score:.2f}', (x1, y1 - 10),
|
|
cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
|
|
|
|
# Log detection to console
|
|
print(f"Detected: {label} | Confidence: {score:.2f} | Box: ({x1}, {y1}), ({x2}, {y2})")
|
|
|
|
cv2.imshow('Logo Detection', frame)
|
|
|
|
if cv2.waitKey(1) & 0xFF == 27:
|
|
break
|
|
|
|
cap.release()
|
|
cv2.destroyAllWindows()
|