无人机识别

航拍俯拍，垃圾与无人船的距离

image-20250303020149155

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
61
62
63
64
65
66
67
68
69
70
import torch
import cv2

def process_video(model, specified_class='mouse'):
    specified_class_center = None
    cap = cv2.VideoCapture(0)
    diffs = []  # List to store diff_x and diff_y values

    while True:
        ret, frame = cap.read()
        with torch.no_grad():
            image = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            results = model(image)

        for detection in results.xyxy[0]:
            conf = detection[4].item()
            if conf > 0.5:
                class_id = int(detection[5].item())
                class_name = model.names[class_id]
                bbox = detection[:4].cpu().numpy().astype(int)

                if class_name == specified_class:
                    specified_class_center = ((bbox[0] + bbox[2]) // 2, (bbox[1] + bbox[3]) // 2)

                center_x = (bbox[0] + bbox[2]) // 2
                center_y = (bbox[1] + bbox[3]) // 2

                cv2.circle(frame, (center_x, center_y), 5, (255, 0, 0), -1)
                cv2.rectangle(frame, (bbox[0], bbox[1]), (bbox[2], bbox[3]), (0, 255, 0), 4)

        if specified_class_center is not None:
            for detection in results.xyxy[0]:
                conf = detection[4].item()
                if conf > 0.5:
                    class_id = int(detection[5].item())
                    class_name = model.names[class_id]
                    bbox = detection[:4].cpu().numpy().astype(int)

                    if class_name != specified_class:
                        other_center_x = (bbox[0] + bbox[2]) // 2
                        other_center_y = (bbox[1] + bbox[3]) // 2

                        cv2.line(frame, (other_center_x, other_center_y), specified_class_center, (0, 0, 255), 2)
                        diff_x = center_x - specified_class_center[0]
                        diff_y = center_y - specified_class_center[1]
                        diffs.append((diff_x, diff_y))  # Append diff_x and diff_y values to the list

        cv2.imshow("Processed Image", frame)

        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

    cap.release()
    cv2.destroyAllWindows()

    return diffs

# 加载YOLOv5模型
model_path = 'D:/YOLOLAST/yolov5-master/runs/train/exp/weights/best.pt'
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = torch.hub.load('D:/YOLOLAST/yolov5-master', 'custom', model_path, source='local')
model.to(device)
model.eval()

# 在主函数中调用封装好的处理函数，并获取返回值
diffs = process_video(model, specified_class='mouse')

# 打印每一次的 diff_x 和 diff_y 的值
for i, (diff_x, diff_y) in enumerate(diffs):
    print(f"Iteration {i+1}: Δx = {diff_x}, Δy = {diff_y}")