tools/eval_boxes.py

import json
import warnings

from detectron2.modeling.meta_arch.rcnn import GeneralizedRCNN

from cubercnn.data.build import build_detection_train_loader
warnings.filterwarnings("ignore", message="Overwriting tiny_vit_21m_512 in registry")
warnings.filterwarnings("ignore", message="Overwriting tiny_vit_21m_384 in registry")
warnings.filterwarnings("ignore", message="Overwriting tiny_vit_21m_224 in registry")
warnings.filterwarnings("ignore", message="Overwriting tiny_vit_11m_224 in registry")
warnings.filterwarnings("ignore", message="Overwriting tiny_vit_5m_224 in registry")

# Copyright (c) Meta Platforms, Inc. and affiliates
from contextlib import ExitStack
import logging
import os
from detectron2.data.detection_utils import convert_image_to_rgb
from detectron2.evaluation.evaluator import inference_context
from detectron2.utils.visualizer import Visualizer
from matplotlib import pyplot as plt
import numpy as np
import torch
import datetime
import detectron2.utils.comm as comm
from detectron2.checkpoint import DetectionCheckpointer
from detectron2.config import get_cfg
from detectron2.data import MetadataCatalog
from detectron2.engine import (
    default_argument_parser, 
    default_setup, 
)
from detectron2.utils.logger import setup_logger
import torch.nn as nn
from tqdm import tqdm
import pickle

from ProposalNetwork.utils.utils import show_mask2
from cubercnn.data.dataset_mapper import DatasetMapper3D
from cubercnn.evaluation.omni3d_evaluation import instances_to_coco_json

logger = logging.getLogger("cubercnn")

from cubercnn.config import get_cfg_defaults
from cubercnn.data import (
    build_detection_test_loader,
    simple_register
)
from cubercnn.evaluation import (
    Omni3DEvaluationHelper,
)
from cubercnn.modeling.meta_arch import build_model
from cubercnn import util, vis, data
# even though this import is unused, it initializes the backbone registry
from cubercnn.modeling.backbone import build_dla_from_vision_fpn_backbone

from matplotlib.patches import PathPatch
from matplotlib.path import Path

def create_striped_patch(ax, x_start, x_end, color, alpha=0.3):
    ylim = ax.get_ylim()
    stripe_height = (ylim[1] - ylim[0]) / 20  # Adjust stripe height as needed
    vertices = []
    codes = []
    for y in np.arange(ylim[0], ylim[1], stripe_height * 2):
        vertices.extend([(x_start, y), (x_end, y + stripe_height), (x_end, y + stripe_height * 2), (x_start, y + stripe_height)])
        codes.extend([Path.MOVETO, Path.LINETO, Path.LINETO, Path.LINETO])
    path = Path(vertices, codes)
    patch = PathPatch(path, facecolor=color, edgecolor='none', alpha=alpha, hatch='/')
    ax.add_patch(patch)

color_palette = ['#008dff','#ff73bf','#c701ff','#4ecb8d','#ff9d3a','#f0c571','#384860','#d83034']

def inference_on_dataset(model, data_loader, experiment_type, proposal_function):
    """
    Run model on the data_loader. 
    Also benchmark the inference speed of `model.__call__` accurately.
    The model will be used in eval mode.

    Args:
        model (callable): a callable which takes an object from
            `data_loader` and returns some outputs.

            If it's an nn.Module, it will be temporarily set to `eval` mode.
            If you wish to evaluate a model in `training` mode instead, you can
            wrap the given model and override its behavior of `.eval()` and `.train()`.
        data_loader: an iterable object with a length.
            The elements it generates will be the inputs to the model.

    Returns:
        The return value of `evaluator.evaluate()`
    """
    
    logger.info("Start inference on {} batches".format(len(data_loader)))

    total = len(data_loader)  # inference data loader must have a fixed length

    inference_json = []

    with ExitStack() as stack:
        if isinstance(model, nn.Module):
            stack.enter_context(inference_context(model))
        stack.enter_context(torch.no_grad())

        for idx, inputs in tqdm(enumerate(data_loader), desc="Average Precision", total=total):
            outputs = model(inputs, experiment_type, proposal_function)
            for input, output in zip(inputs, outputs):

                prediction = {
                    "image_id": input["image_id"],
                    "K": input["K"],
                    "width": input["width"],
                    "height": input["height"],
                }

                # convert to json format
                instances = output["instances"].to('cpu')
                # instances = output["instances"].to('cpu')
                prediction["instances"] = instances_to_coco_json(instances, input["image_id"])

                # store in overall predictions
                inference_json.append(prediction)

    return inference_json

