56 lines
1.6 KiB
Python
56 lines
1.6 KiB
Python
import cv2
|
|
from utils.get_image import get_image
|
|
from ultralytics import YOLO
|
|
|
|
model = YOLO(r"best0.pt").to('cuda')
|
|
|
|
def yolo_shibie(im_PIL, detections):
|
|
results = model(im_PIL)
|
|
result = results[0]
|
|
|
|
# ✅ 获取绘制好框的图像
|
|
frame_with_boxes = result.plot()
|
|
|
|
# ✅ 用 OpenCV 动态显示
|
|
cv2.imshow("YOLO实时检测", frame_with_boxes)
|
|
|
|
# ESC 或 Q 键退出
|
|
if cv2.waitKey(1) & 0xFF in [27, ord('q')]:
|
|
return None
|
|
|
|
# ✅ 提取检测信息
|
|
for i in range(len(result.boxes.xyxy)):
|
|
left, top, right, bottom = result.boxes.xyxy[i]
|
|
cls_id = int(result.boxes.cls[i])
|
|
label = result.names[cls_id]
|
|
|
|
if label in ['center', 'next', 'npc1', 'npc2', 'npc3', 'npc4', 'boss', 'zhaozi']:
|
|
player_x = int(left + (right - left) / 2) + 3
|
|
player_y = int(top + (bottom - top) / 2) + 40
|
|
detections[label] = [player_x, player_y]
|
|
elif label in ['daojv', 'gw']:
|
|
player_x = int(left + (right - left) / 2) + 3
|
|
player_y = int(top + (bottom - top) / 2) + 40
|
|
detections[label].append([player_x, player_y])
|
|
|
|
return detections
|
|
|
|
|
|
while True:
|
|
detections = {
|
|
'center': None, 'next': None,
|
|
'npc1': None, 'npc2': None, 'npc3': None, 'npc4': None,
|
|
'boss': None, 'zhaozi': None,
|
|
'daojv': [], 'gw': []
|
|
}
|
|
|
|
im_opencv = get_image.get_frame() # [RGB, PIL]
|
|
detections = yolo_shibie(im_opencv[1], detections)
|
|
|
|
if detections is None: # 用户退出
|
|
break
|
|
|
|
print(detections)
|
|
|
|
cv2.destroyAllWindows()
|