Это мой код прямо сейчас:
Код: Выделить всё
import cv2
import time
from ultralytics import YOLO
from ultralytics.utils.plotting import colors
import threading
class YOLOVisualizer:
"""
A utility class for running YOLO detection on video/webcam and visualizing results with
custom bounding boxes and overlays.
"""
def __init__(self, source=0, model="best_openvino_model/", device="cpu"):
self.source = source
self.model_path = model
self.device = device
# Load model
self.model = YOLO(model, task="detect")
self.names = self.model.names
# Video capture
self.cap = cv2.VideoCapture(source)
self.cap.set(cv2.CAP_PROP_BUFFERSIZE, 1)
# Get FPS from source
w, h, fps = (
int(self.cap.get(x))
for x in (
cv2.CAP_PROP_FRAME_WIDTH,
cv2.CAP_PROP_FRAME_HEIGHT,
cv2.CAP_PROP_FPS,
)
)
# # Video writer
self.video_writer = cv2.VideoWriter(
"yolo-output.avi",
cv2.VideoWriter_fourcc(*"mp4v"),
fps,
(640, 480),
)
self.avg_fps = []
@staticmethod
def draw_text_with_bg(
img,
text,
pos,
font_scale=0.6,
bg_color=(0, 0, 0),
fg_color=(104, 31, 17),
padding=15,
thickness=4,
):
"""Draw text with a solid background rectangle."""
font = cv2.FONT_HERSHEY_SIMPLEX
(text_width, text_height), baseline = cv2.getTextSize(
text, font, font_scale, thickness
)
rect_width = text_width + 2 * padding
rect_height = text_height + baseline + 2 * padding
rect_x1, rect_y1 = pos
rect_x2, rect_y2 = rect_x1 + rect_width, rect_y1 + rect_height
cv2.rectangle(img, (rect_x1, rect_y1), (rect_x2, rect_y2), bg_color, -1)
text_x = rect_x1 + (rect_width - text_width) // 2
text_y = rect_y1 + (rect_height + text_height) // 2
cv2.putText(
img,
text,
(text_x, text_y),
font,
font_scale,
fg_color,
thickness,
cv2.LINE_AA,
)
@staticmethod
def draw_bbox(img, box, label="", conf=None, cls_id=0):
"""Draw a styled bounding box with label."""
x1, y1, x2, y2 = map(int, box)
# Pick consistent color
color = colors(11 if cls_id == 2 else cls_id, True)
thickness = 3
cv2.rectangle(img, (x1, y1), (x2, y2), color, thickness)
if label:
text_scale = 1.4
thickness_text = 2
(tw, th), baseline = cv2.getTextSize(
label, cv2.FONT_HERSHEY_SIMPLEX, text_scale, thickness_text
)
pad = 8
lx1, ly1 = x1, y1 - (th + baseline + 2 * pad)
lx2, ly2 = x1 + tw + 2 * pad, y1
# Keep label inside frame
if ly1 < 0:
ly1 = y1
ly2 = y1 + th + baseline + 2 * pad
cv2.rectangle(img, (lx1, ly1), (lx2, ly2), color, -1)
text_x = lx1 + (lx2 - lx1 - tw) // 2
text_y = ly1 + (ly2 - ly1 + th) // 2
cv2.putText(
img,
label,
(text_x, text_y),
cv2.FONT_HERSHEY_SIMPLEX,
text_scale,
(255, 255, 255),
thickness_text,
cv2.LINE_AA,
)
# Main inference loop
def run(self):
"""Run YOLO inference and visualization loop."""
while self.cap.isOpened():
self.cap.grab()
ret, frame = self.cap.read()
if not ret:
break
frame = cv2.resize(frame, (640, 480))
start_time = time.time()
results = self.model.predict(
frame,
device="cpu",
conf=0.25,
#imgsz=640,
half=False, # important for CPU
)[0]
process_time = time.time() - start_time
boxes = results.boxes.xyxy.tolist()
clss = results.boxes.cls.tolist()
confs = results.boxes.conf.cpu().numpy()
for box, cls, conf in zip(boxes, clss, confs):
cls = int(cls)
self.draw_bbox(
frame,
box,
label=self.names[cls],
conf=conf,
cls_id=cls,
)
# FPS calculation
self.avg_fps.append(1.0 / process_time if process_time > 0 else 0)
if len(self.avg_fps) > 30:
self.avg_fps.pop(0)
fps = sum(self.avg_fps) / len(self.avg_fps)
# Overlays
self.draw_text_with_bg(
frame,
f"FPS: {int(fps)}",
(5, 25),
bg_color=(104, 31, 17),
fg_color=(255, 255, 255),
font_scale=1,
)
self.draw_text_with_bg(
frame,
f"Time: {process_time * 1000:.0f}ms",
(360, 25),
bg_color=(104, 0, 123),
fg_color=(255, 255, 255),
font_scale=1,
)
cv2.imshow(f"YOLO - {self.source}", frame)
#self.video_writer.write(frame)
if cv2.waitKey(1) & 0xFF == ord("q"):
break
time.sleep(0.05)
self.cap.release()
self.video_writer.release()
cv2.destroyAllWindows()
# Run Example
if __name__ == "__main__":
#YOLO11 test
# Export command: yolo export format=onnx model=yolo11n.onnx
# visualizer = YOLOVisualizer(source=r"horse-rider.mp4", model="yolo11n.onnx")
# visualizer.run()
#YOLO26 test
# Export command: yolo export format=onnx model=yolo26n.onnx
#shared_model = YOLO("best_openvino_model", task="detect")
vis1 = YOLOVisualizer(
source="http://192.168.100.8:4747/video",
model="best_openvino_model"
)
vis2 = YOLOVisualizer(
source="http://192.168.100.13:4747/video",
model="best_openvino_model"
)
t1 = threading.Thread(target=vis1.run)
t2 = threading.Thread(target=vis2.run)
t1.start()
t2.start()
t1.join()
t2.join()
