import torch

# 检查CUDA是否可用
print(torch.cuda.is_available())

# 如果CUDA可用，列出CUDA设备
if torch.cuda.is_available():
    print("CUDA is available!")
    print("CUDA Device Name:", torch.cuda.get_device_name(0))
else:
    print("CUDA is not available.")

pip install -r requirements.txt
pip install ultralytics
# ultralytics中包含了yolov8，不需要额外pip install yolo，这里的yolo下载下来居然是一个管理amp的包，会造成yolo命令冲突而失效

from ultralytics import YOLO

def main():
   # Load model
   model = YOLO("yolov8n.pt")

   # Train
   model.train(data="datasets/pvz/pvztrain.yaml", epochs=250,patience=150)

   # Validate
   model.val()

if __name__ == "__main__":
   main()
#这里设置轮数为250轮以提高精度，但其实100轮后提升效果就已经区别不大，设置patience=150，即在150轮后检测如果已经无提升则直接结束训练

from ultralytics import YOLO
import cv2
import numpy as np

# 加载YOLOv8模型
model = YOLO('runs/detect/train/weights/best.pt')

# 打开视频文件
cap = cv2.VideoCapture('test/test2.mp4')

# 循环遍历视频帧
while cap.isOpened():
    # 从视频读取一帧
    success, frame = cap.read()
    if not success:
        break

    # 在帧上运行YOLOv8检测
    results = model.predict(frame)

    # 检查是否有检测结果
    if results:
        # 获取框和类别信息
        boxes = results[0].boxes.xyxy.cpu().numpy()  # 修改为获取xyxy格式的边界框，并转换为numpy数组
        classes = results[0].boxes.cls.cpu().numpy()  # 获取类别索引，并转换为numpy数组

        # 在帧上展示结果
        annotated_frame = results[0].plot()  # 绘制检测结果

        # 展示带注释的帧
        annotated_frame = cv2.resize(annotated_frame, (640, 480))
        cv2.imshow('YOLOv8 Detection', annotated_frame)
    else:
        # 如果没有检测结果，直接展示原始帧
        cv2.imshow('YOLOv8 Detection', frame)

    # 如果按下'q'则退出循环
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

# 释放视频捕获对象并关闭显示窗口
cap.release()
cv2.destroyAllWindows()

import argparse
import os
import platform
import sys
from pathlib import Path
import cv2
import torch

from ultralytics import YOLO

import pyautogui
from PIL import ImageGrab
import win32gui, win32con, win32com.client
import numpy as np
import time


def cilck_init():
    hwnd = win32gui.FindWindow(None, '植物大战僵尸中文版')
    print(hwnd)
    shell = win32com.client.Dispatch("WScript.Shell")
    shell.SendKeys('%')
    win32gui.SetForegroundWindow(hwnd)
    window_x, window_y, right, bottom = win32gui.GetWindowRect(hwnd)
    box = (window_x, window_y, right, bottom)
    print(box)
    return box


def run(weights='runs/detect/train/weights/best.pt', source='self_data/pvz', imgsz=640, conf_thres=0.25,
        iou_thres=0.45):
    # Load model
    model = YOLO(weights)

    # Initialize click function
    box = cilck_init()

    num_pic = 1
    while num_pic:
        # Grab screenshot
        background_bgr = np.array(ImageGrab.grab(box))
        background = background_bgr[:, :, [2, 1, 0]]  # Convert BGR to RGB
        img_path = 'datasets/data/pvz/test.jpg'
        cv2.imwrite(img_path, background)

        # Perform inference
        results = model.predict(img_path, imgsz=imgsz, conf=conf_thres, iou=iou_thres)

        # Process results
        for result in results:
            for det in result.boxes.data:
                xyxy = det[:4].cpu().numpy().astype(int)
                conf = det[4].cpu().numpy()
                cls = int(det[5].cpu().numpy())
                print(xyxy, conf, cls)

                # Use pyautogui to click on detected coordinates
                pyautogui.click(box[0] + xyxy[0], box[1] + xyxy[1] + 20)

                # Additional processing or saving results can be added here

        num_pic += 1
        #time.sleep(1)  # Add a delay to avoid excessive clicking


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--weights', type=str, default='runs/detect/train/weights/best.pt', help='model path')
    parser.add_argument('--source', type=str, default='self_data/pvz', help='source')
    parser.add_argument('--imgsz', type=int, default=640, help='image size')
    parser.add_argument('--conf-thres', type=float, default=0.25, help='confidence threshold')
    parser.add_argument('--iou-thres', type=float, default=0.45, help='IOU threshold')
    opt = parser.parse_args()

    run(opt.weights, opt.source, opt.imgsz, opt.conf_thres, opt.iou_thres)


if __name__ == "__main__":
    main()