'''
author ：  Hatcher W
Email： hcwang@semi.ac.cn
use: 目标检测Y5推理
input:需要推理的文件目录
output：推理结果
'''
import time
import cv2
import torch,os
import numpy as np
from numpy import random
import os,sys
root_path = os.getcwd()
sys.path.insert(0,root_path+"/Y5")
from Y5.models.experimental import attempt_load
from Y5.utils.datasets import  letterbox
from Y5.utils.general import   non_max_suppression, \
    scale_coords
from Y5.utils.plots import plot_one_box



class Detector:

    def __init__(self, device):
        super(Detector, self).__init__()
        self.img_size = 640
        self.conf_thres = 0.3
        self.iou_thres = 0.4
        self.stride = 1

        self.weights = '../data_MPKFNet/weigth/detection/Y/best.pt'
        self.device = device
        #self.device = select_device(self.device)
        model = attempt_load(self.weights, map_location=self.device)
        model.to(self.device).eval()
        model.half()
        self.model = model
        self.names = model.module.names if hasattr(
            model, 'module') else model.names

    def preprocess(self, img):

        img0 = img.copy()
        img = letterbox(img, new_shape=self.img_size, stride=self.stride)[0]
        img = img[:, :, ::-1].transpose(2, 0, 1)
        img = np.ascontiguousarray(img)
        img = torch.from_numpy(img).to(self.device)
        img = img.half()  # °ë¾«¶È
        img /= 255.0  # ͼÏñ¹éÒ»»¯
        if img.ndimension() == 3:
            img = img.unsqueeze(0)

        return img0, img

    def detect(self, im):

        im0, img = self.preprocess(im)

        pred = self.model(img, augment=False)[0]
        pred = pred.float()
        pred = non_max_suppression(pred, self.conf_thres, self.iou_thres)

        pred_boxes = []
        scores = []
        labels = []
        for det in pred:

            if det is not None and len(det):
                det[:, :4] = scale_coords(
                    img.shape[2:], det[:, :4], im0.shape).round()

                pred_boxes.append(det[:,:4])
                scores.append(det[:,4])
                labels.append((det[:,-1]))
                pred_boxes = pred_boxes[0]
                scores = scores[0]
                labels = labels[0]
                # for *x, conf, cls_id in det:
                #     lbl = self.names[int(cls_id)]
                #     if not lbl in ['person', 'car', 'truck']:
                #         continue
                #     x1, y1 = int(x[0]), int(x[1])
                #     x2, y2 = int(x[2]), int(x[3])
                #     pred_boxes.append(
                #         (x1, y1, x2, y2, lbl, conf))

        return im, pred_boxes,scores,labels


if __name__ == '__main__':

    os.environ['CUDA_VISIBLE_DEVICES'] = '4'
    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
    os.environ['CUDA_VISIBLE_DEVICES'] = '4'
    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
    #path = '/data1/Hatcher/model/data_MPKFNet/data/VOC/VOCdevkit/VOC5/JPEGImages/'
    path = "/data/Hatcher/ouput1/0001_20210414151147/RawImage/"
    outpath = "/data/Hatcher/detectout3"
    if not os.path.exists(outpath):
        os.makedirs(outpath)
    detector = Detector(device)
    path_lists = os.listdir(path)
    for image_path in path_lists:
        img = cv2.imread(os.path.join(path,image_path))  # BGR

        # Get names and colors
        names = [ 'background','aeroplane', 'bicycle', 'bird', 'boat', 'bottle', 'bus', 'car', 'cat', 'chair', 'cow', 'diningtable', 'dog',
         'horse', 'motorbike', 'person', 'pottedplant', 'sheep', 'sofa', 'train', 'tvmonitor' ]#detector.model.module.names if hasattr(detector.model, 'module') else detector.model.names
        colors = [[random.randint(0, 255) for _ in range(3)] for _ in names]
        t3 = time.time()
        im0, pred_boxes ,scores,labels = detector.detect(img)
        t4 = time.time()
        print("推理： ", (t4 - t3) * 1000)
        print(type(labels))
        if len(labels) == 0:
            continue
        indices_label_left = labels < 4
        indices_label_right = labels > 2
        indices = []
        for i in range(len(indices_label_left)):
            indices.append(indices_label_left[i] and indices_label_right[i])

        pred_boxes = pred_boxes[indices]
        labels = labels[indices]
        scores = scores[indices]    #因为莫名奇妙多了一个list维度，所以加了这个0
        pred_boxes = pred_boxes[0]
        labels = labels[0]
        scores = scores[0]
        print(len(scores))
        for i in range(len(scores)):
            xyxy = pred_boxes[i]
            conf = scores[i]
            cls = labels[i]
            print(cls)
            label = f'{names[int(cls)]} {conf:.2f}'
            plot_one_box(xyxy, im0, label=label, color=colors[int(cls)], line_thickness=1)
        #drawn_image = draw_boxes(im0, pred_boxes, labels, scores, names).astype(np.uint8)
        # for *xyxy, conf, cls in pred_boxes,scores,labels:
        #     label = f'{names[int(cls)]} {conf:.2f}'
        #     plot_one_box(xyxy, im0, label=label, color=colors[int(cls)], line_thickness=3)

        # Stream results
        #cv2.imshow(str(path), im0)
        t1 = time.time()
        cv2.imwrite(os.path.join(outpath, image_path),im0)
        t2 = time.time()
        print("写入： ",(t2-t1)*1000)
        #cv2.waitKey(-1)