import model_elem

import cv2
import torch
import torchvision
import numpy as np
import torchvision.transforms as transforms
from torchvision.models.detection import KeypointRCNN_ResNet50_FPN_Weights
# 加载pytorch提供的keypointrcnn_resnet50_fpn()网络模型，可以对17个人体关键点进行检测。
#如果有可用的GPU，则使用GPU，否则使用CPU。
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
#pytorch模型的python接口
model = torchvision.models.detection.keypointrcnn_resnet50_fpn(weights=KeypointRCNN_ResNet50_FPN_Weights.DEFAULT)
model.to(device)
model.eval()

def Capture_Point(image, confidence=0.9):
    image = cv2.cvtColor(image, cv2.COLOR_RGBA2RGB)
    # 准备需要检测的图像
    transform_d = transforms.Compose([transforms.ToTensor()])
    image_t = transform_d(image)    ## 对图像进行变换
    pred = model([image_t.to(device)])         ## 将模型作用到图像上
    # 检测出目标的类别和得分
    pred_class = [model_elem.OBJECT_LIST[ii] for ii in list(pred[0]['labels'].cpu().numpy())]
    pred_score = list(pred[0]['scores'].detach().cpu().numpy())
    # 检测出目标的边界框
    pred_boxes = [[ii[0], ii[1], ii[2], ii[3]] for ii in list(pred[0]['boxes'].detach().cpu().numpy())]
    ## 只保留识别的概率大约 confidence 的结果。
    pred_index = [pred_score.index(x) for x in pred_score if x > confidence]
    for index in pred_index:
        box = pred_boxes[index]
        box = [int(i) for i in box]
        cv2.rectangle(image,(int(box[0]),int(box[1])),(int(box[2]),int(box[3])),(0,255,255))
        texts = pred_class[index] + ":" + str(np.round(pred_score[index], 2))
        cv2.putText(image, texts,(box[0], box[1]), cv2.FONT_HERSHEY_SIMPLEX, 0.75, (0, 255, 255), 2)

    pred_keypoint = pred[0]["keypoints"]
    # 检测到实例的关键点
    pred_keypoint = pred_keypoint[pred_index].detach().cpu().numpy()
    # 对实例数量索引
    my_result = {}
    for index in range(pred_keypoint.shape[0]):
        # 对每个实例的关键点索引
        keypoints = pred_keypoint[index]
        for ii in range(keypoints.shape[0]): ##ii为第几个坐标点
            x = int(keypoints[ii, 0]) #x坐标
            y = int(keypoints[ii, 1]) #y坐标
            visi = keypoints[ii, 2] #置信度
            if visi > 0.:
                cv2.circle(image, (int(x),int(y)), 1, (0,0,255),4)
                texts = str(ii+1)
                cv2.putText(image,texts, (int(x), int(y)), cv2.FONT_HERSHEY_SIMPLEX, 0.75, (0, 255, 255), 2)
                my_result[texts] = (int(x), int(y))
    return image,my_result