def percent_of_boxes(model, data_loader, experiment_type, proposal_functions):
    '''make the detection that have a certain 3D IoU score plots
    if you give the proposal function as input to argparser as:
    
    `PLOT.PROPOSAL_FUNC, ['random', 'z', 'xy', 'dim', 'rotation', 'aspect' ,'full']`
    
    it will work
    '''
    total = len(data_loader)  # inference data loader must have a fixed length
    torch.set_float32_matmul_precision('high')

    with ExitStack() as stack:
        if isinstance(model, nn.Module):
            stack.enter_context(inference_context(model))
        stack.enter_context(torch.no_grad())

        # if not os.path.exists('ProposalNetwork/output/outputs.pkl'):
        if True:
            outputs = []
            for i, inputs in tqdm(enumerate(data_loader), desc=f"IoU3D plots, proposal method: {proposal_functions}", total=total):
                output = model(inputs, experiment_type, proposal_functions)
                outputs.append(output.cpu().numpy())
        #     with open('ProposalNetwork/output/outputs.pkl', 'wb') as f:
        #         pickle.dump(outputs, f)
        # else:
        #     with open('ProposalNetwork/output/outputs_10k.pkl', 'rb') as f:
        #         outputs = pickle.load(f)
        xlim = [0.2,1]
        IoUat = [0.25, 0.4, 0.6]
        n_proposals = outputs[0].shape[-1]
        
        fig, axes = plt.subplots(1, figsize=(7.5,5))
        fig2, axes2 = plt.subplots(1, len(IoUat), figsize=(20,5))
        axes.set_ylabel('Detection rate')
        axes.set_ylim([0,1])
        axes.grid(True)
        axes.set_xlabel('3D Intersection over Union')
        axes.set_xlim(xlim)
        axes.set_title(f'Varying proposal method, {n_proposals} proposals')
        for Iou, ax in zip(IoUat, axes2):
            ax.set_ylabel('Detection rate')
            ax.set_ylim([0,1])
            ax.grid(True)
            ax.set_title(f'Variants, IoU3D = {Iou}')
            ax.set_xlim([1,n_proposals])
            ax.set_xlabel('Number of Proposals')

        for k, proposal_function in enumerate(proposal_functions):
            IoU3Ds = np.concatenate([x[:,k,:] for x in outputs])
            maxIOU_per_instance = np.max(IoU3Ds, axis=1)
            np.random.shuffle(IoU3Ds.T) #transpose to shuffle along the proposal axis
            # detection rate vs. IoU3D
            thresholds = np.arange(xlim[0],xlim[1],0.025)
            detection_rate = np.zeros(len(thresholds))
            for i in range(len(thresholds)):
                detection_rate[i] = np.mean(maxIOU_per_instance > thresholds[i],axis=0)

            # detection rate vs. no. of proposals
            detection_rates = np.zeros((len(IoUat), IoU3Ds.shape[1]))
            for i, IoU in enumerate(IoUat):
                mask_positive_values = IoU3Ds >= IoU
                first_above_threshold = np.argmax(mask_positive_values, axis=1)
                any_above_threshold = mask_positive_values.any(axis=1)
                detection_rate_at_IoU = np.zeros_like(IoU3Ds)

                for j in range(IoU3Ds.shape[0]):
                    if any_above_threshold[j]:
                        detection_rate_at_IoU[j, :first_above_threshold[j]] = 0
                        detection_rate_at_IoU[j, first_above_threshold[j]:] = 1
                
                detection_rates[i] = np.mean(detection_rate_at_IoU, axis=0)

            axes.plot(thresholds, detection_rate, label=f'{proposal_function}', color=color_palette[k])
            for j, ax in enumerate(axes2):
                ax.semilogx(list(range(1, n_proposals+1)), detection_rates[j], label=f'{proposal_function}', color=color_palette[k])
    for ax in axes2:
        ax.legend()
    axes.legend()
    print("saved to 'ProposalNetwork/output/'")
    fig.savefig(f'ProposalNetwork/output/detection_rate_{n_proposals}.png', dpi=300, bbox_inches='tight')
    fig2.savefig(f'ProposalNetwork/output/IoU_varying_{n_proposals}.png', dpi=300, bbox_inches='tight')


    return
            

