# -*- encoding:utf-8 -*-
# @Time    : 2021/1/3 15:15
# @Author  : gfjiang
import os.path as osp
import mmcv
import numpy as np
import cvtools
import matplotlib.pyplot as plt
import cv2.cv2 as cv
from functools import partial
import torch
import math
from cvtools.utils.path import add_prefix_filename_suffix

from mmdet.ops import nms
from mmdet.apis import init_detector, inference_detector


def draw_features(module, input, output, work_dir='./'):
    x = output.cpu().numpy()
    out_channels = list(output.shape)[1]
    height = int(math.sqrt(out_channels))
    width = height
    if list(output.shape)[2] < 128:
        return
    fig = plt.figure(figsize=(32, 32))
    fig.subplots_adjust(left=0.05, right=0.95, bottom=0.05, top=0.95, wspace=0.05, hspace=0.05)
    for i in range(height * width):
        plt.subplot(height, width, i + 1)
        plt.axis('off')
        img = x[0, i, :, :]
        pmin = np.min(img)
        pmax = np.max(img)
        img = ((img - pmin) / (pmax - pmin + 0.000001))*255  # float在[0，1]之间，转换成0-255
        img = img.astype(np.uint8)  # 转成unit8
        img = cv.applyColorMap(img, cv.COLORMAP_JET)  # 生成heat map
        img = img[:, :, ::-1]  # 注意cv2（BGR）和matplotlib(RGB)通道是相反的
        plt.imshow(img)
        # print("{}/{}".format(i,width*height))
    savename = get_image_name_for_hook(module, work_dir)
    fig.savefig(savename, dpi=100)
    fig.clf()
    plt.close()


def get_image_name_for_hook(module, work_dir='./'):
    """
    Generate image filename for hook function

    Parameters:
    -----------
    module: module of neural network
    """
    # os.makedirs(work_dir, exist_ok=True)
    module_name = str(module)
    base_name = module_name.split('(')[0]
    index = 0
    image_name = '.'  # '.' is surely exist, to make first loop condition True
    while osp.exists(image_name):
        index += 1
        image_name = osp.join(
            work_dir, 'feats', '%s_%d.png' % (base_name, index))
    return image_name


class AerialDetectionOBB(object):

    def __init__(self, config, pth):
        self.imgs = []
        self.cfg = mmcv.Config.fromfile(config)
        self.pth = pth
        print('loading model {} ...'.format(pth))
        self.model = init_detector(self.cfg, self.pth, device='cuda:0')
        self.results = []
        self.img_detected = []
        # self.vis_feats((torch.nn.Conv2d, torch.nn.MaxPool2d))

    def __call__(self,
                 imgs_or_path,
                 det_thrs=0.5,
                 vis=False,
                 vis_thr=0.5,
                 save_root=''):
        if isinstance(imgs_or_path, str):
            self.imgs += cvtools.get_files_list(imgs_or_path)
        else:
            self.imgs += imgs_or_path
        prog_bar = mmcv.ProgressBar(len(self.imgs))
        for _, img in enumerate(self.imgs):
            self.detect(img, det_thrs=det_thrs, vis=vis,
                        vis_thr=vis_thr, save_root=save_root)
            prog_bar.update()

    def detect(self,
               img,
               det_thrs=0.5,
               vis=False,
               vis_thr=0.5,
               save_root=''):
        result = inference_detector(self.model, img)
        # result = self.nms(result)
        if isinstance(det_thrs, float):
            det_thrs = [det_thrs] * len(result)
        if vis:
            to_file = osp.join(save_root, osp.basename(img))
            to_file = add_prefix_filename_suffix(to_file, suffix='_obb')
            self.vis(img, result, vis_thr=vis_thr, to_file=to_file)
        result = [det[det[..., -1] > det_thr] for det, det_thr
                  in zip(result, det_thrs)]
        if len(result) == 0:
            print('detect: image {} has no object.'.format(img))
        self.img_detected.append(img)
        self.results.append(result)
        return result

    def nms(self, result, nms_th=0.3):
        dets_num = [len(det_cls) for det_cls in result]
        result = np.vstack(result)
        _, ids = nms(result, nms_th)
        total_num = 0
        nms_result = []
        for num in dets_num:
            ids_cls = ids[np.where((total_num <= ids) & (ids < num))[0]]
            nms_result.append(result[ids_cls])
            total_num += num
        return nms_result

    def vis(self, img, bbox_result, vis_thr=0.5,
            to_file='vis.jpg'):
        bboxes = np.vstack(bbox_result)
        labels = [
            np.full(bbox.shape[0], i, dtype=np.int32)
            for i, bbox in enumerate(bbox_result)
        ]
        labels = np.concatenate(labels)
        inds = np.where(bboxes[:, -1] > vis_thr)[0]
        bboxes = bboxes[inds]
        labels = labels[inds]
        texts = [self.model.CLASSES[index]+'|'+str(round(bbox[-1], 2))
                 for index, bbox in zip(labels, bboxes)]
        img = cvtools.draw_boxes_texts(
            img, bboxes[:, :-1], box_format='polygon', line_width=2)
        cvtools.imwrite(img, to_file)
    
    def vis_feats(self, modules_for_plot):
        h, w = self.cfg.data.train.img_scale
        for name, module in self.model.named_modules():
            if isinstance(module, modules_for_plot):
                draw_features_func = partial(
                    draw_features, work_dir=self.cfg.work_dir)
                module.register_forward_hook(draw_features_func)

    def save_results(self, save):
        str_results = ''
        for i, img in enumerate(self.img_detected):
            result = self.results[i]
            img = osp.basename(img)
            for cls_index, dets in enumerate(result):
                cls = self.model.CLASSES[cls_index]
                for box in dets:
                    bbox_str = ','.join(map(str, map(int, box[:4])))
                    str_results += ' '.join([img, cls, bbox_str]) + '\n'
        with open(save, 'w') as f:
            f.write(str_results)


