<reponame>zzq96/detectron2
import os
os.environ["CUDA_VISIBLE_DEVICES"] = '0'
os.environ["DETECTRON2_DATASETS"] = 'data'
import logging
from collections import OrderedDict
import torch
from torch.nn.parallel import DistributedDataParallel
import cv2
import detectron2.utils.comm as comm
from detectron2.data import MetadataCatalog, build_detection_train_loader
from detectron2.engine import DefaultTrainer, default_argument_parser, default_setup, hooks, launch
from detectron2.utils.events import EventStorage
from detectron2.evaluation import (
    CityscapesInstanceEvaluator,
    CityscapesSemSegEvaluator,
    COCOEvaluator,
    COCOPanopticEvaluator,
    DatasetEvaluators,
    LVISEvaluator,
    PascalVOCDetectionEvaluator,
    SemSegEvaluator,
    verify_results,
)
from detectron2.modeling import GeneralizedRCNNWithTTA
from detectron2.utils.logger import setup_logger

from adet.data.dataset_mapper import DatasetMapperWithBasis
from adet.data.dataset_mapper_depth import DatasetMapperWithDepth
from adet.config import get_cfg
from adet.checkpoint import AdetCheckpointer
from adet.evaluation import TextEvaluator
import matplotlib.pyplot as plt
import numpy as np
from detectron2.engine import DefaultPredictor
import time


from detectron2.data import (
    build_detection_test_loader,
)
cfg = get_cfg()
# add project-specific config (e.g., TensorMask) here if you're not running a model in detectron2's core library
cfg.merge_from_file(r"/home/zqzhou/disk/Code/python/AdelaiDet/configs/BoxData/CondInst/CondInst_R_50_3x_sem.yaml")
cfg.MODEL.PROPOSAL_GENERATOR.NAME = 'FCOSM'
cfg.MODEL.META_ARCHITECTURE = 'CondInstM'
cfg.MODEL.FCOSM.USE_BOX_NMS = False
# cfg.MODEL.FCOSM.USE_SINGLE_NMS = False
cfg.MODEL.WEIGHTS = "/home/zqzhou/disk/Code/python/AdelaiDet/output/BoxData/CondInst/condinst_MS_R_50_3x_sem/old_0.0050/model_final.pth"
predictor_old = DefaultPredictor(cfg)


cfg = get_cfg()
# add project-specific config (e.g., TensorMask) here if you're not running a model in detectron2's core library
cfg.merge_from_file(r"/home/zqzhou/disk/Code/python/AdelaiDet/configs/BoxData/CondInst/CondInst_R_50_3x_sem.yaml")
cfg.MODEL.PROPOSAL_GENERATOR.NAME = 'FCOSM2'
cfg.MODEL.META_ARCHITECTURE = 'CondInstM'
cfg.MODEL.FCOSM.USE_BOX_NMS = False
# cfg.MODEL.FCOSM.NMS_TH = 0.8
cfg.MODEL.FCOSM.BOX_LEVEL_NUM_CONVS = 2
cfg.MODEL.FCOSM.NUM_BOX_CONVS = 2
cfg.MODEL.FCOSM.NUM_CLS_CONVS = 2
# cfg.MODEL.FCOSM.USE_SINGLE_NMS = False
cfg.MODEL.WEIGHTS = "/home/zqzhou/disk/Code/python/AdelaiDet/output/BoxData/CondInst/condinst_MS_R_50_3x_sem/new_222_0.0050/model_final.pth"
predictor = DefaultPredictor(cfg)
data_loader = build_detection_test_loader(cfg, 'BoxData_val')