def mean_average_best_overlap(model, data_loader, experiment_type, proposal_function):
        
    total = len(data_loader)  # inference data loader must have a fixed length
    os.makedirs('output/pkl_files', exist_ok=True)
    with ExitStack() as stack:
        if isinstance(model, nn.Module):
            stack.enter_context(inference_context(model))
        stack.enter_context(torch.no_grad())

        outputs = []
        for i, inputs in tqdm(enumerate(data_loader), desc="Mean average best overlap plots", total=total):
            logger.info('iteration %s',i)
            output = model(inputs, experiment_type, proposal_function)
            # p_info, IoU3D, score_IoU2D, score_seg, score_dim, score_combined, score_random, score_point_cloud, stat_empty_boxes, stats_im, stats_off
            if output is not None:
                outputs.append(output)
        with open('output/pkl_files/outputs_'+str(proposal_function)+'.pkl', 'wb') as f:
            pickle.dump(outputs, f)

        # Create output folder
        if not os.path.exists('ProposalNetwork/output/MABO_'+str(proposal_function)):
            os.makedirs('ProposalNetwork/output/MABO_'+str(proposal_function)) # This is maybe unnecessary
            os.makedirs('ProposalNetwork/output/MABO_'+str(proposal_function)+'/vis/')

        # mean over all the outputs
        Iou2D             = np.concatenate([np.array(sublist) for sublist in (x[1] for x in outputs)])
        score_seg         = np.concatenate([np.array(sublist) for sublist in (x[2] for x in outputs)])
        score_dim         = np.concatenate([np.array(sublist) for sublist in (x[3] for x in outputs)])
        score_combined    = np.concatenate([np.array(sublist) for sublist in (x[4] for x in outputs)])
        score_random      = np.concatenate([np.array(sublist) for sublist in (x[5] for x in outputs)])
        score_point_cloud = np.concatenate([np.array(sublist) for sublist in (x[6] for x in outputs)])
        score_seg_mod     = np.concatenate([np.array(sublist) for sublist in (x[10] for x in outputs)])
        score_corners     = np.concatenate([np.array(sublist) for sublist in (x[11] for x in outputs)])
        stat_empty_boxes  = np.array([x[7] for x in outputs])
        combinations      = np.mean(np.concatenate([np.array(sublist) for sublist in (x[12] for x in outputs)]),axis=0)
        #logger.info('Percentage of cubes with no intersection:',np.mean(stat_empty_boxes))
        print('Percentage of cubes with no intersection:',np.mean(stat_empty_boxes))
        logger.info('combination scores:%s',combinations)
        #print('combination scores:',combinations)
        print('best combination is C'+str(np.argmax(combinations)+1))

        

        Iou2D = Iou2D.mean(axis=0)
        score_seg = score_seg.mean(axis=0)
        score_dim = score_dim.mean(axis=0)
        score_combined = score_combined.mean(axis=0)
        score_random = score_random.mean(axis=0)
        score_point_cloud = score_point_cloud.mean(axis=0)
        score_seg_mod = score_seg_mod.mean(axis=0)
        score_corners = score_corners.mean(axis=0)
        total_num_instances = np.sum([x[0].gt_boxes3D.shape[0] for x in outputs])
        
        print('Avg IoU of chosen cube:', score_combined[0])
        print('Best possible IoU:', score_combined[-1])

        x_range = np.arange(1,1001)
        plt.figure(figsize=(8,5))
        plt.plot(x_range,score_combined, linestyle='-',c=color_palette[6], label='combined') 
        plt.plot(x_range,score_dim, linestyle='-',c=color_palette[5],label='dim') 
        plt.plot(x_range,score_seg, linestyle='-',c=color_palette[2],label='segment')
        plt.plot(x_range,Iou2D, linestyle='-',c=color_palette[4],label='2D IoU') 
        plt.plot(x_range,score_corners, linestyle='-',c=color_palette[7],label='corner dist')
        plt.plot(x_range,score_random, linestyle='-',c='grey',label='random') 
        plt.plot(x_range,score_point_cloud, linestyle='-',c=color_palette[3],label='point cloud')
        plt.plot(x_range,score_seg_mod, linestyle='-',c=color_palette[0],label='mod segment')
        plt.grid(True)
        plt.xscale('log')
        plt.xticks([1, 10, 100, 1000], ['1', '10', '100', '1000'])
        plt.xlim(left=1, right=len(Iou2D))
        plt.xlabel('Number of Proposals')
        plt.ylabel('3D IoU')
        plt.legend()
        plt.title('Average Best Overlap vs Number of Proposals ({} images, {} instances)'.format(1+i,total_num_instances))
        f_name = os.path.join('ProposalNetwork/output/MABO_'+str(proposal_function), 'MABO_'+str(proposal_function)+'.png')
        plt.savefig(f_name, dpi=300, bbox_inches='tight')
        plt.close()
        #logger.info('saved to ', f_name)
        print('saved to ', f_name)
        
        # Statistics
        stats = torch.cat([x[8] for x in outputs],dim=0)
        num_bins = 40
        titles = ['x','y','z']
        plt.figure(figsize=(15, 5))
        plt.suptitle("Distribution of Ground Truths in Normalised Searched Range", fontsize=20)
        for i,title in enumerate(titles):
            ax = plt.subplot(1, 3, 1+i)
            plt.hist(stats[:,i].numpy(), bins=num_bins, color=color_palette[6],density=True, zorder=2)
            plt.axvline(x=0, color='#97a6c4',zorder=2)
            plt.axvline(x=1, color='#97a6c4',zorder=2)
            create_striped_patch(ax, 0, 1, '#97a6c4', alpha=0.8)
            plt.title(title)
        f_name = os.path.join('ProposalNetwork/output/MABO_'+str(proposal_function), 'stats_center_'+str(proposal_function)+'.png')
        plt.savefig(f_name, dpi=300, bbox_inches='tight')
        plt.close()
        print('saved to ', f_name)
        num_bins = 120
        plt.figure(figsize=(15, 5))
        plt.suptitle("Distribution of Ground Truths in Normalised Searched Range", fontsize=20)
        for i,title in enumerate(titles):
            ax = plt.subplot(1, 3, 1+i)
            plt.hist(stats[:,i].numpy(), bins=num_bins, color=color_palette[6],density=True, zorder=2)
            plt.axvline(x=0, color='#97a6c4', zorder=2)
            plt.axvline(x=1, color='#97a6c4', zorder=2)
            create_striped_patch(ax, 0, 1, '#97a6c4', alpha=0.8)
            plt.xlim([-2,2])
            plt.title(title)
        f_name = os.path.join('ProposalNetwork/output/MABO_'+str(proposal_function), 'stats_center_zoom_'+str(proposal_function)+'.png')
        plt.savefig(f_name, dpi=300, bbox_inches='tight')
        plt.close()
        print('saved to ', f_name)
        num_bins = 40
        titles = ['w','h','l']
        plt.figure(figsize=(15, 5))
        plt.suptitle("Distribution of Ground Truths in Normalised Searched Range", fontsize=20)
        for i,title in enumerate(titles):
            ax = plt.subplot(1, 3, 1+i)
            plt.hist(stats[:,3+i].numpy(), bins=num_bins, color=color_palette[6],density=True, zorder=2)
            plt.axvline(x=0, color='#97a6c4', zorder=2)
            plt.axvline(x=1, color='#97a6c4', zorder=2)
            create_striped_patch(ax, 0, 1, '#97a6c4', alpha=0.8)
            plt.title(title)
        f_name = os.path.join('ProposalNetwork/output/MABO_'+str(proposal_function), 'stats_dim_'+str(proposal_function)+'.png')
        plt.savefig(f_name, dpi=300, bbox_inches='tight')
        plt.close()
        print('saved to ', f_name)
        titles = ['rx','ry','rz']
        plt.figure(figsize=(15, 5))
        plt.suptitle("Distribution of Ground Truths in Normalised Searched Range", fontsize=20)
        for i,title in enumerate(titles):
            ax = plt.subplot(1, 3, 1+i)
            plt.hist(stats[:,6+i].numpy(), bins=num_bins, color=color_palette[6],density=True, zorder=2)
            plt.axvline(x=0, color='#97a6c4', zorder=2)
            plt.axvline(x=1, color='#97a6c4', zorder=2)
            create_striped_patch(ax, 0, 1, '#97a6c4', alpha=0.8)
            plt.title(title)
        f_name = os.path.join('ProposalNetwork/output/MABO_'+str(proposal_function), 'stats_rot_'+str(proposal_function)+'.png')
        plt.savefig(f_name, dpi=300, bbox_inches='tight')
        plt.close()
        print('saved to ', f_name)

        stats_off = np.concatenate([np.array(sublist) for sublist in (x[9] for x in outputs)])
        plt.figure(figsize=(15, 15))
        for i,title in enumerate(titles):
            plt.subplot(3, 3, 1+i)
            plt.scatter(stats_off[:,1+i],stats_off[:,0], color=color_palette[6])
            plt.title(title)
        f_name = os.path.join('ProposalNetwork/output/MABO_'+str(proposal_function), 'stats_off_'+str(proposal_function)+'.png')
        plt.savefig(f_name, dpi=300, bbox_inches='tight')
        plt.close()
        #logger.info('saved to ', f_name)
        print('saved to ', f_name)

        plt.figure(figsize=(15, 15))
        for i,title in enumerate(titles):
            plt.subplot(3, 3, 1+i)
            plt.scatter(stats_off[:,1+i],stats_off[:,0], color=color_palette[6])
            plt.title(title)
            plt.xlim([0,2])
            plt.ylim([0,1])
        f_name = os.path.join('ProposalNetwork/output/MABO_'+str(proposal_function), 'stats_off_zoom_'+str(proposal_function)+'.png')
        plt.savefig(f_name, dpi=300, bbox_inches='tight')
        plt.close()
        #logger.info('saved to ', f_name)
        print('saved to ', f_name)
        
        # ## for vis
        d_iter = iter(data_loader)
        for i , _ in tqdm(enumerate(outputs), desc="Plotting every single image", total=len(outputs)):
            p_info = outputs[i][0]
            pred_box_classes_names = [util.MetadataCatalog.get('omni3d_model').thing_classes[label] for label in p_info.pred_cubes.labels.cpu().numpy()]
            box_size = p_info.pred_cubes.num_instances
            for x in range(box_size-len(pred_box_classes_names)):
                pred_box_classes_names.append(f'z={p_info.pred_cubes[x].dimensions[2]}, s={p_info.pred_cubes[x].scores}')
            colors = [np.concatenate([np.random.random(3), np.array([0.6])], axis=0) for _ in range(box_size)]
            fig, (ax, ax1) = plt.subplots(2,1, figsize=(14, 10))
            input = next(d_iter)[0]
            images_raw = input['image']
            
            org_img = convert_image_to_rgb(images_raw.permute(1,2,0).cpu().numpy(), 'BGR').copy()
            v_pred = Visualizer(org_img, None)
            v_pred = v_pred.overlay_instances(
                boxes=p_info.pred_boxes[0:box_size].tensor.cpu().numpy()
                , assigned_colors=colors
            )
            prop_img = v_pred.get_image()
            pred_cube_meshes = [p_info.pred_cubes[j].get_cubes().__getitem__(0).detach() for j in range(box_size)]
            img_3DPR, img_novel, _ = vis.draw_scene_view(prop_img, p_info.K, pred_cube_meshes, text=pred_box_classes_names, blend_weight=0.5, blend_weight_overlay=0.85,scale = prop_img.shape[0],colors=colors)
            vis_img_3d = img_3DPR.astype(np.uint8)
            vis_img_3d = show_mask2(p_info.mask_per_image.cpu().numpy(), vis_img_3d, random_color=colors) # NOTE Uncomment to add segmentation mask to pred image
            #vis_img_3d = np.concatenate((vis_img_3d, np.zeros((vis_img_3d.shape[0],vis_img_3d.shape[1],1))), axis=-1)
            ax.set_title('Predicted')
            # expand_img_novel to have alpha channel
            img_novel = np.concatenate((img_novel, np.ones_like(img_novel[:,:,0:1])*255), axis=-1)/255
            ax.imshow(np.concatenate((vis_img_3d, img_novel), axis=1))
            box_size = len(p_info.gt_cube_meshes)
            v_pred = Visualizer(org_img, None)
            v_pred = v_pred.overlay_instances(boxes=p_info.gt_boxes[0:box_size].tensor.cpu().numpy(), assigned_colors=colors)
            prop_img = v_pred.get_image()
            gt_box_classes_names = [util.MetadataCatalog.get('omni3d_model').thing_classes[i] for i in p_info.gt_box_classes]
            img_3DPR, img_novel, _ = vis.draw_scene_view(prop_img, p_info.K, p_info.gt_cube_meshes,text=gt_box_classes_names, blend_weight=0.5, blend_weight_overlay=0.85,scale = prop_img.shape[0],colors=colors)
            vis_img_3d = img_3DPR.astype(np.uint8)
            im_concat = np.concatenate((vis_img_3d, img_novel), axis=1)
            ax1.set_title('GT')
            ax1.imshow(im_concat)
            f_name = os.path.join('ProposalNetwork/output/MABO_'+str(proposal_function)+'/vis/', f'vis_{i}.png')
            plt.savefig(f_name, dpi=300, bbox_inches='tight')
            plt.close()

            # with open(f'ProposalNetwork/output/MABO/vis/out_{i}.pkl', 'wb') as f:
            #     out = images_raw.permute(1,2,0).cpu().numpy(), K, p_info.mask_per_image.cpu().numpy(), p_info.gt_boxes3D, p_info.gt_boxes[0], pred_box_classes_names
            #     # im, K, mask, gt_boxes3D, gt_boxes, pred_box_classes_names
            #     pickle.dump(out, f)




