钓鱼小代码

自己的代码记录，基于yolov9能钓鱼的小程序

# 此程序开启摄像头进行目标检测,并且可以获取label名和坐标太棒了
# 导包
import queue
import threading
import time
import pydirectinput as pydirectinput
import pyautogui as pyautogui
import torch
import cv2

# 下面两个是yolov5文件夹里面的代码
import utils.general as general
from models.experimental import attempt_load

# 创建一个队列，1代表maxSize
q = queue.Queue(1)  # 创建一个先进先出的队列


# 定义生产者线程 Producer
class Producer(threading.Thread):

    def run(self):
        global q
        pyautogui.FAILSAFE = False
        # 加载YOLOv9模型
        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

        # 检查是否可用CUDA加速。如果CUDA可用，则将设备设置为'cuda'，否则设置为'cpu'。
        model = attempt_load('D:\\yolov9-main\\best.pt', device)

        # half = device.type != 'cpu'
        # if half:
        #     model.half()  # 将模型的权重从浮点数格式转换为半精度浮点数格式

        # 设置置信度阈值和IoU阈值
        model.conf = 0.3
        model.iou = 0.5
        # 设置模型为推理模式
        model.eval()
        # 打开视频流
        # cap = cv2.VideoCapture(0)
        # cap.set(3, 640)  # 设置分辨率
        # cap.set(4, 360)
        while True:

            target_directory = r"D:\yolov9-main\data\fishData\images\gameshot"
            # 字典1
            # dic1 = {'herox': 0.0, 'heroy': 0.0, 'gox': 0.0, 'goy': 0.0}
            # 计数器，用于命名截图文件
            screenshot = pyautogui.screenshot(region=[486, 283, 1000, 640])  # x,y,w,h
            screenshot.save(target_directory + "\\" + "pic{}.png".format(0))
            # 暂时阻塞进程，文件保存完毕继续执行
            time.sleep(0.03)

            # 图像预处理
            img = cv2.imread(target_directory + "\\" + "pic{}.png".format(0))
            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
            img = torch.from_numpy(img).to(device).float() / 255.0
            img = img.permute(2, 0, 1).unsqueeze(0)
            img = torch.nn.functional.interpolate(img, size=(640, 640), mode='bilinear', align_corners=False)

            # size = (320, 320)：表示目标输出图像的大小为320x320像素。将输入图像按照此大小进行调整。
            # mode = 'bilinear'：表示插值方法采用双线性插值。双线性插值是一种常用的插值方法，它通过对最近的四个像素进行加权平均来计算新像素值。
            # align_corners = False：表示在计算插值时不对齐角点。这个参数的具体影响取决于插值方法和具体的实现方式。
            # 进行物体检测
            outputs = model(img)
            results = general.non_max_suppression(outputs, conf_thres=model.conf, iou_thres=model.iou)
            # tensor_cuda = torch.tensor(img.shape)[[1, 0, 1, 0]]
            # 将CUDA张量移动到CPU上
            # tensor_cpu = tensor_cuda.cpu()
            #
            # # 将CPU张量转换为NumPy数组
            # imgArray = tensor_cpu.numpy()

            # 通过非最大抑制算法对模型输出的边界框进行处理，以选择最佳的边界框并过滤掉冗余的检测结果。
            # outputs：模型的输出结果，通常是一个包含预测边界框及其置信度的数组。
            # conf_thres：表示置信度阈值，用于过滤掉置信度低于此阈值的边界框。
            # iou_thres：表示IoU（Intersection over Union）阈值，用于合并重叠度高于此阈值的边界框。
            # 算法的步骤如下：
            # 1、根据置信度阈值过滤掉置信度低于阈值的边界框。
            # 2、对剩余的边界框按照置信度进行排序，置信度高的排在前面。
            # 3、从排好序的边界框列表中选择置信度最高的边界框，并将其添加到最终的结果列表中。
            # 4、遍历剩余的边界框列表，计算当前边界框与已选择的边界框的IoU值（重叠度），如果大于IoU阈值，则将其从列表中移除。
            # 5、重复步骤3和4，直到所有边界框都被处理完毕。
            # 6、最终，non_max_suppression()函数返回经过非最大抑制处理后的边界框结果。
            # if results[0] is not None and len(results[0]) > 0:
            #     result = results[0]
            # 绘制边界框和标签
            for result in results:  # 逐个访问并获取
                if result is not None and len(result) > 0:
                    for idx, (*xyxy, conf, cls) in enumerate(reversed(result)):
                        # result是一个数组或张量，每一行代表一个边界框，包含边界框的坐标、类别标签和置信度等信息。
                        # *xyxy表示将边界框的前四个元素解包为一个名为xyxy的列表。这里的xyxy表示边界框的坐标信息，通常是左上角和右下角的坐标值。
                        # conf表示边界框的置信度，通常是模型对该边界框所属类别的预测置信度。
                        # cls表示边界框的类别标签，通常是模型对该边界框所属类别的预测结果。
                        # label = f'{model.names[int(cls)]}'
                        labelRat = f'{conf:.2f}'
                        labelRatFloat = float(labelRat)
                        # print(label + '01:' + labelRat)
                        # 标签准确率为50%再行动
                        if labelRatFloat > 0.68:
                            # print(f"收杆时间:{current_time}")
                            pydirectinput.keyDown('x')
                            time.sleep(0.5)
                            pydirectinput.keyUp('x')
                            msg = '++下杆吧已经收杆了++'
                            q.put(msg)  # 默认阻塞
                            print(msg)


# sleep就是让资源,因此收鱼竿尽量要多的资源，而放鱼竿资源尽量要少，即放鱼竿的sleep尽量长点
class Consumer(threading.Thread):
    def run(self):
        global q
        while True:
            msg = q.get()  # 默认阻塞
            parts = msg.split(':')
            # 收杆提醒
            fishWords = parts[0]
            fishWords.strip()

            if '下杆吧' in fishWords:
                pydirectinput.keyDown('z')
                time.sleep(0.1)
                pydirectinput.keyUp('z')
                q.put('已经下杆')
                print('！！！已经下杆：{}'.format(msg))
            else:
                pydirectinput.keyDown('z')
                time.sleep(0.1)
                pydirectinput.keyUp('z')


if __name__ == '__main__':
    time.sleep(3)
    t1 = Producer()
    t2 = Consumer()
    t1.start()
    t2.start()