import os

import numpy as np
import cv2
from fastapi import FastAPI, File, UploadFile, Form

from ultralytics import YOLO
import io
from PIL import Image, ImageGrab
import uvicorn

current_dir = os.path.dirname(os.path.abspath(__file__))
onnx_model = YOLO(os.path.join(current_dir,'ai','best.pt'))

# 初始化 FastAPI 应用
app = FastAPI()

@app.post('/detect')
async def detect_objects(file: UploadFile = File(...), conf: float = Form(0.6)):

    try:
        # 将图片数据转换为PIL图像对象
        img = Image.open(io.BytesIO(await file.read()))

        # 使用YOLO模型进行检测
        results = onnx_model(img)

        # 获取第一个图片的结果
        result = results[0]

        # 获取所有的边界框
        boxes = result.boxes
        # 所有的分类名字
        categorys = result.names
        # print(categorys)

        # 提取每个检测框的坐标
        coordinates_list = []

        for box in boxes:
            # 获取边界框的坐标 (x1, y1, x2, y2)
            coordinates = box.xyxy[0].tolist()
            confidence = box.conf.item()
            # 获取类别ID
            class_id = int(box.cls.item())
            # print(class_id, confidence)

            if confidence >= conf:
                # 获取类别名称
                class_name = categorys[class_id]
                # 存储检测结果
                # 四舍五入到整数
                rounded_coordinates = [round(coord) for coord in coordinates]
                coordinates_list.append({
                    'coordinates': rounded_coordinates,
                    'class_name': class_name,
                    'confidence': confidence
                })

        return {'code':200,'data':{'list':coordinates_list}}

    except Exception as e:
        return {'code':500,'data':{}}

@app.get('/getImgMatch')
def getImgMatch(imageName: str, x1: int, y1: int, x2: int, y2: int, conf: float = 0.6):

    screenshot = np.array(ImageGrab.grab(bbox=(x1, y1, x2, y2)))[:, :, ::-1]
    # cv2.imwrite(f'screenshot.jpg', screenshot)
    # 2. 将图片数据转换为 PIL.Image 对象
    img1 = cv2.imread(f"./img/{imageName}")
    img2 = np.array(screenshot)
    # img2 = cv2.cvtColor(np.array(i2), cv2.COLOR_RGB2BGR)
    # 灰度化
    g1 = cv2.cvtColor(img1, cv2.COLOR_BGR2GRAY)
    g2 = cv2.cvtColor(img2, cv2.COLOR_BGR2GRAY)
    result = cv2.matchTemplate(g2, g1, cv2.TM_CCOEFF_NORMED)
    # 设定阈值
    if conf == '' or conf is None:
        conf = 0.6
    # 找到匹配区域
    loc = np.where(result >= float(conf))
    if loc[0].size == 0:
        return {'code':400,'data':{'x': -1, 'y': -1}}
    else:
        # 取第一个匹配点的位置
        pt = zip(*loc[::-1]).__next__()
        center_x = pt[0] + img1.shape[1] // 2
        center_y = pt[1] + img1.shape[0] // 2
        # 在截图上绘制方框
        # cv2.rectangle(img2, pt, (pt[0] + img1.shape[1], pt[1] + img1.shape[0]), (0, 0, 255), 2)
        #
        # # 保存带有方框的截图
        # cv2.imwrite('screenshot_marked.jpg', img2)
        # print(f"第一个匹配区域的中心点坐标: ({center_x}, {center_y})")
        return {'code':200,'data':{'x': round(center_x), 'y': round(center_y)}}

@app.post('/getMatch')
async def getMatch(imgFile: UploadFile = File(...),imgFileBig: UploadFile = File(...), conf: float = Form(0.6)):
    # 读取上传的小图文件内容
    file_bytes = await imgFile.read()
    # 使用 OpenCV 将二进制数据转换为图像
    img1 = cv2.imdecode(np.frombuffer(file_bytes, np.uint8), cv2.IMREAD_COLOR)

    # 读取上传的大图文件内容
    file_big_bytes = await imgFileBig.read()
    # 同样将大图转换为图像
    img2 = cv2.imdecode(np.frombuffer(file_big_bytes, np.uint8), cv2.IMREAD_COLOR)

    # img2 = cv2.cvtColor(np.array(i2), cv2.COLOR_RGB2BGR)
    # 灰度化
    g1 = cv2.cvtColor(img1, cv2.COLOR_BGR2GRAY)
    g2 = cv2.cvtColor(img2, cv2.COLOR_BGR2GRAY)
    result = cv2.matchTemplate(g2, g1, cv2.TM_CCOEFF_NORMED)
    # 设定阈值
    if conf == '' or conf is None:
        conf = 0.6
    # 找到匹配区域
    loc = np.where(result >= float(conf))
    if loc[0].size == 0:
        return {'code':400,'data':{'x': -1, 'y': -1}}
    else:
        # 取第一个匹配点的位置
        pt = zip(*loc[::-1]).__next__()
        center_x = pt[0] + img1.shape[1] // 2
        center_y = pt[1] + img1.shape[0] // 2
        # 在截图上绘制方框
        # cv2.rectangle(img2, pt, (pt[0] + img1.shape[1], pt[1] + img1.shape[0]), (0, 0, 255), 2)
        #
        # # 保存带有方框的截图
        # cv2.imwrite('screenshot_marked.jpg', img2)
        # print(f"第一个匹配区域的中心点坐标: ({center_x}, {center_y})")
        return {'code':200,'data':{'x': round(center_x), 'y': round(center_y)}}

if __name__ == '__main__':
    print("启动成功...^_^...监听端口：9235")
    uvicorn.run(app, host="127.0.0.1", port=9235, log_level="critical", workers=1)