def do_test(cfg, model, iteration='final', storage=None):
        
    filter_settings = data.get_filter_settings_from_cfg(cfg)    
    filter_settings['visibility_thres'] = cfg.TEST.VISIBILITY_THRES
    filter_settings['truncation_thres'] = cfg.TEST.TRUNCATION_THRES
    filter_settings['min_height_thres'] = 0.0625
    filter_settings['max_depth'] = 1e8

    dataset_names_test = cfg.DATASETS.TEST
    only_2d = cfg.MODEL.ROI_CUBE_HEAD.LOSS_W_3D == 0.0
    output_folder = os.path.join(cfg.OUTPUT_DIR, "inference", 'iter_{}'.format(iteration))

 
    for dataset_name in dataset_names_test:
        """
        Cycle through each dataset and test them individually.
        This loop keeps track of each per-image evaluation result, 
        so that it doesn't need to be re-computed for the collective.
        """

        '''
        Distributed Cube R-CNN inference
        '''
        dataset_paths = [os.path.join('datasets', 'Omni3D', name + '.json') for name in cfg.DATASETS.TEST]
        datasets = data.Omni3D(dataset_paths, filter_settings=filter_settings)

        # determine the meta data given the datasets used. 
        data.register_and_store_model_metadata(datasets, cfg.OUTPUT_DIR, filter_settings)

        thing_classes = MetadataCatalog.get('omni3d_model').thing_classes
        dataset_id_to_contiguous_id = MetadataCatalog.get('omni3d_model').thing_dataset_id_to_contiguous_id
        
        infos = datasets.dataset['info']

        if type(infos) == dict:
            infos = [datasets.dataset['info']]

        dataset_id_to_unknown_cats = {}
        possible_categories = set(i for i in range(cfg.MODEL.ROI_HEADS.NUM_CLASSES + 1))
        
        dataset_id_to_src = {}

        for info in infos:
            dataset_id = info['id']
            known_category_training_ids = set()

            if not dataset_id in dataset_id_to_src:
                dataset_id_to_src[dataset_id] = info['source']

            for id in info['known_category_ids']:
                if id in dataset_id_to_contiguous_id:
                    known_category_training_ids.add(dataset_id_to_contiguous_id[id])
            
            # determine and store the unknown categories.
            unknown_categories = possible_categories - known_category_training_ids
            dataset_id_to_unknown_cats[dataset_id] = unknown_categories


        # we need the dataset mapper to get 
        data_mapper = DatasetMapper3D(cfg, is_train=False, mode='get_depth_maps')
        data_mapper.dataset_id_to_unknown_cats = dataset_id_to_unknown_cats

        data_loader = build_detection_test_loader(cfg, dataset_name, mapper=data_mapper, batch_size=cfg.SOLVER.IMS_PER_BATCH, num_workers=4)

        experiment_type = {}

        if cfg.PLOT.EVAL == 'MABO': experiment_type['output_recall_scores'] = True
        else: experiment_type['output_recall_scores'] = False
        # either use pred_boxes or GT boxes
        if cfg.PLOT.MODE2D == 'PRED': experiment_type['use_pred_boxes'] = True
        else: experiment_type['use_pred_boxes'] = False
        if cfg.PLOT.SCORING_FUNC == False:
            experiment_type['scoring_func'] = False
        # define proposal function to use
        if experiment_type['output_recall_scores']:
            _ = mean_average_best_overlap(model, data_loader, experiment_type, cfg.PLOT.PROPOSAL_FUNC)
        elif not cfg.PLOT.SCORING_FUNC:
            _ = percent_of_boxes(model, data_loader, experiment_type, cfg.PLOT.PROPOSAL_FUNC)
        else:
            results_json = inference_on_dataset(model, data_loader, experiment_type, cfg.PLOT.PROPOSAL_FUNC)

            eval_helper = Omni3DEvaluationHelper(
                dataset_names_test, 
                filter_settings, 
                output_folder, 
                iter_label=iteration,
                only_2d=only_2d,
            )
            '''
            Individual dataset evaluation
            '''
            eval_helper.add_predictions(dataset_name, results_json)
            eval_helper.save_predictions(dataset_name)
            eval_helper.evaluate(dataset_name)

            '''
            Optionally, visualize some instances
            '''
            instances = torch.load(os.path.join(output_folder, dataset_name, 'instances_predictions.pth'))
            log_str = vis.visualize_from_instances(
                instances, data_loader.dataset, dataset_name, 
                cfg.INPUT.MIN_SIZE_TEST, os.path.join(output_folder, dataset_name), 
                MetadataCatalog.get('omni3d_model').thing_classes, iteration, visualize_every=1
            )
            logger.info(log_str)

        
    if cfg.PLOT.EVAL == 'AP':
        '''
        Summarize each Omni3D Evaluation metric
        '''  
        eval_helper.summarize_all()