# img_files = ["210120_0293", "210120_0283", "210120_0281", "210120_0269", "210120_0261", "210120_0113", "210104_0262"]
draw_location = False
isdraw_box = False
for idx, inputs in enumerate(data_loader):
    # if not inputs[0]['file_name'].split('/')[-1][:-4]in img_files:
    #     continue
    im = inputs[0]['image'].cpu().numpy().transpose([1,2,0])
    outputs =predictor(im)
    outputs_old =predictor_old(im)
    # We can use `Visualizer` to draw the predictions on the image.
    from detectron2.utils.visualizer import Visualizer
    output = outputs['instances'][outputs['instances'].scores > 0.3]
    output_old = outputs_old['instances'][outputs_old['instances'].scores > 0.3]
    # output = output[output.scores < 1.54]
    # output_1 = output
    output_1 = output[output.box_level==1]
    output_0 = output[output.box_level==0]

    # print(output.pred_classes)
    # print(output.pred_boxes)
    # print(output.scores)
    # print(output.ori_scores)
    # print(output.box_level)
    # print(output.ctrness_pred)
    font = cv2.FONT_HERSHEY_SIMPLEX
    v = Visualizer(im[:, :, ::-1].copy(), MetadataCatalog.get(cfg.DATASETS.TRAIN[0]), scale=1, use_polygon=False)

    out_1 = v.draw_instance_predictions(output_1.to("cpu"), isdraw_box)
    img_1 = out_1.get_image().copy()
    if draw_location:
        for location, fpn_level in zip(output_1.locations.cpu().numpy().astype(np.int32), output_1.fpn_levels.cpu().numpy()):
            # print(location)
            cv2.putText(img_1, str(int(fpn_level)), (location[0], location[1]), 1,1,(255, 0, 0), 2)
            img_1[location[1], location[0]] = [255, 0, 0]
            img_1[location[1] - 1, location[0] - 1] = [255, 0, 0]
            img_1[location[1] , location[0] -1] = [255, 0, 0]
            img_1[location[1] + 1, location[0] -1] = [255, 0, 0]
            img_1[location[1] - 1, location[0] ] = [255, 0, 0]
            img_1[location[1] + 1, location[0] ] = [255, 0, 0]
            img_1[location[1] - 1, location[0]] = [255, 0, 0]
            img_1[location[1] - 1, location[0] -1] = [255, 0, 0]
            img_1[location[1] - 1, location[0] -1] = [255, 0, 0]

    v = Visualizer(im[:, :, ::-1].copy(), MetadataCatalog.get(cfg.DATASETS.TRAIN[0]), scale=1, use_polygon=False)
    out_0 = v.draw_instance_predictions(output_0.to("cpu"), isdraw_box)
    img_0 = out_0.get_image().copy()
    if draw_location:
        for location, fpn_level in zip(output_0.locations.cpu().numpy().astype(np.int32), output_0.fpn_levels.cpu().numpy()):
            # print(location)
            cv2.putText(img_0, str(int(fpn_level)), (location[0], location[1]), 1,1,(255, 0, 0), 2)
            # print(location)
            if location[1] ==0:
                location[1] += 2
            if location[1] ==360:
                location[1] -= 2
            if location[0] ==0:
                location[0] += 2
            if location[0] ==640:
                location[0] -= 2
            img_0[location[1], location[0]] = [255, 0, 0]
            img_0[location[1] - 1, location[0] - 1] = [255, 0, 0]
            img_0[location[1] , location[0] -1] = [255, 0, 0]
            img_0[location[1] + 1, location[0] -1] = [255, 0, 0]
            img_0[location[1] - 1, location[0] ] = [255, 0, 0]
            img_0[location[1] + 1, location[0] ] = [255, 0, 0]
            img_0[location[1] - 1, location[0]] = [255, 0, 0]
            img_0[location[1] - 1, location[0] -1] = [255, 0, 0]
            img_0[location[1] - 1, location[0] -1] = [255, 0, 0]

    v = Visualizer(im[:, :, ::-1].copy(), MetadataCatalog.get(cfg.DATASETS.TRAIN[0]), scale=1, use_polygon=False)
    out_old = v.draw_instance_predictions(output_old.to("cpu"), isdraw_box)
    img_old = out_old.get_image().copy()
    if draw_location:
        for location, fpn_level in zip(output_old.locations.cpu().numpy().astype(np.int32), output_old.fpn_levels.cpu().numpy()):
            # print(location)
            cv2.putText(img_old, str(int(fpn_level)), (location[0], location[1]), 1,1,(255, 0, 0), 2)
            # print(location)
            if location[1] ==0:
                location[1] += 2
            if location[1] ==360:
                location[1] -= 2
            if location[0] ==0:
                location[0] += 2
            if location[0] ==640:
                location[0] -= 2
            img_old[location[1], location[0]] = [255, 0, 0]
            img_old[location[1] - 1, location[0] - 1] = [255, 0, 0]
            img_old[location[1] , location[0] -1] = [255, 0, 0]
            img_old[location[1] + 1, location[0] -1] = [255, 0, 0]
            img_old[location[1] - 1, location[0] ] = [255, 0, 0]
            img_old[location[1] + 1, location[0] ] = [255, 0, 0]
            img_old[location[1] - 1, location[0]] = [255, 0, 0]
            img_old[location[1] - 1, location[0] -1] = [255, 0, 0]
            img_old[location[1] - 1, location[0] -1] = [255, 0, 0]

    v = Visualizer(im[:, :, ::-1].copy(), MetadataCatalog.get(cfg.DATASETS.TRAIN[0]), scale=1, use_polygon=False)
    out = v.draw_instance_predictions(output.to("cpu"), isdraw_box)
    row_0 = np.concatenate([im[:, :, ::-1], img_old], axis=1)
    row_1 = np.concatenate([img_0, img_1], axis=1)
    img_all = np.concatenate([row_0, row_1], axis=0)
    plt.imsave(os.path.join('tmp_dir','test', inputs[0]['file_name'].split('/')[-1]), img_all)
    plt.imsave(os.path.join('tmp_dir','test', inputs[0]['file_name'].split('/')[-1] + "all.png"), out.get_image())
    # plt.imsave(os.path.join('tmp_dir','paper', inputs[0]['file_name'].split('/')[-1][:-4]+"base.png"), img_old)
    # plt.imsave(os.path.join('tmp_dir','paper', inputs[0]['file_name'].split('/')[-1][:-4]+"ours0.png"), img_0)
    # plt.imsave(os.path.join('tmp_dir','paper', inputs[0]['file_name'].split('/')[-1][:-4]+"ours1.png"), img_1)
    a = 1