'''
注释：
    author is leilei
    由于最新版本的darknet将images和video预测分开了，并不是很符合自己的额需求，
    因此参考darknet_images.py修改成属于自己的预测函数
注：
    darknet.py 核心函数：load_network、detect_image draw_boxes bbox2points
    darknet_images.py 核心函数： image_detection,此函数需要修改成输入图像
    darknet官方写的预测图像输出依旧为正方形，而非原图！因此要转换成原图
'''
import os
import cv2
import numpy as np
import darknet

class Detect:
    def __init__(self, metaPath, configPath, weightPath, gpu_id=2, batch=1):
        '''
        :param metaPath:   ***.data 存储各种参数
        :param configPath: ***.cfg  网络结构文件
        :param weightPath: ***.weights yolo的权重
        :param batch:      ########此类只支持batch=1############
        '''
        assert batch == 1, "batch必须为1"
        # 设置gpu_id
        darknet.set_gpu(gpu_id)
        # 网络
        network, class_names, class_colors = darknet.load_network(
            configPath,
            metaPath,
            weightPath,
            batch_size=batch
        )
        self.network = network
        self.class_names = class_names
        self.class_colors = class_colors

    def bbox2point(self, bbox):
        x, y, w, h = bbox
        xmin = x - (w / 2)
        xmax = x + (w / 2)
        ymin = y - (h / 2)
        ymax = y + (h / 2)
        return (xmin, ymin, xmax, ymax)

    def point2bbox(self, point):
        x1, y1, x2, y2 = point
        x = (x1 + x2) / 2
        y = (y1 + y2) / 2
        w = (x2 - x1)
        h = (y2 - y1)
        return (x, y, w, h)

    def image_detection(self, image_bgr, network, class_names, class_colors, thresh=0.25):
        # 判断输入图像是否为3通道
        if len(image_bgr.shape) == 2:
            image_bgr = np.stack([image_bgr] * 3, axis=-1)
        # 获取原始图像大小
        orig_h, orig_w = image_bgr.shape[:2]

        width = darknet.network_width(network)
        height = darknet.network_height(network)
        darknet_image = darknet.make_image(width, height, 3)

        # image = cv2.imread(image_path)
        image_rgb = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB)
        image_resized = cv2.resize(image_rgb, (width, height), interpolation=cv2.INTER_LINEAR)

        darknet.copy_image_from_bytes(darknet_image, image_resized.tobytes())
        detections = darknet.detect_image(network, class_names, darknet_image, thresh=thresh)
        darknet.free_image(darknet_image)
        '''注意：这里原始代码依旧是608*608，而不是原图大小，因此我们需要转换'''
        new_detections = []
        for detection in detections:
            pred_label, pred_conf, (x, y, w, h) = detection
            new_x = x / width * orig_w
            new_y = y / height * orig_h
            new_w = w / width * orig_w
            new_h = h / height * orig_h

            # 可以约束一下
            (x1, y1, x2, y2) = self.bbox2point((new_x, new_y, new_w, new_h))
            x1 = x1 if x1 > 0 else 0
            x2 = x2 if x2 < orig_w else orig_w
            y1 = y1 if y1 > 0 else 0
            y2 = y2 if y2 < orig_h else orig_h

            (new_x, new_y, new_w, new_h) = self.point2bbox((x1, y1, x2, y2))

            new_detections.append((pred_label, pred_conf, (new_x, new_y, new_w, new_h)))

        image = darknet.draw_boxes(new_detections, image_rgb, class_colors)
        return cv2.cvtColor(image, cv2.COLOR_RGB2BGR), new_detections

    def predict_image(self, image_bgr, thresh=0.25, is_show=True, save_path=''):
        '''
        :param image_bgr:  输入图像
        :param thresh:     置信度阈值
        :param is_show:   是否将画框之后的原始图像返回
        :param save_path: 画框后的保存路径, eg='/home/aaa.jpg'
        :return:
        '''
        draw_bbox_image, detections = self.image_detection(image_bgr, self.network, self.class_names, self.class_colors,
                                                           thresh)
        if is_show:
            if save_path:
                cv2.imwrite(save_path, draw_bbox_image)
            return draw_bbox_image
        return detections


if __name__ == '__main__':
    # gpu 通过环境变量设置

    detect = Detect(metaPath=r'/home/cfg/sg.data',
                    configPath=r'/home/cfg/yolov4-sg.cfg',
                    weightPath=r'/home/yolov4-sg_best.weights',
                    gpu_id=1)

    # 读取单张图像
    # image_path = r'/home/aa.jpg'
    # image = cv2.imread(image_path, -1)
    # draw_bbox_image = detect.predict_image(image, save_path='./pred.jpg')

    # 读取文件夹
    image_root = r'/home/Datasets/image/'
    save_root = r'./output'
    if not os.path.exists(save_root):
        os.makedirs(save_root)
    for name in os.listdir(image_root):
        print(name)
        image = cv2.imread(os.path.join(image_root, name), -1)
        draw_bbox_image = detect.predict_image(image, save_path=os.path.join(save_root, name))

    # 读取视频
    # video_path = r'/home/Datasets/SHIJI_Fire/20200915_2.mp4'
    # video_save_path = r'/home/20200915_3_pred.mp4'
    # cap = cv2.VideoCapture(video_path)
    # # 获取视频的fps， width height
    # fps = int(cap.get(cv2.CAP_PROP_FPS))
    # width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    # height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    # count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    # print(count)
    # # 创建视频
    # fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    # # fourcc = cv2.VideoWriter_fourcc('I', '4', '2', '0')
    # video_writer = cv2.VideoWriter(video_save_path, fourcc=fourcc, fps=fps, frameSize=(width,height))
    # ret, frame = cap.read()  # ret表示下一帧还有没有 有为True
    # while ret:
    #     # 预测每一帧
    #     pred = detect.predict_image(frame)
    #     video_writer.write(pred)
    #     cv2.waitKey(fps)
    #     # 读取下一帧
    #     ret, frame = cap.read()
    #     print(ret)
    # cap.release()
    # cv2.destroyAllWindows()