def do_train(cfg, model):
    """
    Run model on the data_loader. 
    Also benchmark the inference speed of `model.__call__` accurately.
    The model will be used in train mode.

    Args:
        model (callable): a callable which takes an object from
            `data_loader` and returns some outputs.

            If it's an nn.Module, it will be temporarily set to `eval` mode.
            If you wish to evaluate a model in `training` mode instead, you can
            wrap the given model and override its behavior of `.eval()` and `.train()`.
        data_loader: an iterable object with a length.
            The elements it generates will be the inputs to the model.

    Returns:
        The return value of `evaluator.evaluate()`
    """

    filter_settings = data.get_filter_settings_from_cfg(cfg)

    # setup and join the data.
    dataset_paths = [os.path.join('datasets', 'Omni3D', name + '.json') for name in cfg.DATASETS.TRAIN]
    datasets = data.Omni3D(dataset_paths, filter_settings=filter_settings)

    # determine the meta data given the datasets used. 
    data.register_and_store_model_metadata(datasets, cfg.OUTPUT_DIR, filter_settings)

    thing_classes = MetadataCatalog.get('omni3d_model').thing_classes
    dataset_id_to_contiguous_id = MetadataCatalog.get('omni3d_model').thing_dataset_id_to_contiguous_id
    
    '''
    It may be useful to keep track of which categories are annotated/known
    for each dataset in use, in case a method wants to use this information.
    '''

    infos = datasets.dataset['info']

    if type(infos) == dict:
        infos = [datasets.dataset['info']]

    dataset_id_to_unknown_cats = {}
    possible_categories = set(i for i in range(cfg.MODEL.ROI_HEADS.NUM_CLASSES + 1))
    
    dataset_id_to_src = {}

    for info in infos:
        dataset_id = info['id']
        known_category_training_ids = set()

        if not dataset_id in dataset_id_to_src:
            dataset_id_to_src[dataset_id] = info['source']

        for id in info['known_category_ids']:
            if id in dataset_id_to_contiguous_id:
                known_category_training_ids.add(dataset_id_to_contiguous_id[id])
        
        # determine and store the unknown categories.
        unknown_categories = possible_categories - known_category_training_ids
        dataset_id_to_unknown_cats[dataset_id] = unknown_categories

    # we need the dataset mapper to get 
    dataset_names = cfg.DATASETS.TRAIN
    data_mapper = DatasetMapper3D(cfg, is_train=False, mode='get_depth_maps')
    data_mapper.dataset_id_to_unknown_cats = dataset_id_to_unknown_cats
    assert cfg.TRAIN.pseudo_gt in ['learn', 'pseudo'], "control what kind of proposal should be saved by setting TRAIN.pseudo_gt to either 'learn' or 'pseudo'"
    experiment_type = {}
    experiment_type['use_pred_boxes'] = cfg.PLOT.MODE2D if cfg.PLOT.MODE2D != '' else False
    experiment_type['pseudo_gt'] = cfg.TRAIN.pseudo_gt
    os.makedirs(f'datasets/proposals_{cfg.TRAIN.pseudo_gt}',exist_ok=True)

    # lol I think we have to hardcode this part in
    
    dataset_json = {}
    dataset_json.update({"categories": [{"supercategory": "nan", "id": 18, "name": "chair"}, {"supercategory": "nan", "id": 31, "name": "door"}, {"supercategory": "nan", "id": 37, "name": "table"}, {"supercategory": "nan", "id": 26, "name": "shelves"}, {"supercategory": "nan", "id": 51, "name": "kitchen pan"}, {"supercategory": "nan", "id": 52, "name": "bin"}, {"supercategory": "nan", "id": 38, "name": "counter"}, {"supercategory": "nan", "id": 29, "name": "cabinet"}, {"supercategory": "nan", "id": 53, "name": "stove"}, {"supercategory": "nan", "id": 28, "name": "sink"}, {"supercategory": "nan", "id": 14, "name": "books"}, {"supercategory": "nan", "id": 49, "name": "refrigerator"}, {"supercategory": "nan", "id": 54, "name": "microwave"}, {"supercategory": "nan", "id": 15, "name": "bottle"}, {"supercategory": "nan", "id": 55, "name": "plates"}, {"supercategory": "nan", "id": 56, "name": "bowl"}, {"supercategory": "nan", "id": 57, "name": "oven"}, {"supercategory": "nan", "id": 58, "name": "vase"}, {"supercategory": "nan", "id": 59, "name": "faucet"}, {"supercategory": "nan", "id": 22, "name": "towel"}, {"supercategory": "nan", "id": 60, "name": "tissues"}, {"supercategory": "nan", "id": 61, "name": "machine"}, {"supercategory": "nan", "id": 62, "name": "printer"}, {"supercategory": "nan", "id": 33, "name": "desk"}, {"supercategory": "nan", "id": 63, "name": "monitor"}, {"supercategory": "nan", "id": 64, "name": "podium"}, {"supercategory": "nan", "id": 35, "name": "bookcase"}, {"supercategory": "nan", "id": 41, "name": "dresser"}, {"supercategory": "nan", "id": 65, "name": "cart"}, {"supercategory": "nan", "id": 66, "name": "projector"}, {"supercategory": "nan", "id": 67, "name": "electronics"}, {"supercategory": "nan", "id": 68, "name": "computer"}, {"supercategory": "nan", "id": 34, "name": "box"}, {"supercategory": "nan", "id": 36, "name": "picture"}, {"supercategory": "nan", "id": 20, "name": "laptop"}, {"supercategory": "nan", "id": 42, "name": "pillow"}, {"supercategory": "nan", "id": 39, "name": "bed"}, {"supercategory": "nan", "id": 69, "name": "air conditioner"}, {"supercategory": "nan", "id": 25, "name": "lamp"}, {"supercategory": "nan", "id": 40, "name": "night stand"}, {"supercategory": "nan", "id": 50, "name": "board"}, {"supercategory": "nan", "id": 43, "name": "sofa"}, {"supercategory": "nan", "id": 71, "name": "coffee maker"}, {"supercategory": "nan", "id": 72, "name": "toaster"}, {"supercategory": "nan", "id": 73, "name": "potted plant"}, {"supercategory": "nan", "id": 48, "name": "stationery"}, {"supercategory": "nan", "id": 74, "name": "painting"}, {"supercategory": "nan", "id": 75, "name": "bag"}, {"supercategory": "nan", "id": 76, "name": "tray"}, {"supercategory": "nan", "id": 19, "name": "cup"}, {"supercategory": "nan", "id": 70, "name": "drawers"}, {"supercategory": "nan", "id": 77, "name": "keyboard"}, {"supercategory": "nan", "id": 21, "name": "shoes"}, {"supercategory": "vehicle & road", "id": 11, "name": "bicycle"}, {"supercategory": "nan", "id": 78, "name": "blanket"}, {"supercategory": "nan", "id": 44, "name": "television"}, {"supercategory": "nan", "id": 79, "name": "rack"}, {"supercategory": "nan", "id": 27, "name": "mirror"}, {"supercategory": "nan", "id": 47, "name": "clothes"}, {"supercategory": "nan", "id": 80, "name": "phone"}, {"supercategory": "nan", "id": 81, "name": "mouse"}, {"supercategory": "person", "id": 7, "name": "person"}, {"supercategory": "nan", "id": 82, "name": "fire extinguisher"}, {"supercategory": "nan", "id": 83, "name": "toys"}, {"supercategory": "nan", "id": 84, "name": "ladder"}, {"supercategory": "nan", "id": 85, "name": "fan"}, {"supercategory": "nan", "id": 32, "name": "toilet"}, {"supercategory": "nan", "id": 30, "name": "bathtub"}, {"supercategory": "nan", "id": 86, "name": "glass"}, {"supercategory": "nan", "id": 87, "name": "clock"}, {"supercategory": "nan", "id": 88, "name": "toilet paper"}, {"supercategory": "nan", "id": 89, "name": "closet"}, {"supercategory": "nan", "id": 46, "name": "curtain"}, {"supercategory": "nan", "id": 24, "name": "window"}, {"supercategory": "nan", "id": 90, "name": "fume hood"}, {"supercategory": "nan", "id": 91, "name": "utensils"}, {"supercategory": "nan", "id": 45, "name": "floor mat"}, {"supercategory": "nan", "id": 92, "name": "soundsystem"}, {"supercategory": "nan", "id": 93, "name": "fire place"}, {"supercategory": "nan", "id": 94, "name": "shower curtain"}, {"supercategory": "nan", "id": 23, "name": "blinds"}, {"supercategory": "nan", "id": 95, "name": "remote"}, {"supercategory": "nan", "id": 96, "name": "pen"}]})

    d_id_to_contiguous = MetadataCatalog.get('omni3d_model').thing_dataset_id_to_contiguous_id
    contiguous_to_id = {v:k for k,v in d_id_to_contiguous.items()}

    global_id = 1
    # this controls the flow of the program in the model class
    model.train()
    for dataset_name in dataset_names:
        idd = 12
        if 'val' in dataset_name:
            idd = 13
        dataset_json.update({"info": {"id": idd, "source": "SUNRGBD", "name": "SUNRGBD Train", "split": "Train", "version": "0.1", "url": "https://rgbd.cs.princeton.edu/"}})
        data_loader = build_detection_test_loader(cfg, dataset_name, mapper=data_mapper, num_workers=4)

        total = len(data_loader)  # inference data loader must have a fixed length

        annotations = []
        images = []
        for idx, inputs in tqdm(enumerate(data_loader), desc="Generating pseudo GT", total=total):
            cubes = model(inputs, experiment_type)
            instances = cubes[0]['instances']
            input_ = inputs[0]
            img_id = input_['image_id']
            input_['instances'].proposal_boxes = input_['instances'].gt_boxes
            bboxes = GeneralizedRCNN._postprocess([input_['instances']], [input_], [input_['instances']._image_size])
            bboxes = bboxes[0]['instances'].proposal_boxes
            # build json for each image
            img = {'width':input_['width'], 'height':input_['height'], 'file_path':input_['file_name'][9:], 'K':input_['K'], 'src_90_rotate':False, 'src_flagged':False, 'incomplete':False, 'id':img_id, 'dataset_id':12}
            for bbox, gt_class, center, dimensions, bbox3D, rotation in zip(bboxes, input_['instances'].gt_classes,
                                                                            instances.pred_center_cam.tolist(), instances.pred_dimensions.tolist(), 
                                                                            instances.pred_bbox3D.tolist(), instances.pred_pose.tolist()):
                c_id = util.MetadataCatalog.get('omni3d_model').thing_classes[gt_class]
                ann = {'behind_camera':False, 'truncation': 0, 'bbox2D_proj':bbox.tolist(), 'bbox2D_tight':-1, 'visibility':1.0, 'segmentation_pts':-1, 'lidar_pts':-1,\
                       'valid3D':True, 'category_id':contiguous_to_id[gt_class.tolist()], 'category_name':c_id, \
                       'id':global_id, 'image_id':img_id, 'dataset_id':idd, 'depth_error':-1, 'center_cam':center,\
                       'dimensions':dimensions, 'bbox3D_cam':bbox3D, 'R_cam':rotation}
                annotations.append(ann)
                global_id += 1

            images.append(img)

        dataset_json.update({'images':images, 'annotations':annotations})

        with open(f'datasets/Omni3D/SUNRGBD_pseudo_gt_{dataset_name}.json', 'w') as f:
            json.dump(dataset_json, f)

    return True