if __name__ == '__main__':
    config_file = 'configs/DOTA/faster_rcnn_RoITrans_r50_fpn_1x_dota_1gpus_mdanet2.py'
    pth_file = 'work_dirs/faster_rcnn_RoITrans_r50_fpn_1x_dota_1gpus_mdanet2/epoch_12.pth'
    detector = AerialDetectionOBB(config_file, pth_file)
    detector('/media/data/DOTA/crop/P2701_2926_1597_3949_2620.png', vis=True, 
             save_root='work_dirs/attention_vis/')
    detector.save_results('work_dirs/faster_rcnn_RoITrans_r50_fpn_1x_dota_1gpus_mdanet2/detect_result.txt')

[((26, 10, 26, 38), 'matplotlib.pyplot.figure', 'plt.figure', (), '', True, 'import matplotlib.pyplot as plt\n'), ((43, 4, 43, 15), 'matplotlib.pyplot.close', 'plt.close', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((59, 10, 59, 32), 'os.path.exists', 'osp.exists', ({(59, 21, 59, 31): 'image_name'}, {}), '(image_name)', True, 'import os.path as osp\n'), ((22, 17, 22, 40), 'math.sqrt', 'math.sqrt', ({(22, 27, 22, 39): 'out_channels'}, {}), '(out_channels)', False, 'import math\n'), ((29, 8, 29, 41), 'matplotlib.pyplot.subplot', 'plt.subplot', ({(29, 20, 29, 26): 'height', (29, 28, 29, 33): 'width', (29, 35, 29, 40): '(i + 1)'}, {}), '(height, width, i + 1)', True, 'import matplotlib.pyplot as plt\n'), ((30, 8, 30, 23), 'matplotlib.pyplot.axis', 'plt.axis', ({(30, 17, 30, 22): '"""off"""'}, {}), "('off')", True, 'import matplotlib.pyplot as plt\n'), ((32, 15, 32, 26), 'numpy.min', 'np.min', ({(32, 22, 32, 25): 'img'}, {}), '(img)', True, 'import numpy as np\n'), ((33, 15, 33, 26), 'numpy.max', 'np.max', ({(33, 22, 33, 25): 'img'}, {}), '(img)', True, 'import numpy as np\n'), ((36, 14, 36, 52), 'cv2.cv2.applyColorMap', 'cv.applyColorMap', ({(36, 31, 36, 34): 'img', (36, 36, 36, 51): 'cv.COLORMAP_JET'}, {}), '(img, cv.COLORMAP_JET)', True, 'import cv2.cv2 as cv\n'), ((38, 8, 38, 23), 'matplotlib.pyplot.imshow', 'plt.imshow', ({(38, 19, 38, 22): 'img'}, {}), '(img)', True, 'import matplotlib.pyplot as plt\n'), ((61, 21, 62, 64), 'os.path.join', 'osp.join', ({(62, 12, 62, 20): 'work_dir', (62, 22, 62, 29): '"""feats"""', (62, 31, 62, 63): "'%s_%d.png' % (base_name, index)"}, {}), "(work_dir, 'feats', '%s_%d.png' % (base_name, index))", True, 'import os.path as osp\n'), ((70, 19, 70, 47), 'mmcv.Config.fromfile', 'mmcv.Config.fromfile', ({(70, 40, 70, 46): 'config'}, {}), '(config)', False, 'import mmcv\n'), ((73, 21, 73, 71), 'mmdet.apis.init_detector', 'init_detector', (), '', False, 'from mmdet.apis import init_detector, inference_detector\n'), ((100, 17, 100, 52), 'mmdet.apis.inference_detector', 'inference_detector', ({(100, 36, 100, 46): 'self.model', (100, 48, 100, 51): 'img'}, {}), '(self.model, img)', False, 'from mmdet.apis import init_detector, inference_detector\n'), ((118, 17, 118, 34), 'numpy.vstack', 'np.vstack', ({(118, 27, 118, 33): 'result'}, {}), '(result)', True, 'import numpy as np\n'), ((119, 17, 119, 36), 'mmdet.ops.nms', 'nms', ({(119, 21, 119, 27): 'result', (119, 29, 119, 35): 'nms_th'}, {}), '(result, nms_th)', False, 'from mmdet.ops import nms\n'), ((130, 17, 130, 39), 'numpy.vstack', 'np.vstack', ({(130, 27, 130, 38): 'bbox_result'}, {}), '(bbox_result)', True, 'import numpy as np\n'), ((135, 17, 135, 39), 'numpy.concatenate', 'np.concatenate', ({(135, 32, 135, 38): 'labels'}, {}), '(labels)', True, 'import numpy as np\n'), ((141, 14, 142, 68), 'cvtools.draw_boxes_texts', 'cvtools.draw_boxes_texts', (), '', False, 'import cvtools\n'), ((143, 8, 143, 37), 'cvtools.imwrite', 'cvtools.imwrite', ({(143, 24, 143, 27): 'img', (143, 29, 143, 36): 'to_file'}, {}), '(img, to_file)', False, 'import cvtools\n'), ((85, 25, 85, 61), 'cvtools.get_files_list', 'cvtools.get_files_list', ({(85, 48, 85, 60): 'imgs_or_path'}, {}), '(imgs_or_path)', False, 'import cvtools\n'), ((106, 22, 106, 72), 'cvtools.utils.path.add_prefix_filename_suffix', 'add_prefix_filename_suffix', (), '', False, 'from cvtools.utils.path import add_prefix_filename_suffix\n'), ((132, 12, 132, 53), 'numpy.full', 'np.full', (), '', True, 'import numpy as np\n'), ((136, 15, 136, 48), 'numpy.where', 'np.where', ({(136, 24, 136, 47): '(bboxes[:, (-1)] > vis_thr)'}, {}), '(bboxes[:, (-1)] > vis_thr)', True, 'import numpy as np\n'), ((157, 18, 157, 35), 'os.path.basename', 'osp.basename', ({(157, 31, 157, 34): 'img'}, {}), '(img)', True, 'import os.path as osp\n'), ((105, 42, 105, 59), 'os.path.basename', 'osp.basename', ({(105, 55, 105, 58): 'img'}, {}), '(img)', True, 'import os.path as osp\n'), ((149, 37, 150, 62), 'functools.partial', 'partial', (), '', False, 'from functools import partial\n'), ((123, 26, 123, 68), 'numpy.where', 'np.where', ({(123, 35, 123, 67): '((total_num <= ids) & (ids < num))'}, {}), '((total_num <= ids) & (ids < num))', True, 'import numpy as np\n')]