Programmiererforum

Code: Select all

import numpy as np

import vot
import sys
import cv2

from ultralytics import YOLO

def iou(box_a, box_b):
# Compute intersection-over-union for two boxes [x1, y1, x2, y2]
x_a = max(box_a[0], box_b[0])
y_a = max(box_a[1], box_b[1])
x_b = min(box_a[2], box_b[2])
y_b = min(box_a[3], box_b[3])
interArea = max(0, x_b - x_a) * max(0, y_b - y_a)
boxAArea = (box_a[2] - box_a[0]) * (box_a[3] - box_a[1])
boxBArea = (box_b[2] - box_b[0]) * (box_b[3] - box_b[1])
IoU = interArea / float(boxAArea + boxBArea - interArea)
return IoU

def get_best_match(selection_box, boxes):
# selection_box: [x1, y1, x2, y2]
# boxes: iterable of box objects with .xyxy attribute
max_iou = 0
best_id = None
best_box = None
for box in boxes:
box_xyxy = box.xyxy.cpu().numpy().tolist()[0]
iou_score = iou(selection_box, box_xyxy)
if iou_score > max_iou:
max_iou = iou_score
best_id = int(box.id.item()) if box.id is not None else None
best_box = box
if max_iou == 0 and len(boxes) > 0:
# Fallback: pick closest box by center
sel_cx = (selection_box[0] + selection_box[2]) / 2
sel_cy = (selection_box[1] + selection_box[3]) / 2
min_dist = float('inf')
for box in boxes:
box_xyxy = box.xyxy.cpu().numpy().tolist()[0]
bx = (box_xyxy[0] + box_xyxy[2]) / 2
by = (box_xyxy[1] + box_xyxy[3]) / 2
dist = (sel_cx - bx) ** 2 + (sel_cy - by) ** 2
if dist < min_dist:
min_dist = dist
best_id = int(box.id.item()) if box.id is not None else None
best_box = box
return best_id, best_box

detector = YOLO('/home/uuu/jupyter_projects/MyProjects/baka/yolo11n.pt')

handle = vot.VOT("rectangle")
selection = handle.region()
selection_box = [selection.x, selection.y, selection.x + selection.width, selection.y + selection.height]

max_iou = 0
target_id = None
score = 0
x1, y1, x2, y2 = 0, 0, 0, 0

results = detector.track(handle.frame(), tracker='botsort.yaml', stream=True, persist=True)

for result in results:
boxes = result.boxes
try:
if target_id is None:
for box in boxes:
target_id, _ = get_best_match(selection_box, boxes)
if target_id is None:
print("Target ID not found -_-")
handle.report(vot.Rectangle(0, 0, 0, 0), 0.0)

found = False
for box in boxes:
curr_id = int(box.id.item()) if box.id is not None else None
if curr_id == target_id:
box_xyxy = box.xyxy.cpu().numpy().tolist()[0]
x1, y1, x2, y2 = map(int, box_xyxy)
score = float(box.conf.item())
found = True
break

if found:
handle.report(vot.Rectangle(x1, y1, x2 - x1, y2 - y1), score)
else:
handle.report(vot.Rectangle(0, 0, 0, 0), 0.0)

except Exception as e:
# Always respond, even if something goes wrong
handle.report(vot.Rectangle(0, 0, 0, 0), 0.0)
print(f"Error: {e}", file=sys.stderr)