def setup(args):
    """
    Create configs and perform basic setups.
    """
    cfg = get_cfg()
    get_cfg_defaults(cfg)

    config_file = args.config_file
    
    # store locally if needed
    if config_file.startswith(util.CubeRCNNHandler.PREFIX):    
        config_file = util.CubeRCNNHandler._get_local_path(util.CubeRCNNHandler, config_file)

    cfg.merge_from_file(config_file)
    cfg.merge_from_list(args.opts)
    device = "cuda" if torch.cuda.is_available() else "cpu"
    cfg.MODEL.DEVICE = device
    cfg.SEED = 13
    cfg.freeze()
    default_setup(cfg, args)

    setup_logger(output=cfg.OUTPUT_DIR, distributed_rank=comm.get_rank(), name="cubercnn")
    
    filter_settings = data.get_filter_settings_from_cfg(cfg)

    for dataset_name in cfg.DATASETS.TRAIN:
        simple_register(dataset_name, filter_settings, filter_empty=True)
    
    dataset_names_test = cfg.DATASETS.TEST

    # filter_ = True if cfg.PLOT.EVAL == 'MABO' else False
    for dataset_name in dataset_names_test:
        if not(dataset_name in cfg.DATASETS.TRAIN):
            # empties should be filtering in test normally
            simple_register(dataset_name, filter_settings, filter_empty=True)
    
    return cfg


def main(args):
    
    cfg = setup(args)

    if args.eval_only:
        assert cfg.PLOT.MODE2D in ['GT', 'PRED'], 'MODE2D must be either GT or PRED'
        assert cfg.PLOT.EVAL in ['AP', 'MABO', 'IoU3D'], 'EVAL must be either AP, MABO or IoU3D'
        if cfg.PLOT.EVAL == 'MABO':
            assert cfg.PLOT.MODE2D == 'GT', 'MABO only works with GT boxes'
    
    name = f'cube {datetime.datetime.now().isoformat()}'
    # wandb.init(project="cube", sync_tensorboard=True, name=name, config=cfg)

    priors = None
    with open('tools/priors.pkl', 'rb') as f:
        priors, _ = pickle.load(f)

    category_path = 'output/Baseline_sgd/category_meta.json'
    # category_path = os.path.join(util.file_parts(args.opts[1])[0], 'category_meta.json')
    
    # store locally if needed
    if category_path.startswith(util.CubeRCNNHandler.PREFIX):
        category_path = util.CubeRCNNHandler._get_local_path(util.CubeRCNNHandler, category_path)

    metadata = util.load_json(category_path)

    # register the categories
    thing_classes = metadata['thing_classes']
    id_map = {int(key):val for key, val in metadata['thing_dataset_id_to_contiguous_id'].items()}
    MetadataCatalog.get('omni3d_model').thing_classes = thing_classes
    MetadataCatalog.get('omni3d_model').thing_dataset_id_to_contiguous_id  = id_map


    # build the  model.
    model = build_model(cfg, priors=priors)


    if args.eval_only:
        # skip straight to eval mode
        # load the saved model if using eval boxes
        if cfg.PLOT.MODE2D == 'PRED':
            DetectionCheckpointer(model, save_dir=cfg.OUTPUT_DIR).resume_or_load(
                cfg.MODEL.WEIGHTS, resume=False)
        return do_test(cfg, model)
    else:
        logger.info('Making pseudo GT')
        return do_train(cfg, model)


if __name__ == "__main__":
    args = default_argument_parser().parse_args()
    print("Command Line Args:", args)

    main(args)