3D-OWIS is capable of detecting unknown instances in inference, and progressively learning novel classes in the process of training.
# Evaluates semantic instance task
# Adapted from the CityScapes evaluation: https://github.com/mcordts/cityscapesScripts/tree/master/cityscapesscripts/evaluation
# Input:
#   - path to .txt prediction files
#   - path to .txt ground truth files
#   - output file to write results to
# Each .txt prediction file look like:
#    [(pred0) rel. path to pred. mask over verts as .txt] [(pred0) label id] [(pred0) confidence]
#    [(pred1) rel. path to pred. mask over verts as .txt] [(pred1) label id] [(pred1) confidence]
#    [(pred2) rel. path to pred. mask over verts as .txt] [(pred2) label id] [(pred2) confidence]
#    ...
#
# NOTE: The prediction files must live in the root of the given prediction path.
#       Predicted mask .txt files must live in a subfolder.
#       Additionally, filenames must not contain spaces.
# The relative paths to predicted masks must contain one integer per line,
# where each line corresponds to vertices in the *_vh_clean_2.ply (in that order).
# Non-zero integers indicate part of the predicted instance.
# The label ids specify the class of the corresponding mask.
# Confidence is a float confidence score of the mask.
#
# Note that only the valid classes are used for evaluation,
# i.e., any ground truth label not in the valid label set
# is ignored in the evaluation.
#
# example usage: evaluate_semantic_instance.py --scan_path [path to scan data] --output_file [output file]

# python imports
import math
import os, sys, argparse
import inspect
from copy import deepcopy
from uuid import uuid4
import yaml
import torch
import json
try:
    import numpy as np
except:
    print("Failed to import numpy package.")
    sys.exit(-1)

from scipy import stats
import csv
#currentdir = os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe())))
#parentdir = os.path.dirname(currentdir)
#sys.path.insert(0,parentdir)
import benchmark.util as util
import benchmark.util_3d as util_3d

from datasets.scannet200.scannet200_splits import HEAD_CATS_SCANNET_200
from datasets.scannet200.scannet200_constants import VALID_CLASS_IDS_200, CLASS_LABELS_200
from datasets.scannet200.owis_splits import KNOWN_CLASSES_LABELS, PREV_KNOWN_CLASSES_LABELS
#parser = argparse.ArgumentParser()
#parser.add_argument('--gt_path', default='', help='path to directory of gt .txt files')
#parser.add_argument('--output_file', default='', help='output file [default: ./semantic_instance_evaluation.txt]')
#opt = parser.parse_args()

#if opt.output_file == '':
#    opt.output_file = os.path.join(os.getcwd(), 'semantic_instance_evaluation.txt')

# ---------- Label info ---------- #
CLASS_LABELS = ['cabinet', 'bed', 'chair', 'sofa', 'table', 'door', 'window', 'bookshelf', 'picture', 'counter', 'desk', 'curtain', 'refrigerator', 'shower curtain', 'toilet', 'sink', 'bathtub', 'otherfurniture']
VALID_CLASS_IDS = np.array([3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 14, 16, 24, 28, 33, 34, 36, 39])
ID_TO_LABEL = {}
LABEL_TO_ID = {}
for i in range(len(VALID_CLASS_IDS)):
    LABEL_TO_ID[CLASS_LABELS[i]] = VALID_CLASS_IDS[i]
    ID_TO_LABEL[VALID_CLASS_IDS[i]] = CLASS_LABELS[i]
    
ID_TO_LABEL_200 = {}
LABEL_TO_ID_200 = {}
for i in range(len(VALID_CLASS_IDS_200)):
    LABEL_TO_ID_200[CLASS_LABELS_200[i]] = VALID_CLASS_IDS_200[i]
    ID_TO_LABEL_200[VALID_CLASS_IDS_200[i]] = CLASS_LABELS_200[i]
# ---------- Evaluation params ---------- #


# overlaps for evaluation
opt = {}
opt['overlaps']             = np.append(np.arange(0.5,0.95,0.05), 0.25)
# minimum region size for evaluation [verts]
opt['min_region_sizes']     = np.array( [ 100 ] ) # 100 for s3dis, scannet
# distance thresholds [m]
opt['distance_threshes']    = np.array( [  float('inf') ] )
# distance confidences
opt['distance_confs']       = np.array( [ -float('inf') ] )


def evaluate_matches(matches):
    overlaps = opt['overlaps']
    min_region_sizes = [ opt['min_region_sizes'][0] ]
    dist_threshes = [ opt['distance_threshes'][0] ]
    dist_confs = [ opt['distance_confs'][0] ]
    # results: class x overlap
    ap = np.zeros( (len(dist_threshes) , len(CLASS_LABELS) , len(overlaps)) , float )
    U_Recall = {}
    WI = dict(zip([round(ov,2) for ov in overlaps],len(overlaps)*[0.0]))
    A_OSE = dict(zip([round(ov,2) for ov in overlaps],len(overlaps)*[0.0]))
    FPo = dict(zip([round(ov,2) for ov in overlaps],len(overlaps)*[0.0]))
    FPc = dict(zip([round(ov,2) for ov in overlaps],len(overlaps)*[0.0]))
    TPc = dict(zip([round(ov,2) for ov in overlaps],len(overlaps)*[0.0]))
    
    for di, (min_region_size, distance_thresh, distance_conf) in enumerate(zip(min_region_sizes, dist_threshes, dist_confs)):
        for oi, overlap_th in enumerate(overlaps):
            pred_visited = {}
            for m in matches:
                for p in matches[m]['pred']:
                    for label_name in CLASS_LABELS:
                        for p in matches[m]['pred'][label_name]:
                            if 'uuid' in p:
                                pred_visited[p['uuid']] = False
            
            
            for li, label_name in enumerate(CLASS_LABELS):
                y_true = np.empty(0)
                y_score = np.empty(0)
                
                # if save_WI:
                #     y_true_wo_unkn = np.empty(0)
                #     y_score_wo_unkn = np.empty(0)
                    
                hard_false_negatives = 0
                has_gt = False
                has_pred = False
                dic = {}
                
                for m in matches:
                    pred_instances = matches[m]['pred'][label_name]
                    gt_instances = matches[m]['gt'][label_name]
                    # filter groups in ground truth
                    gt_instances = [ gt for gt in gt_instances if gt['instance_id']>=1000 and gt['vert_count']>=min_region_size and gt['med_dist']<=distance_thresh and gt['dist_conf']>=distance_conf ]
                    if gt_instances:
                        has_gt = True
                    if pred_instances:
                        has_pred = True

                    cur_true  = np.ones ( len(gt_instances) )
                    cur_score = np.ones ( len(gt_instances) ) * (-float("inf"))
                    cur_match = np.zeros( len(gt_instances) , dtype=bool )
                    
                    # if save_WI:
                    #     cur_true_wo_unkn  = np.ones ( len(gt_instances) )
                    #     cur_score_wo_unkn = np.ones ( len(gt_instances) ) * (-float("inf"))
                    #     cur_match_wo_unkn = np.zeros( len(gt_instances) , dtype=bool )
                    # collect matches
                    for (gti,gt) in enumerate(gt_instances):
                        found_match = False
                        num_pred = len(gt['matched_pred'])
                        for pred in gt['matched_pred']:
                            # greedy assignments
                            if pred_visited[pred['uuid']]:
                                continue
                            overlap = float(pred['intersection']) / (gt['vert_count']+pred['vert_count']-pred['intersection'])
                            if overlap > overlap_th:
                                confidence = pred['confidence']
                                # if already have a prediction for this gt,
                                # the prediction with the lower score is automatically a false positive
                                if cur_match[gti]:
                                    max_score = max( cur_score[gti] , confidence )
                                    min_score = min( cur_score[gti] , confidence )
                                    cur_score[gti] = max_score
                                    # append false positive with unknown
                                    cur_true  = np.append(cur_true,0)
                                    cur_score = np.append(cur_score,min_score)
                                    cur_match = np.append(cur_match,True)
                                    if pred['pred_id'] != 3000 and gt['label_id'] == 3000:
                                        A_OSE[round(overlap_th,2)] += 1
                                    if pred['pred_id'] != 3000 and gt['label_id'] == 3000 and pred['confidence']>0.5:
                                        FPo[round(overlap_th,2)] += 1
                                    if pred['pred_id'] != 3000 and gt['label_id'] != 3000 and gt['label_id'] != pred['pred_id'] and pred['confidence']>0.5:
                                        FPc[round(overlap_th,2)] += 1
                                    if pred['pred_id'] == gt['label_id'] and gt['label_id'] != 3000 and pred['confidence']>0.5:
                                        TPc[round(overlap_th,2)] += 1
                                    # if pred['pred_id'] != 3000 and save_WI:
                                    #     cur_true_wo_unkn  = np.append(cur_true_wo_unkn,0)
                                    #     cur_score_wo_unkn = np.append(cur_score_wo_unkn,min_score)
                                    #     cur_match_wo_unkn = np.append(cur_match_wo_unkn,True)
                                    
                                # otherwise set score
                                else:
                                    found_match = True
                                    cur_match[gti] = True
                                    cur_score[gti] = confidence
                                    pred_visited[pred['uuid']] = True
                                    # if pred['pred_id'] != 3000 and save_WI:
                                    #     cur_match_wo_unkn[gti] = True
                                    #     cur_score_wo_unkn[gti] = confidence

                        if not found_match:
                            hard_false_negatives += 1
                    # remove non-matched ground truth instances
                    
                    cur_true  = cur_true [ cur_match==True ]
                    cur_score = cur_score[ cur_match==True ]
                    
                    # if save_WI:
                    #     cur_true_wo_unkn  = cur_true_wo_unkn [ cur_match_wo_unkn==True ]
                    #     cur_score_wo_unkn = cur_score_wo_unkn[ cur_match_wo_unkn==True ]
                    
                    # collect non-matched predictions as false positive
                    for pred in pred_instances:
                        found_gt = False
                        for gt in pred['matched_gt']:
                            overlap = float(gt['intersection']) / (gt['vert_count']+pred['vert_count']-gt['intersection'])
                            if overlap > overlap_th:
                                found_gt = True
                                break
                        if not found_gt:
                            num_ignore = pred['void_intersection']
                            for gt in pred['matched_gt']:
                                # group?
                                if gt['instance_id'] < 1000:
                                    num_ignore += gt['intersection']
                                # small ground truth instances
                                if gt['vert_count'] < min_region_size or gt['med_dist']>distance_thresh or gt['dist_conf']<distance_conf:
                                    num_ignore += gt['intersection']
                            proportion_ignore = float(num_ignore)/pred['vert_count']
                            # if not ignored append false positive
                            if proportion_ignore <= overlap_th:
                                cur_true = np.append(cur_true,0)
                                confidence = pred["confidence"]
                                cur_score = np.append(cur_score,confidence)
                                
                                # if pred['pred_id'] != 3000 and save_WI:
                                #         cur_true_wo_unkn  = np.append(cur_true_wo_unkn,0)
                                #         confidence_wo_unkn = pred["confidence"]
                                #         cur_score_wo_unkn = np.append(cur_score_wo_unkn,confidence_wo_unkn)

                    # append to overall results
                    y_true  = np.append(y_true,cur_true)
                    y_score = np.append(y_score,cur_score)
                    
                    # if save_WI:
                    #     y_true_wo_unkn  = np.append(y_true_wo_unkn,cur_true_wo_unkn)
                    #     y_score_wo_unkn = np.append(y_score_wo_unkn,cur_score_wo_unkn)

                # compute average precision
                if has_gt and has_pred:
                    # compute precision recall curve first

                    # sorting and cumsum
                    score_arg_sort      = np.argsort(y_score)
                    y_score_sorted      = y_score[score_arg_sort]
                    y_true_sorted       = y_true[score_arg_sort]
                    y_true_sorted_cumsum = np.cumsum(y_true_sorted)

                    # unique thresholds
                    (thresholds,unique_indices) = np.unique( y_score_sorted , return_index=True )
                    num_prec_recall = len(unique_indices) + 1
                    
                    

                    # prepare precision recall
                    num_examples      = len(y_score_sorted)
                    
                    
                    # https://github.com/ScanNet/ScanNet/pull/26
                    # all predictions are non-matched but also all of them are ignored and not counted as FP
                    # y_true_sorted_cumsum is empty
                    # num_true_examples = y_true_sorted_cumsum[-1]
                    num_true_examples = y_true_sorted_cumsum[-1] if len(y_true_sorted_cumsum) > 0 else 0
                    precision         = np.zeros(num_prec_recall)
                    recall            = np.zeros(num_prec_recall) 

                    # deal with the first point
                    y_true_sorted_cumsum = np.append( y_true_sorted_cumsum , 0 )
                    # deal with remaining
                    for idx_res,idx_scores in enumerate(unique_indices):
                        cumsum = y_true_sorted_cumsum[idx_scores-1]
                        tp = num_true_examples - cumsum
                        fp = num_examples      - idx_scores - tp
                        fn = cumsum + hard_false_negatives
                        p  = float(tp)/(tp+fp)
                        r  = float(tp)/(tp+fn)
                        precision[idx_res] = p
                        recall   [idx_res] = r
                    
                    #Getting the Urecall
                    if label_name == "unknown":
                        U_Recall[int(overlap_th*100)] = recall[0]            
                    # first point in curve is artificial
                    precision[-1] = 1.
                    recall   [-1] = 0.
                    
                        
                    # compute average of precision-recall curve
                    recall_for_conv = np.copy(recall)
                    recall_for_conv = np.append(recall_for_conv[0], recall_for_conv)
                    recall_for_conv = np.append(recall_for_conv, 0.)
                    
                    
                    stepWidths = np.convolve(recall_for_conv,[-0.5,0,0.5],'valid')
                    # integrate is now simply a dot product
                    ap_current = np.dot(precision, stepWidths)

                elif has_gt:
                    ap_current = 0.0
                else:
                    ap_current = float('nan')
                ap[di,li,oi] = ap_current
            
            try:
                WI[round(overlap_th,2)] = FPo[round(overlap_th,2)]/(TPc[round(overlap_th,2)]+FPc[round(overlap_th,2)])
            except:
                pass
            
    
    return U_Recall,WI,A_OSE, ap

def compute_averages(aps):
    d_inf = 0
    o50   = np.where(np.isclose(opt['overlaps'],0.5))
    o25   = np.where(np.isclose(opt['overlaps'],0.25))
    oAllBut25  = np.where(np.logical_not(np.isclose(opt['overlaps'],0.25)))
    avg_dict = {}
    #avg_dict['all_ap']     = np.nanmean(aps[ d_inf,:,:  ])
    avg_dict['all_ap']     = np.nanmean(aps[ d_inf,:,oAllBut25])
    avg_dict['all_ap_50%'] = np.nanmean(aps[ d_inf,:,o50])
    avg_dict['all_ap_25%'] = np.nanmean(aps[ d_inf,:,o25])
    avg_dict["classes"]  = {}
    for (li,label_name) in enumerate(CLASS_LABELS):
        avg_dict["classes"][label_name]             = {}
        #avg_dict["classes"][label_name]["ap"]       = np.average(aps[ d_inf,li,  :])
        avg_dict["classes"][label_name]["ap"]       = np.average(aps[ d_inf,li,oAllBut25])
        avg_dict["classes"][label_name]["ap50%"]    = np.average(aps[ d_inf,li,o50])
        avg_dict["classes"][label_name]["ap25%"]    = np.average(aps[ d_inf,li,o25])
    return avg_dict

def make_pred_info(pred: dict):
    # pred = {'pred_scores' = 100, 'pred_classes' = 100 'pred_masks' = Nx100}
    pred_info = {}
    assert(pred['pred_classes'].shape[0] == pred['pred_scores'].shape[0] == pred['pred_masks'].shape[1])
    for i in range(len(pred['pred_classes'])):
        info = {}
        info["label_id"] = pred['pred_classes'][i]
        info["conf"] = pred['pred_scores'][i]#<=================================
        info["mask"] = pred['pred_masks'][:,i]
        pred_info[uuid4()] = info # we later need to identify these objects
    return pred_info

def assign_instances_for_scan(split,task, pred: dict, gt_file: str):
    pred_info = make_pred_info(pred)
    try:
        gt_ids = util_3d.load_ids(gt_file)
    except Exception as e:
        util.print_error('unable to load ' + gt_file + ': ' + str(e))

    # get gt instances
    gt_instances = util_3d.get_instances(split,task, gt_ids, VALID_CLASS_IDS, CLASS_LABELS, ID_TO_LABEL)
    # associate
    gt2pred = deepcopy(gt_instances)
    for label in gt2pred: 
        for gt in gt2pred[label]:
            gt['matched_pred'] = []
    pred2gt = {}
    for label in CLASS_LABELS:
        pred2gt[label] = []
    num_pred_instances = 0
    # mask of void labels in the groundtruth
    bool_void = np.logical_not(np.in1d(gt_ids//1000, VALID_CLASS_IDS))
    # go thru all prediction masks
    for uuid in pred_info:
        label_id = int(pred_info[uuid]['label_id'])
        conf = pred_info[uuid]['conf']
        if not label_id in ID_TO_LABEL:
            continue
        label_name = ID_TO_LABEL[label_id]
        # read the mask
        pred_mask = pred_info[uuid]['mask']
        assert(len(pred_mask) == len(gt_ids))
        # convert to binary
        pred_mask = np.not_equal(pred_mask, 0)
        num = np.count_nonzero(pred_mask)
        if num < opt['min_region_sizes'][0]:
            continue  # skip if empty

        pred_instance = {}
        pred_instance['uuid'] = uuid
        pred_instance['pred_id'] = num_pred_instances
        pred_instance['label_id'] = label_id
        pred_instance['vert_count'] = num
        pred_instance['confidence'] = conf
        pred_instance['void_intersection'] = np.count_nonzero(np.logical_and(bool_void, pred_mask))

        # matched gt instances
        matched_gt = []
        # go thru all gt instances with matching label
        for (gt_num, gt_inst) in enumerate(gt2pred[label_name]):
            intersection = np.count_nonzero(np.logical_and(gt_ids == gt_inst['instance_id'], pred_mask))
            if intersection > 0:
                gt_copy = gt_inst.copy()
                pred_copy = pred_instance.copy()
                gt_copy['intersection']   = intersection
                pred_copy['intersection'] = intersection
                matched_gt.append(gt_copy)
                gt2pred[label_name][gt_num]['matched_pred'].append(pred_copy)
        pred_instance['matched_gt'] = matched_gt
        num_pred_instances += 1
        pred2gt[label_name].append(pred_instance)

    return gt2pred, pred2gt


def print_results(avgs,current_epoch,save_dir,split,task_ow, u_recall, wi, a_ose):
    prev_known_avg = []
    prev_known_avg_25 = []
    prev_known_avg_50 = []
    curr_known_avg = []
    curr_known_avg_25 = []
    curr_known_avg_50 = []
    sep     = ""
    col1    = ":"
    lineLen = 64

    print("")
    print("#"*lineLen)
    line  = ""
    line += "{:<15}".format("what"      ) + sep + col1
    line += "{:>15}".format("AP"        ) + sep
    line += "{:>15}".format("AP_50%"    ) + sep
    line += "{:>15}".format("AP_25%"    ) + sep
    print(line)
    print("#"*lineLen)

    for (li,label_name) in enumerate(CLASS_LABELS):
        ap_avg  = avgs["classes"][label_name]["ap"]
        ap_50o  = avgs["classes"][label_name]["ap50%"]
        ap_25o  = avgs["classes"][label_name]["ap25%"]
        line  = "{:<15}".format(label_name) + sep + col1
        line += sep + "{:>15.3f}".format(ap_avg ) + sep
        line += sep + "{:>15.3f}".format(ap_50o ) + sep
        line += sep + "{:>15.3f}".format(ap_25o ) + sep
        print(line)
        if PREV_KNOWN_CLASSES_LABELS[split][task_ow] != None and label_name in PREV_KNOWN_CLASSES_LABELS[split][task_ow]:
            prev_known_avg.append(avgs["classes"][label_name]["ap"])
            prev_known_avg_25.append(avgs["classes"][label_name]["ap50%"])
            prev_known_avg_50.append(avgs["classes"][label_name]["ap25%"])
        
        if label_name in KNOWN_CLASSES_LABELS[split][task_ow]:
            curr_known_avg.append(avgs["classes"][label_name]["ap"])
            curr_known_avg_25.append(avgs["classes"][label_name]["ap50%"])
            curr_known_avg_50.append(avgs["classes"][label_name]["ap25%"])

    all_ap_avg  = avgs["all_ap"]
    all_ap_50o  = avgs["all_ap_50%"]
    all_ap_25o  = avgs["all_ap_25%"]

    print("-"*lineLen)
    line  = "{:<15}".format("Average") + sep + col1
    line += "{:>15.3f}".format(all_ap_avg)  + sep
    line += "{:>15.3f}".format(all_ap_50o)  + sep
    line += "{:>15.3f}".format(all_ap_25o)  + sep
    print(line)
    print("")
    
    if PREV_KNOWN_CLASSES_LABELS[split][task_ow] != None:
        line  = "{:<15}".format("Previous known") + sep + col1
        line += "{:>15.3f}".format(np.nanmean(np.array(prev_known_avg)))  + sep
        line += "{:>15.3f}".format(np.nanmean(np.array(prev_known_avg_25)))  + sep
        line += "{:>15.3f}".format(np.nanmean(np.array(prev_known_avg_50)))  + sep
        print(line)
        print("")
    
    line  = "{:<15}".format("Current known") + sep + col1
    line += "{:>15.3f}".format(np.nanmean(np.array(curr_known_avg)))  + sep
    line += "{:>15.3f}".format(np.nanmean(np.array(curr_known_avg_25)))  + sep
    line += "{:>15.3f}".format(np.nanmean(np.array(curr_known_avg_50)))  + sep
    print(line)
    print("")
    
    print("")
    print("#"*lineLen)
    line  = ""
    line += "{:<15}".format(""      ) + sep
    line += "{:>30}".format("Unknown class evaluation metrics"        ) + sep
    line += "{:>15}".format(""    ) + sep
    print(line)
    print("#"*lineLen)
    
    line  = "{:<21}".format("U-Recall") + sep + col1
    line += "{:>13}".format("@0.5"        ) + sep
    line += "{:>13}".format("@0.8"    ) + sep
    line += "{:>13}".format("@0.9"    ) + sep
    print(line)
    if type(u_recall).__name__ != 'float':
        line  = "{:<21}".format("") + sep + col1
        line += sep + "{:>13.5f}".format(u_recall[50] ) + sep
        line += sep + "{:>13.5f}".format(u_recall[80]) + sep
        line += sep + "{:>13.5f}".format(u_recall[90]) + sep
    else: 
        line  = "{:<21}".format("") + sep + col1
        line += sep + "{:>13.5f}".format(u_recall) + sep
        line += sep + "{:>13.5f}".format(u_recall) + sep
        line += sep + "{:>13.5f}".format(u_recall) + sep
    
    print(line)
    # line += "{:>12.5f}".format(u_recall)  + sep
    # print(line)
    print("")
    print("-"*lineLen)
    line  = "{:<21}".format("Wilderness Impact") + sep + col1
    line += "{:>13.5f}".format(wi)  + sep
    print(line)
    print("")
    
    line  = "{:<21}".format("A-OSE") + sep + col1
    line += "{:>13d}".format(int(a_ose))  + sep
    print(line)
    print("")
    
    
    results_dict = {
    'ap': all_ap_avg.item(),
    'current_known_ap': np.nanmean(np.array(curr_known_avg)).item(),
    'WI': wi,
    'A-OSE': int(a_ose)
    }
    if type(u_recall).__name__ != 'float':
        results_dict['U_RECALL']=u_recall[50].item()
    else:
        results_dict['U_RECALL']=u_recall
        
    if PREV_KNOWN_CLASSES_LABELS[split][task_ow] != None:
        results_dict['previously_known_ap']=np.nanmean(np.array(prev_known_avg)).item()
        
    results_path = os.path.join(save_dir, 'ow_results')
    
    if not os.path.exists(results_path):
        os.makedirs(results_path)

    suff = str(current_epoch)+'_ow_results.yml'

    with open(os.path.join(results_path,suff), 'w') as outfile:
        yaml.dump(results_dict, outfile, default_flow_style=False)


def write_result_file(avgs, filename):
    _SPLITTER = ','
    with open(filename, 'w') as f:
        f.write(_SPLITTER.join(['class', 'class id', 'ap', 'ap50', 'ap25']) + '\n')
        for i in range(len(VALID_CLASS_IDS)):
            class_name = CLASS_LABELS[i]
            class_id = VALID_CLASS_IDS[i]
            ap = avgs["classes"][class_name]["ap"]
            ap50 = avgs["classes"][class_name]["ap50%"]
            ap25 = avgs["classes"][class_name]["ap25%"]
            f.write(_SPLITTER.join([str(x) for x in [class_name, class_id, ap, ap50, ap25]]) + '\n')


def evaluate(current_epoch,config_general, preds: dict, gt_path: str, output_file: str, dataset: str = "scannet"):
    global CLASS_LABELS
    global VALID_CLASS_IDS
    global ID_TO_LABEL
    global LABEL_TO_ID
    global opt

    if dataset == "stpls3d":
        # global CLASS_LABELS
        # global VALID_CLASS_IDS
        # global ID_TO_LABEL
        # global LABEL_TO_ID

        opt['min_region_sizes'] = np.array([10])

        CLASS_LABELS = ['Build', 'LowVeg', 'MediumVeg', 'HighVeg', 'Vehicle', 'Truck', 'Aircraft', 'MilitaryVeh',
                        'Bike', 'Motorcycle', 'LightPole', 'StreetSign', 'Clutter', 'Fence']
        VALID_CLASS_IDS = np.array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14])

        ID_TO_LABEL = {}
        LABEL_TO_ID = {}
        for i in range(len(VALID_CLASS_IDS)):
            LABEL_TO_ID[CLASS_LABELS[i]] = VALID_CLASS_IDS[i]
            ID_TO_LABEL[VALID_CLASS_IDS[i]] = CLASS_LABELS[i]

    if dataset == "s3dis":
        # global CLASS_LABELS
        # global VALID_CLASS_IDS
        # global ID_TO_LABEL
        # global LABEL_TO_ID

        CLASS_LABELS = ['ceiling', 'floor', 'wall', 'beam', 'column', 'window', 'door', 'table',
                        'chair', 'sofa', 'bookcase', 'board', 'clutter']
        VALID_CLASS_IDS = np.array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13])
        ID_TO_LABEL = {}
        LABEL_TO_ID = {}
        for i in range(len(VALID_CLASS_IDS)):
            LABEL_TO_ID[CLASS_LABELS[i]] = VALID_CLASS_IDS[i]
            ID_TO_LABEL[VALID_CLASS_IDS[i]] = CLASS_LABELS[i]

    if dataset == "scannet200":
#         CLASS_LABELS = (
# 'chair', 'table', 'door', 'couch', 'cabinet', 'shelf', 'desk', 'office chair', 'bed', 'pillow', 'sink', 'picture', 'window', 'toilet', 'bookshelf', 'monitor', 'curtain', 'book', 'armchair', 'coffee table', 'box',
# 'refrigerator', 'lamp', 'kitchen cabinet', 'towel', 'clothes', 'tv', 'nightstand', 'counter', 'dresser', 'stool', 'cushion', 'plant', 'ceiling', 'bathtub', 'end table', 'dining table', 'keyboard', 'bag', 'backpack', 'toilet paper',
# 'printer', 'tv stand', 'whiteboard', 'blanket', 'shower curtain', 'trash can', 'closet', 'stairs', 'microwave', 'stove', 'shoe', 'computer tower', 'bottle', 'bin', 'ottoman', 'bench', 'board', 'washing machine', 'mirror', 'copier',
# 'basket', 'sofa chair', 'file cabinet', 'fan', 'laptop', 'shower', 'paper', 'person', 'paper towel dispenser', 'oven', 'blinds', 'rack', 'plate', 'blackboard', 'piano', 'suitcase', 'rail', 'radiator', 'recycling bin', 'container',
# 'wardrobe', 'soap dispenser', 'telephone', 'bucket', 'clock', 'stand', 'light', 'laundry basket', 'pipe', 'clothes dryer', 'guitar', 'toilet paper holder', 'seat', 'speaker', 'column', 'bicycle', 'ladder', 'bathroom stall', 'shower wall',
# 'cup', 'jacket', 'storage bin', 'coffee maker', 'dishwasher', 'paper towel roll', 'machine', 'mat', 'windowsill', 'bar', 'toaster', 'bulletin board', 'ironing board', 'fireplace', 'soap dish', 'kitchen counter', 'doorframe',
# 'toilet paper dispenser', 'mini fridge', 'fire extinguisher', 'ball', 'hat', 'shower curtain rod', 'water cooler', 'paper cutter', 'tray', 'shower door', 'pillar', 'ledge', 'toaster oven', 'mouse', 'toilet seat cover dispenser',
# 'furniture', 'cart', 'storage container', 'scale', 'tissue box', 'light switch', 'crate', 'power outlet', 'decoration', 'sign', 'projector', 'closet door', 'vacuum cleaner', 'candle', 'plunger', 'stuffed animal', 'headphones', 'dish rack',
# 'broom', 'guitar case', 'range hood', 'dustpan', 'hair dryer', 'water bottle', 'handicap bar', 'purse', 'vent', 'shower floor', 'water pitcher', 'mailbox', 'bowl', 'paper bag', 'alarm clock', 'music stand', 'projector screen', 'divider',
# 'laundry detergent', 'bathroom counter', 'object', 'bathroom vanity', 'closet wall', 'laundry hamper', 'bathroom stall door', 'ceiling light', 'trash bin', 'dumbbell', 'stair rail', 'tube', 'bathroom cabinet', 'cd case', 'closet rod',
# 'coffee kettle', 'structure', 'shower head', 'keyboard piano', 'case of water bottles', 'coat rack', 'storage organizer', 'folded chair', 'fire alarm', 'power strip', 'calendar', 'poster', 'potted plant', 'luggage', 'mattress', 'unknown')

#         VALID_CLASS_IDS = np.array((2, 4, 5, 6, 7, 8, 9, 10, 11, 13, 14, 15, 16, 17, 18, 19, 21, 22, 23, 24, 26, 27, 28, 29, 31, 32, 33, 34, 35, 36, 38, 39, 40, 41, 42, 44, 45, 46, 47, 48, 49, 50, 51, 52, 54, 55, 56, 57, 58, 59, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71,
# 72, 73, 74, 75, 76, 77, 78, 79, 80, 82, 84, 86, 87, 88, 89, 90, 93, 95, 96, 97, 98, 99, 100, 101, 102, 103, 104, 105, 106, 107, 110, 112, 115, 116, 118, 120, 121, 122, 125, 128, 130, 131, 132, 134, 136, 138, 139, 140, 141, 145, 148, 154,
# 155, 156, 157, 159, 161, 163, 165, 166, 168, 169, 170, 177, 180, 185, 188, 191, 193, 195, 202, 208, 213, 214, 221, 229, 230, 232, 233, 242, 250, 261, 264, 276, 283, 286, 300, 304, 312, 323, 325, 331, 342, 356, 370, 392, 395, 399, 408, 417,
# 488, 540, 562, 570, 572, 581, 609, 748, 776, 1156, 1163, 1164, 1165, 1166, 1167, 1168, 1169, 1170, 1171, 1172, 1173, 1174, 1175, 1176, 1178, 1179, 1180, 1181, 1182, 1183, 1184, 1185, 1186, 1187, 1188, 1189, 1190, 1191,3000))
        VALID_CLASS_IDS = list(VALID_CLASS_IDS_200) 
        CLASS_LABELS = list(CLASS_LABELS_200)
        CLASS_LABELS.remove('wall')
        CLASS_LABELS.remove('floor')
        VALID_CLASS_IDS.remove(1)
        VALID_CLASS_IDS.remove(3)
        VALID_CLASS_IDS = np.array(VALID_CLASS_IDS)
        CLASS_LABELS = np.array(CLASS_LABELS)
        
        ID_TO_LABEL = {}
        LABEL_TO_ID = {}
        for i in range(len(VALID_CLASS_IDS)):
            LABEL_TO_ID[CLASS_LABELS[i]] = VALID_CLASS_IDS[i]
            ID_TO_LABEL[VALID_CLASS_IDS[i]] = CLASS_LABELS[i]

    total_true = 0
    total_seen = 0
    NUM_CLASSES = len(VALID_CLASS_IDS)

    true_positive_classes = np.zeros(NUM_CLASSES)
    positive_classes = np.zeros(NUM_CLASSES)
    gt_classes = np.zeros(NUM_CLASSES)

    # precision & recall
    total_gt_ins = np.zeros(NUM_CLASSES)
    at = 0.5
    tpsins = [[] for _ in range(NUM_CLASSES)]
    fpsins = [[] for _ in range(NUM_CLASSES)]
    # mucov and mwcov
    all_mean_cov = [[] for _ in range(NUM_CLASSES)]
    all_mean_weighted_cov = [[] for _ in range(NUM_CLASSES)]

    print('evaluating', len(preds), 'scans...')
    matches = {}
    for i,(k,v) in enumerate(preds.items()):
        gt_file = os.path.join(gt_path, k + ".txt")
        if not os.path.isfile(gt_file):
            util.print_error('Scan {} does not match any gt file'.format(k), user_fault=True)

        if dataset == "s3dis":
            gt_ids = util_3d.load_ids(gt_file)
            gt_sem = (gt_ids // 1000) - 1
            gt_ins = gt_ids - (gt_ids // 1000) * 1000

            # pred_sem = v['pred_classes'] - 1
            pred_sem = np.zeros(v['pred_masks'].shape[0], dtype=np.int)
            # TODO CONTINUE HERE!!!!!!!!!!!!!
            pred_ins = np.zeros(v['pred_masks'].shape[0], dtype=np.int)


            for inst_id in reversed(range(v['pred_masks'].shape[1])):
                point_ids = np.argwhere(v['pred_masks'][:, inst_id] == 1.)[:, 0]
                pred_ins[point_ids] = inst_id + 1
                pred_sem[point_ids] = v['pred_classes'][inst_id] - 1

            # semantic acc
            total_true += np.sum(pred_sem == gt_sem)
            total_seen += pred_sem.shape[0]


            # TODO PARALLELIZ THIS!!!!!!!
            # pn semantic mIoU
            '''
            for j in range(gt_sem.shape[0]):
                gt_l = int(gt_sem[j])
                pred_l = int(pred_sem[j])
                gt_classes[gt_l] += 1
                positive_classes[pred_l] += 1
                true_positive_classes[gt_l] += int(gt_l == pred_l)
            '''


            uniq, counts = np.unique(pred_sem, return_counts=True)
            positive_classes[uniq] += counts

            uniq, counts = np.unique(gt_sem, return_counts=True)
            gt_classes[uniq] += counts

            uniq, counts = np.unique(gt_sem[pred_sem == gt_sem], return_counts=True)
            true_positive_classes[uniq] += counts

            # instance
            un = np.unique(pred_ins)
            pts_in_pred = [[] for _ in range(NUM_CLASSES)]
            for ig, g in enumerate(un):  # each object in prediction
                if g == -1:
                    continue
                tmp = (pred_ins == g)
                sem_seg_i = int(stats.mode(pred_sem[tmp])[0])
                pts_in_pred[sem_seg_i] += [tmp]

            un = np.unique(gt_ins)
            pts_in_gt = [[] for _ in range(NUM_CLASSES)]
            for ig, g in enumerate(un):
                tmp = (gt_ins == g)
                sem_seg_i = int(stats.mode(gt_sem[tmp])[0])
                pts_in_gt[sem_seg_i] += [tmp]

            # instance mucov & mwcov
            for i_sem in range(NUM_CLASSES):
                sum_cov = 0
                mean_cov = 0
                mean_weighted_cov = 0
                num_gt_point = 0
                for ig, ins_gt in enumerate(pts_in_gt[i_sem]):
                    ovmax = 0.
                    num_ins_gt_point = np.sum(ins_gt)
                    num_gt_point += num_ins_gt_point
                    for ip, ins_pred in enumerate(pts_in_pred[i_sem]):
                        union = (ins_pred | ins_gt)
                        intersect = (ins_pred & ins_gt)
                        iou = float(np.sum(intersect)) / np.sum(union)

                        if iou > ovmax:
                            ovmax = iou
                            ipmax = ip

                    sum_cov += ovmax
                    mean_weighted_cov += ovmax * num_ins_gt_point

                if len(pts_in_gt[i_sem]) != 0:
                    mean_cov = sum_cov / len(pts_in_gt[i_sem])
                    all_mean_cov[i_sem].append(mean_cov)

                    mean_weighted_cov /= num_gt_point
                    all_mean_weighted_cov[i_sem].append(mean_weighted_cov)

        if dataset == "s3dis":
            # instance precision & recall
            for i_sem in range(NUM_CLASSES):
                tp = [0.] * len(pts_in_pred[i_sem])
                fp = [0.] * len(pts_in_pred[i_sem])
                gtflag = np.zeros(len(pts_in_gt[i_sem]))
                total_gt_ins[i_sem] += len(pts_in_gt[i_sem])

                for ip, ins_pred in enumerate(pts_in_pred[i_sem]):
                    ovmax = -1.

                    for ig, ins_gt in enumerate(pts_in_gt[i_sem]):
                        union = (ins_pred | ins_gt)
                        intersect = (ins_pred & ins_gt)
                        iou = float(np.sum(intersect)) / np.sum(union)

                        if iou > ovmax:
                            ovmax = iou
                            igmax = ig

                    if ovmax >= at:
                        tp[ip] = 1  # true
                    else:
                        fp[ip] = 1  # false positive

                tpsins[i_sem] += tp
                fpsins[i_sem] += fp

        matches_key = os.path.abspath(gt_file)
        # assign gt to predictions
        gt2pred, pred2gt = assign_instances_for_scan(config_general.split,config_general.OW_task, v, gt_file)
        matches[matches_key] = {}
        matches[matches_key]['gt'] = gt2pred
        matches[matches_key]['pred'] = pred2gt
        
        
        sys.stdout.write("\rscans processed: {}".format(i+1))
        sys.stdout.flush()
    
    print('')
    U_Recall, WI, A_OSE, ap_scores = evaluate_matches(matches)
    avgs = compute_averages(ap_scores)
    
    # print
    if U_Recall and WI and A_OSE:
        print_results(avgs,current_epoch,config_general.save_dir,config_general.split,config_general.OW_task,U_Recall, WI[0.5], A_OSE[0.5])
    else:
        print_results(avgs,current_epoch,config_general.save_dir,config_general.split,config_general.OW_task,0.0, 0.0, 0.0)
    
    write_result_file(avgs, output_file)
       
        
    if dataset == "s3dis":
        MUCov = np.zeros(NUM_CLASSES)
        MWCov = np.zeros(NUM_CLASSES)
        for i_sem in range(NUM_CLASSES):
            MUCov[i_sem] = np.mean(all_mean_cov[i_sem])
            MWCov[i_sem] = np.mean(all_mean_weighted_cov[i_sem])

        precision = np.zeros(NUM_CLASSES)
        recall = np.zeros(NUM_CLASSES)
        for i_sem in range(NUM_CLASSES):
            tp = np.asarray(tpsins[i_sem]).astype(np.float)
            fp = np.asarray(fpsins[i_sem]).astype(np.float)
            tp = np.sum(tp)
            fp = np.sum(fp)
            rec = tp / total_gt_ins[i_sem]
            prec = tp / (tp + fp)

            precision[i_sem] = prec
            recall[i_sem] = rec

        '''
        LOG_FOUT = open(os.path.join('results_a5.txt'), 'w')
    
        def log_string(out_str):
            LOG_FOUT.write(out_str + '\n')
            LOG_FOUT.flush()
            print(out_str)
        '''

        return np.mean(precision), np.mean(recall)
    
    if dataset=="scannet200":
        return U_Recall,WI,A_OSE
    

# TODO: remove this
#import pandas as pd
#def main():
#    print("!!! CLI is only for debugging purposes. use `evaluate()` instead.")
#    evaluate(pd.read_pickle("/globalwork/schult/saved_predictions.pkl"), opt.gt_path, opt.output_file)

#if __name__ == '__main__':
#    main()

benchmark/evaluate_semantic_instance.py
aminebdj/3D-OWIS
import os, sys
import csv
try:
    import numpy as np
except:
    print("Failed to import numpy package.")
    sys.exit(-1)
try:
    import imageio
except:
    print("Please install the module 'imageio' for image processing, e.g.")
    print("pip install imageio")
    sys.exit(-1)

# print an error message and quit
def print_error(message, user_fault=False):
    sys.stderr.write('ERROR: ' + str(message) + '\n')
    if user_fault:
        sys.exit(2)
    sys.exit(-1)


# if string s represents an int
def represents_int(s):
    try:
        int(s)
        return True
    except ValueError:
        return False


def read_label_mapping(filename, label_from='raw_category', label_to='nyu40id'):
    assert os.path.isfile(filename)
    mapping = dict()
    with open(filename) as csvfile:
        reader = csv.DictReader(csvfile, delimiter='\t')
        for row in reader:
            mapping[row[label_from]] = int(row[label_to])
    # if ints convert
    if represents_int(list(mapping.keys())[0]):
        mapping = {int(k):v for k,v in mapping.items()}
    return mapping


# input: scene_types.txt or scene_types_all.txt
def read_scene_types_mapping(filename, remove_spaces=True):
    assert os.path.isfile(filename)
    mapping = dict()
    lines = open(filename).read().splitlines()
    lines = [line.split('\t') for line in lines]
    if remove_spaces:
        mapping = { x[1].strip():int(x[0]) for x in lines }
    else:
        mapping = { x[1]:int(x[0]) for x in lines }
    return mapping


# color by label
def visualize_label_image(filename, image):
    height = image.shape[0]
    width = image.shape[1]
    vis_image = np.zeros([height, width, 3], dtype=np.uint8)
    color_palette = create_color_palette()
    for idx, color in enumerate(color_palette):
        vis_image[image==idx] = color
    imageio.imwrite(filename, vis_image)


# color by different instances (mod length of color palette)
def visualize_instance_image(filename, image):
    height = image.shape[0]
    width = image.shape[1]
    vis_image = np.zeros([height, width, 3], dtype=np.uint8)
    color_palette = create_color_palette()
    instances = np.unique(image)
    for idx, inst in enumerate(instances):
        vis_image[image==inst] = color_palette[inst%len(color_palette)]
    imageio.imwrite(filename, vis_image)


# color palette for nyu40 labels
def create_color_palette():
    return [
        (0, 0, 0),
        (174, 199, 232),		# wall
        (152, 223, 138),		# floor
        (31, 119, 180), 		# cabinet
        (255, 187, 120),		# bed
        (188, 189, 34), 		# chair
        (140, 86, 75),  		# sofa
        (255, 152, 150),		# table
        (214, 39, 40),  		# door
        (197, 176, 213),		# window
        (148, 103, 189),		# bookshelf
        (196, 156, 148),		# picture
        (23, 190, 207), 		# counter
        (178, 76, 76),
        (247, 182, 210),		# desk
        (66, 188, 102),
        (219, 219, 141),		# curtain
        (140, 57, 197),
        (202, 185, 52),
        (51, 176, 203),
        (200, 54, 131),
        (92, 193, 61),
        (78, 71, 183),
        (172, 114, 82),
        (255, 127, 14), 		# refrigerator
        (91, 163, 138),
        (153, 98, 156),
        (140, 153, 101),
        (158, 218, 229),		# shower curtain
        (100, 125, 154),
        (178, 127, 135),
        (120, 185, 128),
        (146, 111, 194),
        (44, 160, 44),  		# toilet
        (112, 128, 144),		# sink
        (96, 207, 209),
        (227, 119, 194),		# bathtub
        (213, 92, 176),
        (94, 106, 211),
        (82, 84, 163),  		# otherfurn
        (100, 85, 144)
    ]

benchmark/util.py
aminebdj/3D-OWIS
import os, sys
import json

try:
    import numpy as np
except:
    print("Failed to import numpy package.")
    sys.exit(-1)

try:
    from plyfile import PlyData, PlyElement
except:
    print("Please install the module 'plyfile' for PLY i/o, e.g.")
    print("pip install plyfile")
    sys.exit(-1)

import benchmark.util as util
############################################################################################################################
from datasets.scannet200.scannet200_splits import HEAD_CATS_SCANNET_200, TAIL_CATS_SCANNET_200, COMMON_CATS_SCANNET_200
from datasets.scannet200.owis_splits import UNKNOWN_CLASSES_LABELS
############################################################################################################################
# matrix: 4x4 np array
# points Nx3 np array
def transform_points(matrix, points):
    assert len(points.shape) == 2 and points.shape[1] == 3
    num_points = points.shape[0]
    p = np.concatenate([points, np.ones((num_points, 1))], axis=1)
    p = np.matmul(matrix, np.transpose(p))
    p = np.transpose(p)
    p[:,:3] /= p[:,3,None]
    return p[:,:3]


def export_ids(filename, ids):
    with open(filename, 'w') as f:
        for id in ids:
            f.write('%d\n' % id)


def load_ids(filename):
    ids = open(filename).read().splitlines()
    ids = np.array(ids, dtype=np.int64)
    return ids


def read_mesh_vertices(filename):
    assert os.path.isfile(filename)
    with open(filename, 'rb') as f:
        plydata = PlyData.read(f)
        num_verts = plydata['vertex'].count
        vertices = np.zeros(shape=[num_verts, 3], dtype=np.float32)
        vertices[:,0] = plydata['vertex'].data['x']
        vertices[:,1] = plydata['vertex'].data['y']
        vertices[:,2] = plydata['vertex'].data['z']
    return vertices


# export 3d instance labels for instance evaluation
def export_instance_ids_for_eval(filename, label_ids, instance_ids):
    assert label_ids.shape[0] == instance_ids.shape[0]
    output_mask_path_relative = 'pred_mask'
    name = os.path.splitext(os.path.basename(filename))[0]
    output_mask_path = os.path.join(os.path.dirname(filename), output_mask_path_relative)
    if not os.path.isdir(output_mask_path):
        os.mkdir(output_mask_path)
    insts = np.unique(instance_ids)
    zero_mask = np.zeros(shape=(instance_ids.shape[0]), dtype=np.int32)
    with open(filename, 'w') as f:
        for idx, inst_id in enumerate(insts):
            if inst_id == 0:  # 0 -> no instance for this vertex
                continue
            output_mask_file = os.path.join(output_mask_path_relative, name + '_' + str(idx) + '.txt')
            loc = np.where(instance_ids == inst_id)
            label_id = label_ids[loc[0][0]]
            f.write('%s %d %f\n' % (output_mask_file, label_id, 1.0))
            # write mask
            mask = np.copy(zero_mask)
            mask[loc[0]] = 1
            export_ids(output_mask_file, mask)


# ------------ Instance Utils ------------ #

class Instance(object):
    instance_id = 0
    label_id = 0
    vert_count = 0
    med_dist = -1
    dist_conf = 0.0

    def __init__(self,mesh_vert_instances, instance_id):
        if (instance_id == -1):
            return
        self.instance_id     = int(instance_id)
        self.label_id    = int(self.get_label_id(instance_id))
        self.vert_count = int(self.get_instance_verts(mesh_vert_instances, instance_id))

    def get_label_id(self, instance_id):
        return int(instance_id // 1000)

    def get_instance_verts(self, mesh_vert_instances, instance_id):
        return (mesh_vert_instances == instance_id).sum()

    def to_json(self):
        return json.dumps(self, default=lambda o: o.__dict__, sort_keys=True, indent=4)

    def to_dict(self):
        dict = {}
        dict["instance_id"] = self.instance_id
        dict["label_id"]    = self.label_id
        dict["vert_count"]  = self.vert_count
        dict["med_dist"]    = self.med_dist
        dict["dist_conf"]   = self.dist_conf
        return dict

    def from_json(self, data):
        self.instance_id     = int(data["instance_id"])
        self.label_id        = int(data["label_id"])
        self.vert_count      = int(data["vert_count"])
        if ("med_dist" in data):
            self.med_dist    = float(data["med_dist"])
            self.dist_conf   = float(data["dist_conf"])

    def __str__(self):
        return "("+str(self.instance_id)+")"

def read_instance_prediction_file(filename, pred_path):
    lines = open(filename).read().splitlines()
    instance_info = {}
    abs_pred_path = os.path.abspath(pred_path)
    for line in lines:
        parts = line.split(' ')
        if len(parts) != 3:
            util.print_error('invalid instance prediction file. Expected (per line): [rel path prediction] [label id prediction] [confidence prediction]')
        if os.path.isabs(parts[0]):
            util.print_error('invalid instance prediction file. First entry in line must be a relative path')
        mask_file = os.path.join(os.path.dirname(filename), parts[0])
        mask_file = os.path.abspath(mask_file)
        # check that mask_file lives inside prediction path
        if os.path.commonprefix([mask_file, abs_pred_path]) != abs_pred_path:
            util.print_error('predicted mask {} in prediction text file {} points outside of prediction path.'.format(mask_file,filename))

        info            = {}
        info["label_id"] = int(float(parts[1]))
        info["conf"]    = float(parts[2])
        instance_info[mask_file]  = info
    return instance_info


def get_instances(split,task, ids, class_ids, class_labels, id2label):
    instances = {}
    for label in class_labels:
        instances[label] = []
    instance_ids = np.unique(ids)
    for id in instance_ids:
        if id == 0:
            continue
        inst = Instance(ids, id)
        #common tail labels
        
        # COMMON_TAIL_LABELS = TAIL_CATS_SCANNET_200+COMMON_CATS_SCANNET_200
        # TAIL_LABELS = TAIL_CATS_SCANNET_200
        unknown_class_labels = UNKNOWN_CLASSES_LABELS[split][task]
        #get their ids
        try:
            # if task == "task1":
            #     if id2label[inst.label_id] in COMMON_TAIL_LABELS:
            #         inst.label_id = 3000
            # elif task == "task2":
            #     if id2label[inst.label_id] in TAIL_LABELS:
            #         inst.label_id = 3000
            if unknown_class_labels is not None:
                if id2label[inst.label_id] in unknown_class_labels:
                    inst.label_id = 3000
        except:
            pass
            # print(inst.label_id)
        #change them to 3000
        
        if inst.label_id in class_ids:
            #assign 3000 uknown
            instances[id2label[inst.label_id]].append(inst.to_dict())
            # if id2label[inst.label_id]=='unknown':
            #     print(inst.vert_count)
            
    return instances

benchmark/util_3d.py
aminebdj/3D-OWIS
import re
import os
import numpy as np
from fire import Fire
from natsort import natsorted
from loguru import logger
from datasets.preprocessing.base_preprocessing import BasePreprocessing


class ScannetPreprocessing(BasePreprocessing):
    def __init__(
            self,
            data_dir: str = "./data/raw/s3dis",
            save_dir: str = "./data/processed/s3dis",
            modes: tuple = ("Area_1", "Area_2", "Area_3", "Area_4", "Area_5", "Area_6"),
            n_jobs: int = -1,
    ):
        super().__init__(data_dir, save_dir, modes, n_jobs)

        self.class_map = {
            'ceiling': 0,
            'floor': 1,
            'wall': 2,
            'beam': 3,
            'column': 4,
            'window': 5,
            'door': 6,
            'table': 7,
            'chair': 8,
            'sofa': 9,
            'bookcase': 10,
            'board': 11,
            'clutter': 12,
            'stairs': 12  # stairs are also mapped to clutter
        }

        self.color_map = [
            [0, 255, 0],        # ceiling
            [0, 0, 255],        # floor
            [0, 255, 255],      # wall
            [255, 255, 0],      # beam
            [255, 0, 255],      # column
            [100, 100, 255],    # window
            [200, 200, 100],    # door
            [170, 120, 200],    # table
            [255, 0, 0],        # chair
            [200, 100, 100],    # sofa
            [10, 200, 100],     # bookcase
            [200, 200, 200],    # board
            [50, 50, 50]]      # clutter

        self.create_label_database()

        for mode in self.modes:
            filepaths = []
            for scene_path in [f.path for f in os.scandir(self.data_dir / mode) if f.is_dir()]:
                filepaths.append(scene_path)
            self.files[mode] = natsorted(filepaths)

    def create_label_database(self):
        label_database = dict()
        for class_name, class_id in self.class_map.items():
            label_database[class_id] = {
                'color': self.color_map[class_id],
                'name': class_name,
                'validation': True
            }

        self._save_yaml(self.save_dir / "label_database.yaml", label_database)
        return label_database

    def _buf_count_newlines_gen(self, fname):
        def _make_gen(reader):
            while True:
                b = reader(2 ** 16)
                if not b: break
                yield b

        with open(fname, "rb") as f:
            count = sum(buf.count(b"\n") for buf in _make_gen(f.raw.read))
        return count

    def process_file(self, filepath, mode):
        """process_file.

        Please note, that for obtaining segmentation labels ply files were used.

        Args:
            filepath: path to the main ply file
            mode: train, test or validation

        Returns:
            filebase: info about file
        """
        filebase = {
            "filepath": filepath,
            "scene": filepath.split("/")[-1],
            "area": mode,
            "raw_filepath": str(filepath),
            "file_len": -1,
        }

        scene_name = filepath.split("/")[-1]
        instance_counter = 0
        scene_points = []
        for instance in [f for f in os.scandir(self.data_dir / mode / scene_name / "Annotations")
                         if f.name.endswith(".txt")]:
            instance_class = self.class_map[instance.name.split("_")[0]]
            instance_points = np.loadtxt(instance.path)

            instance_normals = np.ones((instance_points.shape[0], 3))
            instance_class = np.array(instance_class).repeat(instance_points.shape[0])[..., None]
            instance_id = np.array(instance_counter).repeat(instance_points.shape[0])[..., None]

            instance_points = np.hstack((instance_points,
                                         instance_normals,
                                         instance_class,
                                         instance_id))

            scene_points.append(instance_points)
            instance_counter += 1

        points = np.vstack(scene_points)

        pcd_size = self._buf_count_newlines_gen(f"{filepath}/{scene_name}.txt")
        if points.shape[0] != pcd_size:
            print(f"FILE SIZE DOES NOT MATCH FOR {filepath}/{scene_name}.txt")
            print(f"({points.shape[0]} vs. {pcd_size})")

        filebase["raw_segmentation_filepath"] = ""

        # add segment id as additional feature (DUMMY)
        points = np.hstack((points, np.ones(points.shape[0])[..., None]))
        points[:, [9, 10, -1]] = points[:, [-1, 9, 10]]  # move segments after RGB

        gt_data = (points[:, -2] + 1) * 1000 + points[:, -1] + 1

        file_len = len(points)
        filebase["file_len"] = file_len

        processed_filepath = self.save_dir / mode / f"{scene_name}.npy"
        if not processed_filepath.parent.exists():
            processed_filepath.parent.mkdir(parents=True, exist_ok=True)
        np.save(processed_filepath, points.astype(np.float32))
        filebase["filepath"] = str(processed_filepath)

        processed_gt_filepath = self.save_dir / "instance_gt" / mode / f"{scene_name}.txt"
        if not processed_gt_filepath.parent.exists():
            processed_gt_filepath.parent.mkdir(parents=True, exist_ok=True)
        np.savetxt(processed_gt_filepath, gt_data.astype(np.int32), fmt="%d")
        filebase["instance_gt_filepath"] = str(processed_gt_filepath)

        filebase["color_mean"] = [
            float((points[:, 3] / 255).mean()),
            float((points[:, 4] / 255).mean()),
            float((points[:, 5] / 255).mean()),
        ]
        filebase["color_std"] = [
            float(((points[:, 3] / 255) ** 2).mean()),
            float(((points[:, 4] / 255) ** 2).mean()),
            float(((points[:, 5] / 255) ** 2).mean()),
        ]
        return filebase

    def compute_color_mean_std(
            self, train_database_path: str = ""
    ):
        area_database_paths = [f for f in os.scandir(self.save_dir)
                               if f.name.startswith("Area_") and f.name.endswith(".yaml")]

        for database_path in area_database_paths:
            database = self._load_yaml(database_path.path)
            color_mean, color_std = [], []
            for sample in database:
                color_std.append(sample["color_std"])
                color_mean.append(sample["color_mean"])

            color_mean = np.array(color_mean).mean(axis=0)
            color_std = np.sqrt(np.array(color_std).mean(axis=0) - color_mean ** 2)
            feats_mean_std = {
                "mean": [float(each) for each in color_mean],
                "std": [float(each) for each in color_std],
            }
            self._save_yaml(self.save_dir / f"{database_path.name}_color_mean_std.yaml", feats_mean_std)

        for database_path in area_database_paths:
            all_mean, all_std = [], []
            for let_out_path in area_database_paths:
                if database_path == let_out_path:
                    continue

                database = self._load_yaml(let_out_path.path)
                for sample in database:
                    all_std.append(sample["color_std"])
                    all_mean.append(sample["color_mean"])

            all_color_mean = np.array(all_mean).mean(axis=0)
            all_color_std = np.sqrt(np.array(all_std).mean(axis=0) - all_color_mean ** 2)
            feats_mean_std = {
                "mean": [float(each) for each in all_color_mean],
                "std": [float(each) for each in all_color_std],
            }
            file_path = database_path.name.replace("_database.yaml", "")
            self._save_yaml(self.save_dir / f"{file_path}_color_mean_std.yaml", feats_mean_std)

    @logger.catch
    def fix_bugs_in_labels(self):
        pass

    def joint_database(self, train_modes=("Area_1", "Area_2", "Area_3", "Area_4", "Area_5", "Area_6")):
        for mode in train_modes:
            joint_db = []
            for let_out in train_modes:
                if mode == let_out:
                    continue
                joint_db.extend(self._load_yaml(self.save_dir / (let_out + "_database.yaml")))
            self._save_yaml(self.save_dir / f"train_{mode}_database.yaml", joint_db)

    def _parse_scene_subscene(self, name):
        scene_match = re.match(r"scene(\d{4})_(\d{2})", name)
        return int(scene_match.group(1)), int(scene_match.group(2))


if __name__ == "__main__":
    Fire(ScannetPreprocessing)

datasets/preprocessing/s3dis_preprocessing.py
aminebdj/3D-OWIS
import re
from pathlib import Path
import numpy as np
import pandas as pd
from fire import Fire
from natsort import natsorted
from loguru import logger

from datasets.preprocessing.base_preprocessing import BasePreprocessing
from utils.point_cloud_utils import load_ply_with_normals

from datasets.scannet200.scannet200_constants import VALID_CLASS_IDS_200, SCANNET_COLOR_MAP_200, CLASS_LABELS_200


class ScannetPreprocessing(BasePreprocessing):
    def __init__(
            self,
            data_dir: str = "./data/raw/scannet/scannet",
            save_dir: str = "./data/processed/scannet",
            modes: tuple = ("train", "validation", "test"),
            n_jobs: int = -1,
            git_repo: str = "./data/raw/scannet/ScanNet",
            scannet200: bool = False
    ):
        super().__init__(data_dir, save_dir, modes, n_jobs)

        self.scannet200 = scannet200

        if self.scannet200:
            self.labels_pd = pd.read_csv(self.data_dir / "scannetv2-labels.combined.tsv", sep='\t', header=0)

        git_repo = Path(git_repo)
        self.create_label_database(git_repo)
        for mode in self.modes:
            trainval_split_dir = git_repo / "Tasks" / "Benchmark"
            scannet_special_mode = "val" if mode == "validation" else mode
            with open(
                    trainval_split_dir / (f"scannetv2_{scannet_special_mode}.txt")
            ) as f:
                # -1 because the last one is always empty
                split_file = f.read().split("\n")[:-1]

            scans_folder = "scans_test" if mode == "test" else "scans"
            filepaths = []
            for scene in split_file:
                filepaths.append(
                    self.data_dir / scans_folder / scene / (scene + "_vh_clean_2.ply")
                )
            self.files[mode] = natsorted(filepaths)

    def create_label_database(self, git_repo):
        if self.scannet200:
            label_database = {}
            for row_id, class_id in enumerate(VALID_CLASS_IDS_200):
                label_database[class_id] = {
                    'color': SCANNET_COLOR_MAP_200[class_id],
                    'name': CLASS_LABELS_200[row_id],
                    'validation': True
                }
            self._save_yaml(self.save_dir / "label_database.yaml", label_database)
            return label_database
        else:
            if (self.save_dir / "label_database.yaml").exists():
                return self._load_yaml(self.save_dir / "label_database.yaml")
            df = pd.read_csv(self.data_dir / "scannetv2-labels.combined.tsv", sep="\t")
            df = (
                df[~df[["nyu40class", "nyu40id"]].duplicated()][["nyu40class", "nyu40id"]]
                    .set_index("nyu40id")
                    .sort_index()[["nyu40class"]]
                    .rename(columns={"nyu40class": "name"})
                    .replace(" ", "_", regex=True)
            )
            df = pd.DataFrame([{"name": "empty"}]).append(df)
            df["validation"] = False

            with open(
                    git_repo / "Tasks" / "Benchmark" / "classes_SemVoxLabel-nyu40id.txt"
            ) as f:
                for_validation = f.read().split("\n")
            for category in for_validation:
                index = int(re.split(" +", category)[0])
                df.loc[index, "validation"] = True

            # doing this hack because otherwise I will have to install imageio
            with open(git_repo / "BenchmarkScripts" / "util.py") as f:
                util = f.read()
                color_list = eval("[" + util.split("return [\n")[1])

            df["color"] = color_list

            label_database = df.to_dict("index")
            self._save_yaml(self.save_dir / "label_database.yaml", label_database)
            return label_database

    def process_file(self, filepath, mode):
        """process_file.

        Please note, that for obtaining segmentation labels ply files were used.

        Args:
            filepath: path to the main ply file
            mode: train, test or validation

        Returns:
            filebase: info about file
        """
        scene, sub_scene = self._parse_scene_subscene(filepath.name)
        filebase = {
            "filepath": filepath,
            "scene": scene,
            "sub_scene": sub_scene,
            "raw_filepath": str(filepath),
            "file_len": -1,
        }
        # reading both files and checking that they are fitting
        coords, features, _ = load_ply_with_normals(filepath)
        file_len = len(coords)
        filebase["file_len"] = file_len
        points = np.hstack((coords, features))

        if mode in ["train", "validation"]:
            # getting scene information
            description_filepath = Path(filepath).parent / filepath.name.replace(
                "_vh_clean_2.ply", ".txt"
            )
            with open(description_filepath) as f:
                scene_type = f.read().split("\n")[:-1]
            scene_type = scene_type[-1].split(" = ")[1]
            filebase["scene_type"] = scene_type
            filebase["raw_description_filepath"] = description_filepath

            # getting instance info
            instance_info_filepath = next(
                Path(filepath).parent.glob("*.aggregation.json")
            )
            segment_indexes_filepath = next(
                Path(filepath).parent.glob("*[0-9].segs.json")
            )
            instance_db = self._read_json(instance_info_filepath)
            segments = self._read_json(segment_indexes_filepath)
            segments = np.array(segments["segIndices"])
            filebase["raw_instance_filepath"] = instance_info_filepath
            filebase["raw_segmentation_filepath"] = segment_indexes_filepath

            # add segment id as additional feature
            segment_ids = np.unique(segments, return_inverse=True)[1]
            points = np.hstack((points, segment_ids[..., None]))

            # reading labels file
            label_filepath = filepath.parent / filepath.name.replace(
                ".ply", ".labels.ply"
            )
            filebase["raw_label_filepath"] = label_filepath
            label_coords, label_colors, labels = load_ply_with_normals(label_filepath)
            if not np.allclose(coords, label_coords):
                raise ValueError("files doesn't have same coordinates")

            # adding instance label
            labels = labels[:, np.newaxis]
            empty_instance_label = np.full(labels.shape, -1)
            labels = np.hstack((labels, empty_instance_label))
            for instance in instance_db["segGroups"]:
                segments_occupied = np.array(instance["segments"])
                occupied_indices = np.isin(segments, segments_occupied)
                labels[occupied_indices, 1] = instance["id"]

                if self.scannet200:
                    label200 = instance['label']
                    # Map the category name to id
                    label_ids = self.labels_pd[self.labels_pd['raw_category'] == label200]['id']
                    label_id = int(label_ids.iloc[0]) if len(label_ids) > 0 else 0
                    labels[occupied_indices, 0] = label_id
            points = np.hstack((points, labels))

            # gt_data = (points[:, -2] + 1) * 1000 + points[:, -1] + 1
            gt_data = points[:, -2] * 1000 + points[:, -1] + 1
        else:
            segments_test = "../../data/raw/scannet_test_segments"
            segment_indexes_filepath = filepath.name.replace(".ply", ".0.010000.segs.json")
            segments = self._read_json(f"{segments_test}/{segment_indexes_filepath}")
            segments = np.array(segments["segIndices"])
            # add segment id as additional feature
            segment_ids = np.unique(segments, return_inverse=True)[1]
            points = np.hstack((points, segment_ids[..., None]))

        processed_filepath = self.save_dir / mode / f"{scene:04}_{sub_scene:02}.npy"
        if not processed_filepath.parent.exists():
            processed_filepath.parent.mkdir(parents=True, exist_ok=True)
        np.save(processed_filepath, points.astype(np.float32))
        filebase["filepath"] = str(processed_filepath)

        if mode == "test":
            return filebase

        processed_gt_filepath = self.save_dir / "instance_gt" / mode / f"scene{scene:04}_{sub_scene:02}.txt"
        if not processed_gt_filepath.parent.exists():
            processed_gt_filepath.parent.mkdir(parents=True, exist_ok=True)
        np.savetxt(processed_gt_filepath, gt_data.astype(np.int32), fmt="%d")
        filebase["instance_gt_filepath"] = str(processed_gt_filepath)

        filebase["color_mean"] = [
            float((features[:, 0] / 255).mean()),
            float((features[:, 1] / 255).mean()),
            float((features[:, 2] / 255).mean()),
        ]
        filebase["color_std"] = [
            float(((features[:, 0] / 255) ** 2).mean()),
            float(((features[:, 1] / 255) ** 2).mean()),
            float(((features[:, 2] / 255) ** 2).mean()),
        ]
        return filebase

    def compute_color_mean_std(
            self, train_database_path: str = "./data/processed/scannet/train_database.yaml"
    ):
        train_database = self._load_yaml(train_database_path)
        color_mean, color_std = [], []
        for sample in train_database:
            color_std.append(sample["color_std"])
            color_mean.append(sample["color_mean"])

        color_mean = np.array(color_mean).mean(axis=0)
        color_std = np.sqrt(np.array(color_std).mean(axis=0) - color_mean ** 2)
        feats_mean_std = {
            "mean": [float(each) for each in color_mean],
            "std": [float(each) for each in color_std],
        }
        self._save_yaml(self.save_dir / "color_mean_std.yaml", feats_mean_std)

    @logger.catch
    def fix_bugs_in_labels(self):
        if not self.scannet200:
            logger.add(self.save_dir / "fixed_bugs_in_labels.log")
            found_wrong_labels = {
                tuple([270, 0]): 50,
                tuple([270, 2]): 50,
                tuple([384, 0]): 149,
            }
            for scene, wrong_label in found_wrong_labels.items():
                scene, sub_scene = scene
                bug_file = self.save_dir / "train" / f"{scene:04}_{sub_scene:02}.npy"
                points = np.load(bug_file)
                bug_mask = points[:, -1] != wrong_label
                points = points[bug_mask]
                np.save(bug_file, points)
                logger.info(f"Fixed {bug_file}")

    def _parse_scene_subscene(self, name):
        scene_match = re.match(r"scene(\d{4})_(\d{2})", name)
        return int(scene_match.group(1)), int(scene_match.group(2))


if __name__ == "__main__":
    Fire(ScannetPreprocessing)

datasets/preprocessing/scannet_preprocessing.py
aminebdj/3D-OWIS
import re
import os
import numpy as np
from fire import Fire
from natsort import natsorted
from loguru import logger
import pandas as pd

from datasets.preprocessing.base_preprocessing import BasePreprocessing

class STPLS3DPreprocessing(BasePreprocessing):
    def __init__(
            self,
            data_dir: str = "../../data/raw/stpls3d",
            save_dir: str = "../../data/processed/stpls3d",
            modes: tuple = ("train", "validation", "test"),
            n_jobs: int = -1
    ):
        super().__init__(data_dir, save_dir, modes, n_jobs)

        # https://github.com/meidachen/STPLS3D/blob/main/HAIS/STPLS3DInstanceSegmentationChallenge_Codalab_Evaluate.py#L31
        CLASS_LABELS = ['Build', 'LowVeg', 'MediumVeg', 'HighVeg', 'Vehicle', 'Truck', 'Aircraft', 'MilitaryVeh',
                        'Bike', 'Motorcycle', 'LightPole', 'StreetSign', 'Clutter', 'Fence']
        VALID_CLASS_IDS = np.array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14])

        self.class_map = {
            'Ground': 0,
            'Build': 1,
            'LowVeg': 2,
            'MediumVeg': 3,
            'HighVeg': 4,
            'Vehicle': 5,
            'Truck': 6,
            'Aircraft': 7,
            'MilitaryVeh': 8,
            'Bike': 9,
            'Motorcycle': 10,
            'LightPole': 11,
            'StreetSign': 12,
            'Clutter': 13,
            'Fence': 14
        }

        self.color_map = [
            [0, 255, 0],     # Ground
            [0, 0, 255],     # Build
            [0, 255, 255],   # LowVeg
            [255, 255, 0],   # MediumVeg
            [255, 0, 255],   # HiVeg
            [100, 100, 255], # Vehicle
            [200, 200, 100], # Truck
            [170, 120, 200], # Aircraft
            [255, 0, 0],     # MilitaryVec
            [200, 100, 100], # Bike
            [10, 200, 100],  # Motorcycle
            [200, 200, 200], # LightPole
            [50, 50, 50],    # StreetSign
            [60, 130, 60],   # Clutter
            [130, 30, 60]]   # Fence

        self.create_label_database()

        for mode in self.modes:
            filepaths = []
            for scene_path in [f.path for f in os.scandir(self.data_dir / mode)]:
                filepaths.append(scene_path)
            self.files[mode] = natsorted(filepaths)

    def create_label_database(self):
        label_database = dict()
        for class_name, class_id in self.class_map.items():
            label_database[class_id] = {
                'color': self.color_map[class_id],
                'name': class_name,
                'validation': True
            }

        self._save_yaml(self.save_dir / "label_database.yaml", label_database)
        return label_database

    def process_file(self, filepath, mode):
        """process_file.

        Please note, that for obtaining segmentation labels ply files were used.

        Args:
            filepath: path to the main ply file
            mode: train, test or validation

        Returns:
            filebase: info about file
        """
        filebase = {
            "filepath": filepath,
            "scene": filepath.split("/")[-1],
            "raw_filepath": str(filepath),
            "file_len": -1,
        }

        points = pd.read_csv(filepath, header=None).values

        filebase["raw_segmentation_filepath"] = ""

        # add segment id as additional feature (DUMMY)
        if mode in ["train", "validation"]:
            points = np.hstack((points,
                                np.ones(points.shape[0])[..., None],   # normal 1
                                np.ones(points.shape[0])[..., None],   # normal 2
                                np.ones(points.shape[0])[..., None],   # normal 3
                                np.ones(points.shape[0])[..., None]))  # segments
        else:
            # we need to add dummies for semantics and instances
            points = np.hstack((points,
                                np.ones(points.shape[0])[..., None],   # semantic class
                                np.ones(points.shape[0])[..., None],   # instance id
                                np.ones(points.shape[0])[..., None],   # normal 1
                                np.ones(points.shape[0])[..., None],   # normal 2
                                np.ones(points.shape[0])[..., None],   # normal 3
                                np.ones(points.shape[0])[..., None]))  # segments

        points = points[:, [0, 1, 2, 3, 4, 5, 8, 9, 10, 11, 6, 7]]  # move segments after RGB

        # move point clouds to be in positive range (important for split pointcloud function)
        points[:, :3] = points[:, :3] - points[:, :3].min(0)

        points = points.astype(np.float32)

        if mode == "test":
            points = points[:, :-2]
        else:
            points[points[:, -1] == -100., -1] = -1  # -1 indicates "no instance"

        file_len = len(points)
        filebase["file_len"] = file_len

        processed_filepath = self.save_dir / mode / f"{filebase['scene'].replace('.txt', '')}.npy"
        if not processed_filepath.parent.exists():
            processed_filepath.parent.mkdir(parents=True, exist_ok=True)
        np.save(processed_filepath, points.astype(np.float32))
        filebase["filepath"] = str(processed_filepath)

        if mode in ["validation", "test"]:
            blocks = self.splitPointCloud(points)

            filebase["instance_gt_filepath"] = []
            filebase["filepath_crop"] = []
            for block_id, block in enumerate(blocks):
                if len(block) > 10000:
                    if mode == "validation":
                        new_instance_ids = np.unique(block[:, -1], return_inverse=True)[1]

                        assert new_instance_ids.shape[0] == block.shape[0]
                        # == 0 means -1 == no instance
                        # new_instance_ids[new_instance_ids == 0]
                        assert new_instance_ids.max() < 1000, "we cannot encode when there are more than 999 instances in a block"

                        gt_data = (block[:, -2]) * 1000 + new_instance_ids

                        processed_gt_filepath = self.save_dir / "instance_gt" / mode / f"{filebase['scene'].replace('.txt', '')}_{block_id}.txt"
                        if not processed_gt_filepath.parent.exists():
                            processed_gt_filepath.parent.mkdir(parents=True, exist_ok=True)
                        np.savetxt(processed_gt_filepath, gt_data.astype(np.int32), fmt="%d")
                        filebase["instance_gt_filepath"].append(str(processed_gt_filepath))

                    processed_filepath = self.save_dir / mode / f"{filebase['scene'].replace('.txt', '')}_{block_id}.npy"
                    if not processed_filepath.parent.exists():
                        processed_filepath.parent.mkdir(parents=True, exist_ok=True)
                    np.save(processed_filepath, block.astype(np.float32))
                    filebase["filepath_crop"].append(str(processed_filepath))
                else:
                    print("block was smaller than 1000 points")
                    assert False

        filebase["color_mean"] = [
            float((points[:, 3] / 255).mean()),
            float((points[:, 4] / 255).mean()),
            float((points[:, 5] / 255).mean()),
        ]
        filebase["color_std"] = [
            float(((points[:, 3] / 255) ** 2).mean()),
            float(((points[:, 4] / 255) ** 2).mean()),
            float(((points[:, 5] / 255) ** 2).mean()),
        ]
        return filebase

    def compute_color_mean_std(
            self, train_database_path: str = "./data/processed/stpls3d/train_database.yaml"
    ):
        train_database = self._load_yaml(train_database_path)
        color_mean, color_std = [], []
        for sample in train_database:
            color_std.append(sample["color_std"])
            color_mean.append(sample["color_mean"])

        color_mean = np.array(color_mean).mean(axis=0)
        color_std = np.sqrt(np.array(color_std).mean(axis=0) - color_mean ** 2)
        feats_mean_std = {
            "mean": [float(each) for each in color_mean],
            "std": [float(each) for each in color_std],
        }
        self._save_yaml(self.save_dir / "color_mean_std.yaml", feats_mean_std)

    def splitPointCloud(self, cloud, size=50.0, stride=50):
        limitMax = np.amax(cloud[:, 0:3], axis=0)
        width = int(np.ceil((limitMax[0] - size) / stride)) + 1
        depth = int(np.ceil((limitMax[1] - size) / stride)) + 1
        cells = [(x * stride, y * stride) for x in range(width) for y in range(depth)]
        blocks = []
        for (x, y) in cells:
            xcond = (cloud[:, 0] <= x + size) & (cloud[:, 0] >= x)
            ycond = (cloud[:, 1] <= y + size) & (cloud[:, 1] >= y)
            cond = xcond & ycond
            block = cloud[cond, :]
            blocks.append(block)
        return blocks

    @logger.catch
    def fix_bugs_in_labels(self):
        pass

    def _parse_scene_subscene(self, name):
        scene_match = re.match(r"scene(\d{4})_(\d{2})", name)
        return int(scene_match.group(1)), int(scene_match.group(2))


if __name__ == "__main__":
    Fire(STPLS3DPreprocessing)

datasets/preprocessing/stpls3d_preprocessing.py
aminebdj/3D-OWIS
# Copyright (c) Facebook, Inc. and its affiliates.
import numpy as np
import torch


def check_aspect(crop_range, aspect_min):
    xy_aspect = np.min(crop_range[:2]) / np.max(crop_range[:2])
    xz_aspect = np.min(crop_range[[0, 2]]) / np.max(crop_range[[0, 2]])
    yz_aspect = np.min(crop_range[1:]) / np.max(crop_range[1:])
    return (
        (xy_aspect >= aspect_min)
        or (xz_aspect >= aspect_min)
        or (yz_aspect >= aspect_min)
    )


class RandomCuboid(object):
    """
    RandomCuboid augmentation from DepthContrast [https://arxiv.org/abs/2101.02691]
    We slightly modify this operation to account for object detection.
    This augmentation randomly crops a cuboid from the input and
    ensures that the cropped cuboid contains at least one bounding box
    """

    def __init__(
        self,
        min_points,
        #aspect=0.8,
        crop_length=6.0,
        version1=True
    ):
        #self.aspect = aspect
        self.crop_length = crop_length
        self.min_points = min_points
        self.version1 = version1

    def __call__(self, point_cloud):
        if point_cloud.shape[0] < self.min_points:
            print("too small pcd")
            return np.ones(point_cloud.shape[0], dtype=np.bool)

        range_xyz = np.max(point_cloud[:, :2], axis=0) - np.min(
            point_cloud[:, :2], axis=0
        )

        for _ in range(100):
            #crop_range = self.min_crop + np.random.rand(3) * (
            #    self.max_crop - self.min_crop
            #)
            #crop_range[-1] = 999.
            # if not check_aspect(crop_range, self.aspect):
            #     continue

            sample_center = point_cloud[:, :2].min(axis=0) + range_xyz/2

            if self.version1:
                offset_x = np.random.uniform(-range_xyz[0]/4,range_xyz[0]/4)
                offset_y = np.random.uniform(-range_xyz[1]/4,range_xyz[1]/4)
            else:
                offset_x = np.random.uniform(-(range_xyz[0]/2) + self.crop_length / 4,
                                             +(range_xyz[0]/2) - self.crop_length / 4)
                offset_y = np.random.uniform(-(range_xyz[1]/2) + self.crop_length / 4,
                                             +(range_xyz[1]/2) - self.crop_length / 4)

            sample_center[0] = sample_center[0] + offset_x
            sample_center[1] = sample_center[1] + offset_y

            min_xy = sample_center - self.crop_length / 2
            max_xy = sample_center + self.crop_length / 2

            upper_idx = (
                np.sum((point_cloud[:, :2] <= max_xy).astype(np.int32), 1) == 2
            )
            lower_idx = (
                np.sum((point_cloud[:, :2] >= min_xy).astype(np.int32), 1) == 2
            )

            new_pointidx = (upper_idx) & (lower_idx)

            if np.sum(new_pointidx) < self.min_points:
                print("TOO SMALL")
                continue

            return new_pointidx

        # fallback
        print("FALLBACK")
        return np.ones(point_cloud.shape[0], dtype=np.bool)

datasets/random_cuboid.py
aminebdj/3D-OWIS

#wall and floor are excluded from the training in all splits


SPLIT_A_SUBSET_1_IDS = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 32, 33, 34, 39, 42, 43, 45, 46, 47, 48, 50, 57, 58, 59, 60, 62, 63, 66, 71, 74, 78, 79, 81, 90, 98, 99, 115, 116, 161, 170, 171, 172, 178]
SPLIT_A_SUBSET_2_IDS = [31, 35, 36, 37, 38, 40, 41, 44, 49, 51, 52, 53, 54, 55, 56, 61, 64, 65, 68, 69, 70, 72, 75, 76, 77, 80, 83, 86, 87, 88, 89, 93, 95, 96, 97, 101, 102, 103, 104, 106, 107, 108, 111, 113, 118, 123, 126, 127, 128, 132, 133, 140, 143, 144, 149, 152, 166, 167, 169, 173, 174, 175, 176, 180, 184, 189, 195, 197]
SPLIT_A_SUBSET_3_IDS = [67, 73, 82, 84, 85, 91, 92, 94, 100, 105, 109, 110, 112, 114, 117, 119, 120, 121, 122, 124, 125, 129, 130, 131, 134, 135, 136, 137, 138, 139, 141, 142, 145, 146, 147, 148, 150, 151, 153, 154, 155, 156, 157, 158, 159, 160, 162, 163, 164, 165, 168, 177, 179, 181, 182, 183, 185, 186, 187, 188, 190, 191, 192, 193, 194, 196]
SPLIT_A_SUBSET_1_LABELS = ['tv stand', 'curtain', 'blinds', 'shower curtain', 'bookshelf', 'tv', 'kitchen cabinet', 'pillow', 'lamp', 'dresser', 'monitor', 'object', 'ceiling', 'board', 'stove', 'closet wall', 'couch', 'office chair', 'kitchen counter', 'shower', 'closet', 'doorframe', 'sofa chair', 'mailbox', 'nightstand', 'washing machine', 'picture', 'book', 'sink', 'recycling bin', 'table', 'backpack', 'shower wall', 'toilet', 'copier', 'counter', 'stool', 'refrigerator', 'window', 'file cabinet', 'chair', 'plant', 'coffee table', 'stairs', 'armchair', 'cabinet', 'bathroom vanity', 'bathroom stall', 'mirror', 'blackboard', 'trash can', 'stair rail', 'box', 'towel', 'door', 'clothes', 'whiteboard', 'bed', 'bathtub', 'desk', 'wardrobe', 'clothes dryer', 'radiator', 'shelf']
SPLIT_A_SUBSET_2_LABELS = ["cushion", "end table", "dining table", "keyboard", "bag", "toilet paper", "printer", "blanket", "microwave", "shoe", "computer tower", "bottle", "bin", "ottoman", "bench", "basket", "fan", "laptop", "person", "paper towel dispenser", "oven", "rack", "piano", "suitcase", "rail", "container", "telephone", "stand", "light", "laundry basket", "pipe", "seat", "column", "bicycle", "ladder", "jacket", "storage bin", "coffee maker", "dishwasher", "machine", "mat", "windowsill", "bulletin board", "fireplace", "mini fridge", "water cooler", "shower door", "pillar", "ledge", "furniture", "cart", "decoration", "closet door", "vacuum cleaner", "dish rack", "range hood", "projector screen", "divider", "bathroom counter", "laundry hamper", "bathroom stall door", "ceiling light", "trash bin", "bathroom cabinet", "structure", "storage organizer", "potted plant", "mattress"]
SPLIT_A_SUBSET_3_LABELS = ["paper", "plate", "soap dispenser", "bucket", "clock", "guitar", "toilet paper holder", "speaker", "cup", "paper towel roll", "bar", "toaster", "ironing board", "soap dish", "toilet paper dispenser", "fire extinguisher", "ball", "hat", "shower curtain rod", "paper cutter", "tray", "toaster oven", "mouse", "toilet seat cover dispenser", "storage container", "scale", "tissue box", "light switch", "crate", "power outlet", "sign", "projector", "candle", "plunger", "stuffed animal", "headphones", "broom", "guitar case", "dustpan", "hair dryer", "water bottle", "handicap bar", "purse", "vent", "shower floor", "water pitcher", "bowl", "paper bag", "alarm clock", "music stand", "laundry detergent", "dumbbell", "tube", "cd case", "closet rod", "coffee kettle", "shower head", "keyboard piano", "case of water bottles", "coat rack", "folded chair", "fire alarm", "power strip", "calendar", "poster", "luggage"]

SPLIT_B_SUBSET_1_IDS = [0, 1, 2, 3, 5, 8, 9, 11, 16, 18, 19, 22, 25, 26, 27, 29, 30, 32, 33, 35, 36, 38, 39, 42, 44, 47, 51, 55, 56, 62, 64, 65, 72, 75, 76, 77, 79, 80, 81, 87, 88, 94, 96, 101, 102, 112, 113, 118, 121, 127, 136, 139, 140, 141, 143, 144, 145, 151, 157, 164, 165, 170, 172, 173, 177, 186, 187, 188, 189, 194, 195, 196, 197]
SPLIT_B_SUBSET_2_IDS = [4, 6, 7, 12, 14, 15, 17, 20, 37, 41, 43, 48, 52, 54, 57, 60, 63, 67, 68, 71, 74, 78, 83, 84, 85, 86, 89, 91, 93, 95, 100, 105, 106, 108, 111, 120, 124, 125, 130, 132, 133, 134, 142, 148, 155, 161, 166, 167, 175, 178, 179, 181, 184, 190, 192]
SPLIT_B_SUBSET_3_IDS = [10, 13, 21, 23, 24, 28, 31, 34, 40, 45, 46, 49, 50, 53, 58, 59, 61, 66, 69, 70, 73, 82, 90, 92, 97, 98, 99, 103, 104, 107, 109, 110, 114, 115, 116, 117, 119, 122, 123, 126, 128, 129, 131, 135, 137, 138, 146, 147, 149, 150, 152, 153, 154, 156, 158, 159, 160, 162, 163, 168, 169, 171, 174, 176, 180, 182, 183, 185, 191, 193]
SPLIT_B_SUBSET_1_LABELS = ['alarm clock', 'backpack', 'bag', 'bed', 'blanket', 'case of water bottles', 'ceiling', 'closet', 'closet door', 'closet wall', 'clothes', 'coat rack', 'container', 'curtain', 'door', 'dresser', 'dumbbell', 'fan', 'guitar case', 'hat', 'ironing board', 'lamp', 'laptop', 'laundry basket', 'laundry hamper', 'luggage', 'mattress', 'mini fridge', 'nightstand', 'object', 'pillow', 'poster', 'power outlet', 'purse', 'rack', 'recycling bin', 'shelf', 'shoe', 'sign', 'storage bin', 'storage organizer', 'suitcase', 'tissue box', 'wardrobe', 'decoration', 'armchair', 'bench', 'bicycle', 'candle', 'chair', 'coffee table', 'couch', 'dining table', 'end table', 'fireplace', 'jacket', 'keyboard piano', 'light', 'music stand', 'ottoman', 'piano', 'picture', 'pillar', 'plant', 'potted plant', 'rail', 'sofa chair', 'speaker', 'stool', 'table', 'tv', 'tv stand', 'vacuum cleaner']
SPLIT_B_SUBSET_2_LABELS = ['guitar', 'paper towel roll', 'book', 'bookshelf', 'cart', 'furniture', 'blackboard', 'projector', 'seat', 'folded chair', 'office chair', 'projector screen', 'whiteboard', 'bin', 'bucket', 'bulletin board', 'copier', 'machine', 'mailbox', 'paper cutter', 'printer', 'column', 'storage container', 'blinds', 'structure', 'water bottle', 'ball', 'board', 'box', 'cabinet', 'cd case', 'ceiling light', 'clock', 'computer tower', 'cup', 'desk', 'divider', 'file cabinet', 'headphones', 'keyboard', 'monitor', 'mouse', 'paper', 'person', 'power strip', 'radiator', 'stand', 'telephone', 'tray', 'tube', 'window', 'windowsill', 'pipe', 'stair rail', 'stairs']
SPLIT_B_SUBSET_3_LABELS = ['bar', 'basket', 'bathroom cabinet', 'bathroom counter', 'bathroom stall', 'bathroom stall door', 'bathroom vanity', 'bathtub', 'bottle', 'broom', 'clothes dryer', 'cushion', 'doorframe', 'fire alarm', 'hair dryer', 'handicap bar', 'ledge', 'light switch', 'mat', 'mirror', 'paper towel dispenser', 'plunger', 'scale', 'shower', 'shower curtain', 'shower curtain rod', 'shower door', 'shower floor', 'shower head', 'shower wall', 'sink', 'soap dish', 'soap dispenser', 'toilet', 'toilet paper', 'toilet paper dispenser', 'toilet paper holder', 'toilet seat cover dispenser', 'towel', 'trash bin', 'washing machine', 'closet rod', 'dustpan', 'laundry detergent', 'stuffed animal', 'bowl', 'calendar', 'coffee kettle', 'coffee maker', 'counter', 'dish rack', 'dishwasher', 'fire extinguisher', 'kitchen cabinet', 'kitchen counter', 'microwave', 'oven', 'paper bag', 'plate', 'range hood', 'refrigerator', 'stove', 'toaster', 'toaster oven', 'trash can', 'vent', 'water cooler', 'water pitcher', 'crate', 'ladder']

SPLIT_C_SUBSET_1_IDS = [0, 2, 4, 5, 6, 8, 20, 33, 36, 37, 38, 39, 40, 43, 46, 50, 51, 52, 57, 60, 61, 66, 69, 73, 75, 76, 77, 78, 88, 91, 94, 102, 105, 109, 111, 119, 121, 123, 124, 127, 128, 129, 132, 136, 139, 140, 144, 146, 147, 148, 154, 155, 160, 161, 164, 168, 170, 171, 172, 173, 177, 178, 180, 185, 195, 196]
SPLIT_C_SUBSET_2_IDS = [3, 10, 13, 14, 15, 16, 17, 18, 19, 22, 25, 27, 28, 34, 42, 44, 47, 48, 55, 58, 59, 62, 64, 65, 67, 70, 71, 79, 81, 83, 84, 85, 87, 89, 90, 95, 96, 97, 99, 100, 106, 112, 113, 116, 118, 120, 122, 126, 130, 131, 135, 141, 149, 151, 156, 162, 163, 165, 167, 176, 183, 184, 186, 187, 192, 194]
SPLIT_C_SUBSET_3_IDS = [1, 7, 9, 11, 12, 21, 23, 24, 26, 29, 30, 31, 32, 35, 41, 45, 49, 53, 54, 56, 63, 68, 72, 74, 80, 82, 86, 92, 93, 98, 101, 103, 104, 107, 108, 110, 114, 115, 117, 125, 133, 134, 137, 138, 142, 143, 145, 150, 152, 153, 157, 158, 159, 166, 169, 174, 175, 179, 181, 182, 188, 189, 190, 191, 193, 197]
SPLIT_C_SUBSET_1_LABELS = ['basket', 'trash can', 'stair rail', 'toaster oven', 'laundry hamper', 'bulletin board', 'dining table', 'stuffed animal', 'bathroom vanity', 'box', 'ceiling', 'potted plant', 'luggage', 'closet wall', 'paper cutter', 'desk', 'object', 'rail', 'tissue box', 'plate', 'keyboard', 'hat', 'copier', 'shower head', 'bed', 'paper towel dispenser', 'fire extinguisher', 'paper towel roll', 'backpack', 'water bottle', 'bathroom cabinet', 'stove', 'laundry basket', 'alarm clock', 'headphones', 'piano', 'guitar', 'bag', 'door', 'speaker', 'water cooler', 'shoe', 'water pitcher', 'dumbbell', 'furniture', 'decoration', 'radiator', 'plunger', 'shower', 'bar', 'hair dryer', 'suitcase', 'cabinet', 'chair', 'board', 'laundry detergent', 'whiteboard', 'vacuum cleaner', 'power outlet', 'storage bin', 'computer tower', 'mailbox', 'shelf', 'ledge', 'pillar', 'toilet paper']
SPLIT_C_SUBSET_2_LABELS = ['ironing board', 'divider', 'oven', 'dish rack', 'shower door', 'mini fridge', 'bicycle', 'laptop', 'armchair', 'couch', 'coffee kettle', 'counter', 'structure', 'pipe', 'bowl', 'shower curtain rod', 'sofa chair', 'clothes dryer', 'coffee table', 'stairs', 'toilet seat cover dispenser', 'machine', 'paper bag', 'book', 'blinds', 'monitor', 'shower wall', 'curtain', 'closet', 'telephone', 'fan', 'ball', 'bucket', 'sign', 'mirror', 'clock', 'nightstand', 'tv stand', 'handicap bar', 'poster', 'blanket', 'cup', 'recycling bin', 'lamp', 'scale', 'mouse', 'wardrobe', 'ottoman', 'paper', 'power strip', 'fireplace', 'doorframe', 'toilet', 'trash bin', 'case of water bottles', 'light', 'washing machine', 'guitar case', 'sink', 'bathtub', 'ladder', 'bookshelf', 'column', 'clothes', 'keyboard piano', 'music stand']
SPLIT_C_SUBSET_3_LABELS = ['mattress', 'toaster', 'stool', 'plant', 'folded chair', 'microwave', 'cushion', 'bench', 'soap dispenser', 'storage organizer', 'shower curtain', 'cart', 'kitchen counter', 'towel', 'blackboard', 'tv', 'printer', 'stand', 'rack', 'bathroom counter', 'closet rod', 'bottle', 'range hood', 'purse', 'candle', 'person', 'coffee maker', 'light switch', 'storage container', 'bathroom stall door', 'shower floor', 'kitchen cabinet', 'refrigerator', 'fire alarm', 'tube', 'toilet paper holder', 'ceiling light', 'picture', 'end table', 'closet door', 'file cabinet', 'crate', 'toilet paper dispenser', 'pillow', 'mat', 'bathroom stall', 'broom', 'container', 'seat', 'jacket', 'dresser', 'dustpan', 'table', 'projector', 'window', 'windowsill', 'tray', 'cd case', 'soap dish', 'office chair', 'dishwasher', 'vent', 'coat rack', 'calendar', 'bin', 'projector screen']

"""
IDS for currently known, unknown, and previously known classes for each split and task.

IMPORTANT: IDS in this script correspond to the index of the prediction of the class in the classification head, 
           and NOT the label IDS defined by scannet200 benchmark.
"""

# UNKNOWN_CLASSES_IDS: IDS of the unknown classes
UNKNOWN_CLASSES_IDS = {
    'A':{
        'task1': SPLIT_A_SUBSET_2_IDS+SPLIT_A_SUBSET_3_IDS,
        'task2': SPLIT_A_SUBSET_3_IDS,
        'task3': None
    },
    'B':{
        'task1': SPLIT_B_SUBSET_2_IDS+SPLIT_B_SUBSET_3_IDS,
        'task2': SPLIT_B_SUBSET_3_IDS,
        'task3': None
    },
    'C':{
        'task1': SPLIT_C_SUBSET_2_IDS+SPLIT_C_SUBSET_3_IDS,
        'task2': SPLIT_C_SUBSET_3_IDS,
        'task3': None
    }
}

# KNOWN_CLASSES_IDS: IDS of the currently known classes
KNOWN_CLASSES_IDS = {
    'A': {
        'task1': SPLIT_A_SUBSET_1_IDS,
        'task2': SPLIT_A_SUBSET_2_IDS,
        'task3': SPLIT_A_SUBSET_3_IDS
    },
    'B': {
        'task1': SPLIT_B_SUBSET_1_IDS,
        'task2': SPLIT_B_SUBSET_2_IDS,
        'task3': SPLIT_B_SUBSET_3_IDS
    },
    'C': {
        'task1': SPLIT_C_SUBSET_1_IDS,
        'task2': SPLIT_C_SUBSET_2_IDS,
        'task3': SPLIT_C_SUBSET_3_IDS
    }
}

# PREV_KNOWN_CLASSES_IDS: IDS of the previously known classes
PREV_KNOWN_CLASSES_IDS = {
    'A': {
        'task1': None,
        'task2': SPLIT_A_SUBSET_1_IDS,
        'task3': SPLIT_A_SUBSET_1_IDS+SPLIT_A_SUBSET_2_IDS
    },
    'B': {
        'task1': None,  
        'task2': SPLIT_B_SUBSET_1_IDS,
        'task3': SPLIT_B_SUBSET_1_IDS+SPLIT_B_SUBSET_2_IDS
    },
    'C': {
        'task1': None,  
        'task2': SPLIT_C_SUBSET_1_IDS,
        'task3': SPLIT_C_SUBSET_1_IDS+SPLIT_C_SUBSET_2_IDS
    }
}

"""
LABELS for currently known, unknown, and previously known classes for each split and task.
"""

PREV_KNOWN_CLASSES_LABELS = {
    'A': {
        'task1': None,
        'task2': SPLIT_A_SUBSET_1_LABELS,
        'task3': SPLIT_A_SUBSET_1_LABELS+SPLIT_A_SUBSET_2_LABELS
    },
    'B': {
        'task1': None,  
        'task2': SPLIT_B_SUBSET_1_LABELS,
        'task3': SPLIT_B_SUBSET_1_LABELS+SPLIT_B_SUBSET_2_LABELS
    },
    'C': {
        'task1': None,  
        'task2': SPLIT_C_SUBSET_1_LABELS,
        'task3': SPLIT_C_SUBSET_1_LABELS+SPLIT_C_SUBSET_2_LABELS
    }
}

KNOWN_CLASSES_LABELS = {
    'A' :{
        'task1': SPLIT_A_SUBSET_1_LABELS,
        'task2': SPLIT_A_SUBSET_2_LABELS,
        'task3': SPLIT_A_SUBSET_3_LABELS
    },
    'B': {
        'task1': SPLIT_B_SUBSET_1_LABELS,
        'task2': SPLIT_B_SUBSET_2_LABELS,
        'task3': SPLIT_B_SUBSET_3_LABELS
    },
    
    'C': {
        'task1': SPLIT_C_SUBSET_1_LABELS,
        'task2': SPLIT_C_SUBSET_2_LABELS,
        'task3': SPLIT_C_SUBSET_3_LABELS
    }
}

UNKNOWN_CLASSES_LABELS = {
    'A' :{
        'task1': SPLIT_A_SUBSET_2_LABELS+SPLIT_A_SUBSET_3_LABELS,
        'task2': SPLIT_A_SUBSET_3_LABELS,
        'task3': None
    },
    'B': {
        'task1': SPLIT_B_SUBSET_2_LABELS+SPLIT_B_SUBSET_3_LABELS,
        'task2': SPLIT_B_SUBSET_3_LABELS,
        'task3': None,
    },
    'C': {
        'task1': SPLIT_C_SUBSET_2_LABELS+SPLIT_C_SUBSET_3_LABELS,
        'task2': SPLIT_C_SUBSET_3_LABELS,
        'task3': None,
    }
}


datasets/scannet200/owis_splits.py
aminebdj/3D-OWIS
### ScanNet Benchmark constants ###
VALID_CLASS_IDS_20 = (1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 14, 16, 24, 28, 33, 34, 36, 39)

CLASS_LABELS_20 = ('wall', 'floor', 'cabinet', 'bed', 'chair', 'sofa', 'table', 'door', 'window',
                   'bookshelf', 'picture', 'counter', 'desk', 'curtain', 'refrigerator',
                   'shower curtain', 'toilet', 'sink', 'bathtub', 'otherfurniture')

SCANNET_COLOR_MAP_20 = {
    0: (0., 0., 0.),
    1: (174., 199., 232.),
    2: (152., 223., 138.),
    3: (31., 119., 180.),
    4: (255., 187., 120.),
    5: (188., 189., 34.),
    6: (140., 86., 75.),
    7: (255., 152., 150.),
    8: (214., 39., 40.),
    9: (197., 176., 213.),
    10: (148., 103., 189.),
    11: (196., 156., 148.),
    12: (23., 190., 207.),
    14: (247., 182., 210.),
    15: (66., 188., 102.),
    16: (219., 219., 141.),
    17: (140., 57., 197.),
    18: (202., 185., 52.),
    19: (51., 176., 203.),
    20: (200., 54., 131.),
    21: (92., 193., 61.),
    22: (78., 71., 183.),
    23: (172., 114., 82.),
    24: (255., 127., 14.),
    25: (91., 163., 138.),
    26: (153., 98., 156.),
    27: (140., 153., 101.),
    28: (158., 218., 229.),
    29: (100., 125., 154.),
    30: (178., 127., 135.),
    32: (146., 111., 194.),
    33: (44., 160., 44.),
    34: (112., 128., 144.),
    35: (96., 207., 209.),
    36: (227., 119., 194.),
    37: (213., 92., 176.),
    38: (94., 106., 211.),
    39: (82., 84., 163.),
    40: (100., 85., 144.),
}
### ScanNet200 Benchmark constants ###
VALID_CLASS_IDS_200 = (
1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 13, 14, 15, 16, 17, 18, 19, 21, 22, 23, 24, 26, 27, 28, 29, 31, 32, 33, 34, 35, 36, 38, 39, 40, 41, 42, 44, 45, 46, 47, 48, 49, 50, 51, 52, 54, 55, 56, 57, 58, 59, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71,
72, 73, 74, 75, 76, 77, 78, 79, 80, 82, 84, 86, 87, 88, 89, 90, 93, 95, 96, 97, 98, 99, 100, 101, 102, 103, 104, 105, 106, 107, 110, 112, 115, 116, 118, 120, 121, 122, 125, 128, 130, 131, 132, 134, 136, 138, 139, 140, 141, 145, 148, 154,
155, 156, 157, 159, 161, 163, 165, 166, 168, 169, 170, 177, 180, 185, 188, 191, 193, 195, 202, 208, 213, 214, 221, 229, 230, 232, 233, 242, 250, 261, 264, 276, 283, 286, 300, 304, 312, 323, 325, 331, 342, 356, 370, 392, 395, 399, 408, 417,
488, 540, 562, 570, 572, 581, 609, 748, 776, 1156, 1163, 1164, 1165, 1166, 1167, 1168, 1169, 1170, 1171, 1172, 1173, 1174, 1175, 1176, 1178, 1179, 1180, 1181, 1182, 1183, 1184, 1185, 1186, 1187, 1188, 1189, 1190, 1191,3000)

CLASS_LABELS_200 = (
'wall', 'chair', 'floor', 'table', 'door', 'couch', 'cabinet', 'shelf', 'desk', 'office chair', 'bed', 'pillow', 'sink', 'picture', 'window', 'toilet', 'bookshelf', 'monitor', 'curtain', 'book', 'armchair', 'coffee table', 'box',
'refrigerator', 'lamp', 'kitchen cabinet', 'towel', 'clothes', 'tv', 'nightstand', 'counter', 'dresser', 'stool', 'cushion', 'plant', 'ceiling', 'bathtub', 'end table', 'dining table', 'keyboard', 'bag', 'backpack', 'toilet paper',
'printer', 'tv stand', 'whiteboard', 'blanket', 'shower curtain', 'trash can', 'closet', 'stairs', 'microwave', 'stove', 'shoe', 'computer tower', 'bottle', 'bin', 'ottoman', 'bench', 'board', 'washing machine', 'mirror', 'copier',
'basket', 'sofa chair', 'file cabinet', 'fan', 'laptop', 'shower', 'paper', 'person', 'paper towel dispenser', 'oven', 'blinds', 'rack', 'plate', 'blackboard', 'piano', 'suitcase', 'rail', 'radiator', 'recycling bin', 'container',
'wardrobe', 'soap dispenser', 'telephone', 'bucket', 'clock', 'stand', 'light', 'laundry basket', 'pipe', 'clothes dryer', 'guitar', 'toilet paper holder', 'seat', 'speaker', 'column', 'bicycle', 'ladder', 'bathroom stall', 'shower wall',
'cup', 'jacket', 'storage bin', 'coffee maker', 'dishwasher', 'paper towel roll', 'machine', 'mat', 'windowsill', 'bar', 'toaster', 'bulletin board', 'ironing board', 'fireplace', 'soap dish', 'kitchen counter', 'doorframe',
'toilet paper dispenser', 'mini fridge', 'fire extinguisher', 'ball', 'hat', 'shower curtain rod', 'water cooler', 'paper cutter', 'tray', 'shower door', 'pillar', 'ledge', 'toaster oven', 'mouse', 'toilet seat cover dispenser',
'furniture', 'cart', 'storage container', 'scale', 'tissue box', 'light switch', 'crate', 'power outlet', 'decoration', 'sign', 'projector', 'closet door', 'vacuum cleaner', 'candle', 'plunger', 'stuffed animal', 'headphones', 'dish rack',
'broom', 'guitar case', 'range hood', 'dustpan', 'hair dryer', 'water bottle', 'handicap bar', 'purse', 'vent', 'shower floor', 'water pitcher', 'mailbox', 'bowl', 'paper bag', 'alarm clock', 'music stand', 'projector screen', 'divider',
'laundry detergent', 'bathroom counter', 'object', 'bathroom vanity', 'closet wall', 'laundry hamper', 'bathroom stall door', 'ceiling light', 'trash bin', 'dumbbell', 'stair rail', 'tube', 'bathroom cabinet', 'cd case', 'closet rod',
'coffee kettle', 'structure', 'shower head', 'keyboard piano', 'case of water bottles', 'coat rack', 'storage organizer', 'folded chair', 'fire alarm', 'power strip', 'calendar', 'poster', 'potted plant', 'luggage', 'mattress','unknown')

# CLASS_LABELS_200 = ('wall','tv stand','floor',
#  'curtain',
#  'blinds',
#  'shower curtain',
#  'bookshelf',
#  'tv',
#  'kitchen cabinet',
#  'pillow',
#  'lamp',
#  'dresser',
#  'monitor',
#  'object',
#  'ceiling',
#  'board',
#  'stove',
#  'closet wall',
#  'couch',
#  'office chair',
#  'kitchen counter',
#  'shower',
#  'closet',
#  'doorframe',
#  'sofa chair',
#  'mailbox',
#  'nightstand',
#  'washing machine',
#  'picture',
#  'book',
#  'sink',
#  'recycling bin',
#  'table',
#  'backpack',
#  'shower wall',
#  'toilet',
#  'copier',
#  'counter',
#  'stool',
#  'refrigerator',
#  'window',
#  'file cabinet',
#  'chair',
#  'plant',
#  'coffee table',
#  'stairs',
#  'armchair',
#  'cabinet',
#  'bathroom vanity',
#  'bathroom stall',
#  'mirror',
#  'blackboard',
#  'trash can',
#  'stair rail',
#  'box',
#  'towel',
#  'door',
#  'clothes',
#  'whiteboard',
#  'bed',
#  'bathtub',
#  'desk',
#  'wardrobe',
#  'clothes dryer',
#  'radiator',
#  'shelf',
#  'cushion',
#  'end table',
#  'dining table',
#  'keyboard',
#  'bag',
#  'toilet paper',
#  'printer',
#  'blanket',
#  'microwave',
#  'shoe',
#  'computer tower',
#  'bottle',
#  'bin',
#  'ottoman',
#  'bench',
#  'basket',
#  'fan',
#  'laptop',
#  'person',
#  'paper towel dispenser',
#  'oven',
#  'rack',
#  'piano',
#  'suitcase',
#  'rail',
#  'container',
#  'telephone',
#  'stand',
#  'light',
#  'laundry basket',
#  'pipe',
#  'seat',
#  'column',
#  'bicycle',
#  'ladder',
#  'jacket',
#  'storage bin',
#  'coffee maker',
#  'dishwasher',
#  'machine',
#  'mat',
#  'windowsill',
#  'bulletin board',
#  'fireplace',
#  'mini fridge',
#  'water cooler',
#  'shower door',
#  'pillar',
#  'ledge',
#  'furniture',
#  'cart',
#  'decoration',
#  'closet door',
#  'vacuum cleaner',
#  'dish rack',
#  'range hood',
#  'projector screen',
#  'divider',
#  'bathroom counter',
#  'laundry hamper',
#  'bathroom stall door',
#  'ceiling light',
#  'trash bin',
#  'bathroom cabinet',
#  'structure',
#  'storage organizer',
#  'potted plant',
#  'mattress',
#  'paper',
#  'plate',
#  'soap dispenser',
#  'bucket',
#  'clock',
#  'guitar',
#  'toilet paper holder',
#  'speaker',
#  'cup',
#  'paper towel roll',
#  'bar',
#  'toaster',
#  'ironing board',
#  'soap dish',
#  'toilet paper dispenser',
#  'fire extinguisher',
#  'ball',
#  'hat',
#  'shower curtain rod',
#  'paper cutter',
#  'tray',
#  'toaster oven',
#  'mouse',
#  'toilet seat cover dispenser',
#  'storage container',
#  'scale',
#  'tissue box',
#  'light switch',
#  'crate',
#  'power outlet',
#  'sign',
#  'projector',
#  'candle',
#  'plunger',
#  'stuffed animal',
#  'headphones',
#  'broom',
#  'guitar case',
#  'dustpan',
#  'hair dryer',
#  'water bottle',
#  'handicap bar',
#  'purse',
#  'vent',
#  'shower floor',
#  'water pitcher',
#  'bowl',
#  'paper bag',
#  'alarm clock',
#  'music stand',
#  'laundry detergent',
#  'dumbbell',
#  'tube',
#  'cd case',
#  'closet rod',
#  'coffee kettle',
#  'shower head',
#  'keyboard piano',
#  'case of water bottles',
#  'coat rack',
#  'folded chair',
#  'fire alarm',
#  'power strip',
#  'calendar',
#  'poster',
#  'luggage',
#  'unknown')

# VALID_CLASS_IDS_200 = (1,
#  51,
#  3,
#  21,
#  86,
#  55,
#  18,
#  33,
#  29,
#  13,
#  28,
#  36,
#  19,
#  1163,
#  41,
#  69,
#  62,
#  1165,
#  6,
#  10,
#  159,
#  78,
#  57,
#  161,
#  74,
#  540,
#  34,
#  70,
#  15,
#  22,
#  14,
#  97,
#  4,
#  48,
#  128,
#  17,
#  72,
#  35,
#  38,
#  27,
#  16,
#  75,
#  2,
#  40,
#  24,
#  58,
#  23,
#  7,
#  1164,
#  125,
#  71,
#  89,
#  56,
#  1171,
#  26,
#  31,
#  5,
#  32,
#  52,
#  11,
#  42,
#  9,
#  99,
#  110,
#  96,
#  8,
#  39,
#  44,
#  45,
#  46,
#  47,
#  49,
#  50,
#  54,
#  59,
#  63,
#  64,
#  65,
#  66,
#  67,
#  68,
#  73,
#  76,
#  77,
#  80,
#  82,
#  84,
#  87,
#  90,
#  93,
#  95,
#  98,
#  101,
#  104,
#  105,
#  106,
#  107,
#  116,
#  120,
#  121,
#  122,
#  131,
#  132,
#  134,
#  136,
#  139,
#  140,
#  141,
#  154,
#  156,
#  165,
#  177,
#  188,
#  191,
#  193,
#  213,
#  214,
#  250,
#  276,
#  283,
#  323,
#  342,
#  609,
#  748,
#  1156,
#  1166,
#  1167,
#  1168,
#  1169,
#  1173,
#  1178,
#  1183,
#  1189,
#  1191,
#  79,
#  88,
#  100,
#  102,
#  103,
#  112,
#  115,
#  118,
#  130,
#  138,
#  145,
#  148,
#  155,
#  157,
#  163,
#  166,
#  168,
#  169,
#  170,
#  180,
#  185,
#  195,
#  202,
#  208,
#  221,
#  229,
#  230,
#  232,
#  233,
#  242,
#  261,
#  264,
#  286,
#  300,
#  304,
#  312,
#  325,
#  331,
#  356,
#  370,
#  392,
#  395,
#  399,
#  408,
#  417,
#  488,
#  562,
#  570,
#  572,
#  581,
#  776,
#  1170,
#  1172,
#  1174,
#  1175,
#  1176,
#  1179,
#  1180,
#  1181,
#  1182,
#  1184,
#  1185,
#  1186,
#  1187,
#  1188,
#  1190,
#  3000)

SCANNET_COLOR_MAP_200 = {
0: (0., 0., 0.),
1: (174., 199., 232.),
2: (188., 189., 34.),
3: (152., 223., 138.),
4: (255., 152., 150.),
5: (214., 39., 40.),
6: (91., 135., 229.),
7: (31., 119., 180.),
8: (229., 91., 104.),
9: (247., 182., 210.),
10: (91., 229., 110.),
11: (255., 187., 120.),
13: (141., 91., 229.),
14: (112., 128., 144.),
15: (196., 156., 148.),
16: (197., 176., 213.),
17: (44., 160., 44.),
18: (148., 103., 189.),
19: (229., 91., 223.),
21: (219., 219., 141.),
22: (192., 229., 91.),
23: (88., 218., 137.),
24: (58., 98., 137.),
26: (177., 82., 239.),
27: (255., 127., 14.),
28: (237., 204., 37.),
29: (41., 206., 32.),
31: (62., 143., 148.),
32: (34., 14., 130.),
33: (143., 45., 115.),
34: (137., 63., 14.),
35: (23., 190., 207.),
36: (16., 212., 139.),
38: (90., 119., 201.),
39: (125., 30., 141.),
40: (150., 53., 56.),
41: (186., 197., 62.),
42: (227., 119., 194.),
44: (38., 100., 128.),
45: (120., 31., 243.),
46: (154., 59., 103.),
47: (169., 137., 78.),
48: (143., 245., 111.),
49: (37., 230., 205.),
50: (14., 16., 155.),
51: (196., 51., 182.),
52: (237., 80., 38.),
54: (138., 175., 62.),
55: (158., 218., 229.),
56: (38., 96., 167.),
57: (190., 77., 246.),
58: (208., 49., 84.),
59: (208., 193., 72.),
62: (55., 220., 57.),
63: (10., 125., 140.),
64: (76., 38., 202.),
65: (191., 28., 135.),
66: (211., 120., 42.),
67: (118., 174., 76.),
68: (17., 242., 171.),
69: (20., 65., 247.),
70: (208., 61., 222.),
71: (162., 62., 60.),
72: (210., 235., 62.),
73: (45., 152., 72.),
74: (35., 107., 149.),
75: (160., 89., 237.),
76: (227., 56., 125.),
77: (169., 143., 81.),
78: (42., 143., 20.),
79: (25., 160., 151.),
80: (82., 75., 227.),
82: (253., 59., 222.),
84: (240., 130., 89.),
86: (123., 172., 47.),
87: (71., 194., 133.),
88: (24., 94., 205.),
89: (134., 16., 179.),
90: (159., 32., 52.),
93: (213., 208., 88.),
95: (64., 158., 70.),
96: (18., 163., 194.),
97: (65., 29., 153.),
98: (177., 10., 109.),
99: (152., 83., 7.),
100: (83., 175., 30.),
101: (18., 199., 153.),
102: (61., 81., 208.),
103: (213., 85., 216.),
104: (170., 53., 42.),
105: (161., 192., 38.),
106: (23., 241., 91.),
107: (12., 103., 170.),
110: (151., 41., 245.),
112: (133., 51., 80.),
115: (184., 162., 91.),
116: (50., 138., 38.),
118: (31., 237., 236.),
120: (39., 19., 208.),
121: (223., 27., 180.),
122: (254., 141., 85.),
125: (97., 144., 39.),
128: (106., 231., 176.),
130: (12., 61., 162.),
131: (124., 66., 140.),
132: (137., 66., 73.),
134: (250., 253., 26.),
136: (55., 191., 73.),
138: (60., 126., 146.),
139: (153., 108., 234.),
140: (184., 58., 125.),
141: (135., 84., 14.),
145: (139., 248., 91.),
148: (53., 200., 172.),
154: (63., 69., 134.),
155: (190., 75., 186.),
156: (127., 63., 52.),
157: (141., 182., 25.),
159: (56., 144., 89.),
161: (64., 160., 250.),
163: (182., 86., 245.),
165: (139., 18., 53.),
166: (134., 120., 54.),
168: (49., 165., 42.),
169: (51., 128., 133.),
170: (44., 21., 163.),
177: (232., 93., 193.),
180: (176., 102., 54.),
185: (116., 217., 17.),
188: (54., 209., 150.),
191: (60., 99., 204.),
193: (129., 43., 144.),
195: (252., 100., 106.),
202: (187., 196., 73.),
208: (13., 158., 40.),
213: (52., 122., 152.),
214: (128., 76., 202.),
221: (187., 50., 115.),
229: (180., 141., 71.),
230: (77., 208., 35.),
232: (72., 183., 168.),
233: (97., 99., 203.),
242: (172., 22., 158.),
250: (155., 64., 40.),
261: (118., 159., 30.),
264: (69., 252., 148.),
276: (45., 103., 173.),
283: (111., 38., 149.),
286: (184., 9., 49.),
300: (188., 174., 67.),
304: (53., 206., 53.),
312: (97., 235., 252.),
323: (66., 32., 182.),
325: (236., 114., 195.),
331: (241., 154., 83.),
342: (133., 240., 52.),
356: (16., 205., 144.),
370: (75., 101., 198.),
392: (237., 95., 251.),
395: (191., 52., 49.),
399: (227., 254., 54.),
408: (49., 206., 87.),
417: (48., 113., 150.),
488: (125., 73., 182.),
540: (229., 32., 114.),
562: (158., 119., 28.),
570: (60., 205., 27.),
572: (18., 215., 201.),
581: (79., 76., 153.),
609: (134., 13., 116.),
748: (192., 97., 63.),
776: (108., 163., 18.),
1156: (95., 220., 156.),
1163: (98., 141., 208.),
1164: (144., 19., 193.),
1165: (166., 36., 57.),
1166: (212., 202., 34.),
1167: (23., 206., 34.),
1168: (91., 211., 236.),
1169: (79., 55., 137.),
1170: (182., 19., 117.),
1171: (134., 76., 14.),
1172: (87., 185., 28.),
1173: (82., 224., 187.),
1174: (92., 110., 214.),
1175: (168., 80., 171.),
1176: (197., 63., 51.),
1178: (175., 199., 77.),
1179: (62., 180., 98.),
1180: (8., 91., 150.),
1181: (77., 15., 130.),
1182: (154., 65., 96.),
1183: (197., 152., 11.),
1184: (59., 155., 45.),
1185: (12., 147., 145.),
1186: (54., 35., 219.),
1187: (210., 73., 181.),
1188: (221., 124., 77.),
1189: (149., 214., 66.),
1190: (72., 185., 134.),
1191: (42., 94., 198.),
3000: (100., 100., 100.),
}

### For instance segmentation the non-object categories ###
VALID_PANOPTIC_IDS = (1, 3)

CLASS_LABELS_PANOPTIC = ('wall', 'floor')


MAP_STRING_TO_ID = {'wall': 0,
 'chair': 1,
 'floor': 2,
 'table': 3,
 'door': 4,
 'couch': 5,
 'cabinet': 6,
 'shelf': 7,
 'desk': 8,
 'office chair': 9,
 'bed': 10,
 'pillow': 11,
 'sink': 12,
 'picture': 13,
 'window': 14,
 'toilet': 15,
 'bookshelf': 16,
 'monitor': 17,
 'curtain': 18,
 'book': 19,
 'armchair': 20,
 'coffee table': 21,
 'box': 22,
 'refrigerator': 23,
 'lamp': 24,
 'kitchen cabinet': 25,
 'towel': 26,
 'clothes': 27,
 'tv': 28,
 'nightstand': 29,
 'counter': 30,
 'dresser': 31,
 'stool': 32,
 'cushion': 33,
 'plant': 34,
 'ceiling': 35,
 'bathtub': 36,
 'end table': 37,
 'dining table': 38,
 'keyboard': 39,
 'bag': 40,
 'backpack': 41,
 'toilet paper': 42,
 'printer': 43,
 'tv stand': 44,
 'whiteboard': 45,
 'blanket': 46,
 'shower curtain': 47,
 'trash can': 48,
 'closet': 49,
 'stairs': 50,
 'microwave': 51,
 'stove': 52,
 'shoe': 53,
 'computer tower': 54,
 'bottle': 55,
 'bin': 56,
 'ottoman': 57,
 'bench': 58,
 'board': 59,
 'washing machine': 60,
 'mirror': 61,
 'copier': 62,
 'basket': 63,
 'sofa chair': 64,
 'file cabinet': 65,
 'fan': 66,
 'laptop': 67,
 'shower': 68,
 'paper': 69,
 'person': 70,
 'paper towel dispenser': 71,
 'oven': 72,
 'blinds': 73,
 'rack': 74,
 'plate': 75,
 'blackboard': 76,
 'piano': 77,
 'suitcase': 78,
 'rail': 79,
 'radiator': 80,
 'recycling bin': 81,
 'container': 82,
 'wardrobe': 83,
 'soap dispenser': 84,
 'telephone': 85,
 'bucket': 86,
 'clock': 87,
 'stand': 88,
 'light': 89,
 'laundry basket': 90,
 'pipe': 91,
 'clothes dryer': 92,
 'guitar': 93,
 'toilet paper holder': 94,
 'seat': 95,
 'speaker': 96,
 'column': 97,
 'bicycle': 98,
 'ladder': 99,
 'bathroom stall': 100,
 'shower wall': 101,
 'cup': 102,
 'jacket': 103,
 'storage bin': 104,
 'coffee maker': 105,
 'dishwasher': 106,
 'paper towel roll': 107,
 'machine': 108,
 'mat': 109,
 'windowsill': 110,
 'bar': 111,
 'toaster': 112,
 'bulletin board': 113,
 'ironing board': 114,
 'fireplace': 115,
 'soap dish': 116,
 'kitchen counter': 117,
 'doorframe': 118,
 'toilet paper dispenser': 119,
 'mini fridge': 120,
 'fire extinguisher': 121,
 'ball': 122,
 'hat': 123,
 'shower curtain rod': 124,
 'water cooler': 125,
 'paper cutter': 126,
 'tray': 127,
 'shower door': 128,
 'pillar': 129,
 'ledge': 130,
 'toaster oven': 131,
 'mouse': 132,
 'toilet seat cover dispenser': 133,
 'furniture': 134,
 'cart': 135,
 'storage container': 136,
 'scale': 137,
 'tissue box': 138,
 'light switch': 139,
 'crate': 140,
 'power outlet': 141,
 'decoration': 142,
 'sign': 143,
 'projector': 144,
 'closet door': 145,
 'vacuum cleaner': 146,
 'candle': 147,
 'plunger': 148,
 'stuffed animal': 149,
 'headphones': 150,
 'dish rack': 151,
 'broom': 152,
 'guitar case': 153,
 'range hood': 154,
 'dustpan': 155,
 'hair dryer': 156,
 'water bottle': 157,
 'handicap bar': 158,
 'purse': 159,
 'vent': 160,
 'shower floor': 161,
 'water pitcher': 162,
 'mailbox': 163,
 'bowl': 164,
 'paper bag': 165,
 'alarm clock': 166,
 'music stand': 167,
 'projector screen': 168,
 'divider': 169,
 'laundry detergent': 170,
 'bathroom counter': 171,
 'object': 172,
 'bathroom vanity': 173,
 'closet wall': 174,
 'laundry hamper': 175,
 'bathroom stall door': 176,
 'ceiling light': 177,
 'trash bin': 178,
 'dumbbell': 179,
 'stair rail': 180,
 'tube': 181,
 'bathroom cabinet': 182,
 'cd case': 183,
 'closet rod': 184,
 'coffee kettle': 185,
 'structure': 186,
 'shower head': 187,
 'keyboard piano': 188,
 'case of water bottles': 189,
 'coat rack': 190,
 'storage organizer': 191,
 'folded chair': 192,
 'fire alarm': 193,
 'power strip': 194,
 'calendar': 195,
 'poster': 196,
 'potted plant': 197,
 'luggage': 198,
 'mattress': 199,
 'unknown': 200}

datasets/scannet200/scannet200_constants.py
aminebdj/3D-OWIS
### This file contains the HEAD - COMMON - TAIL split category ids for ScanNet 200

HEAD_CATS_SCANNET_200 = ['tv stand', 'curtain', 'blinds', 'shower curtain', 'bookshelf', 'tv', 'kitchen cabinet', 'pillow', 'lamp', 'dresser', 'monitor', 'object', 'ceiling', 'board', 'stove', 'closet wall', 'couch', 'office chair', 'kitchen counter', 'shower', 'closet', 'doorframe', 'sofa chair', 'mailbox', 'nightstand', 'washing machine', 'picture', 'book', 'sink', 'recycling bin', 'table', 'backpack', 'shower wall', 'toilet', 'copier', 'counter', 'stool', 'refrigerator', 'window', 'file cabinet', 'chair', 'wall', 'plant', 'coffee table', 'stairs', 'armchair', 'cabinet', 'bathroom vanity', 'bathroom stall', 'mirror', 'blackboard', 'trash can', 'stair rail', 'box', 'towel', 'door', 'clothes', 'whiteboard', 'bed', 'floor', 'bathtub', 'desk', 'wardrobe', 'clothes dryer', 'radiator', 'shelf']
COMMON_CATS_SCANNET_200 = ["cushion", "end table", "dining table", "keyboard", "bag", "toilet paper", "printer", "blanket", "microwave", "shoe", "computer tower", "bottle", "bin", "ottoman", "bench", "basket", "fan", "laptop", "person", "paper towel dispenser", "oven", "rack", "piano", "suitcase", "rail", "container", "telephone", "stand", "light", "laundry basket", "pipe", "seat", "column", "bicycle", "ladder", "jacket", "storage bin", "coffee maker", "dishwasher", "machine", "mat", "windowsill", "bulletin board", "fireplace", "mini fridge", "water cooler", "shower door", "pillar", "ledge", "furniture", "cart", "decoration", "closet door", "vacuum cleaner", "dish rack", "range hood", "projector screen", "divider", "bathroom counter", "laundry hamper", "bathroom stall door", "ceiling light", "trash bin", "bathroom cabinet", "structure", "storage organizer", "potted plant", "mattress"]
TAIL_CATS_SCANNET_200 = ["paper", "plate", "soap dispenser", "bucket", "clock", "guitar", "toilet paper holder", "speaker", "cup", "paper towel roll", "bar", "toaster", "ironing board", "soap dish", "toilet paper dispenser", "fire extinguisher", "ball", "hat", "shower curtain rod", "paper cutter", "tray", "toaster oven", "mouse", "toilet seat cover dispenser", "storage container", "scale", "tissue box", "light switch", "crate", "power outlet", "sign", "projector", "candle", "plunger", "stuffed animal", "headphones", "broom", "guitar case", "dustpan", "hair dryer", "water bottle", "handicap bar", "purse", "vent", "shower floor", "water pitcher", "bowl", "paper bag", "alarm clock", "music stand", "laundry detergent", "dumbbell", "tube", "cd case", "closet rod", "coffee kettle", "shower head", "keyboard piano", "case of water bottles", "coat rack", "folded chair", "fire alarm", "power strip", "calendar", "poster", "luggage"]


### Given the different size of the official train and val sets, not all ScanNet200 categories are present in the validation set.
### Here we list of categories with labels and IDs present in both train and validation set, and the remaining categories those are present in train, but not in val
### We dont evaluate on unseen validation categories in this benchmark

VALID_CLASS_IDS_200_VALIDATION = ('wall', 'chair', 'floor', 'table', 'door', 'couch', 'cabinet', 'shelf', 'desk', 'office chair', 'bed', 'pillow', 'sink', 'picture', 'window', 'toilet', 'bookshelf', 'monitor', 'curtain', 'book', 'armchair', 'coffee table', 'box', 'refrigerator', 'lamp', 'kitchen cabinet', 'towel', 'clothes', 'tv', 'nightstand', 'counter', 'dresser', 'stool', 'cushion', 'plant', 'ceiling', 'bathtub', 'end table', 'dining table', 'keyboard', 'bag', 'backpack', 'toilet paper', 'printer', 'tv stand', 'whiteboard', 'blanket', 'shower curtain', 'trash can', 'closet', 'stairs', 'microwave', 'stove', 'shoe', 'computer tower', 'bottle', 'bin', 'ottoman', 'bench', 'board', 'washing machine', 'mirror', 'copier', 'basket', 'sofa chair', 'file cabinet', 'fan', 'laptop', 'shower', 'paper', 'person', 'paper towel dispenser', 'oven', 'blinds', 'rack', 'plate', 'blackboard', 'piano', 'suitcase', 'rail', 'radiator', 'recycling bin', 'container', 'wardrobe', 'soap dispenser', 'telephone', 'bucket', 'clock', 'stand', 'light', 'laundry basket', 'pipe', 'clothes dryer', 'guitar', 'toilet paper holder', 'seat', 'speaker', 'column', 'ladder', 'bathroom stall', 'shower wall', 'cup', 'jacket', 'storage bin', 'coffee maker', 'dishwasher', 'paper towel roll', 'machine', 'mat', 'windowsill', 'bar', 'toaster', 'bulletin board', 'ironing board', 'fireplace', 'soap dish', 'kitchen counter', 'doorframe', 'toilet paper dispenser', 'mini fridge', 'fire extinguisher', 'ball', 'hat', 'shower curtain rod', 'water cooler', 'paper cutter', 'tray', 'shower door', 'pillar', 'ledge', 'toaster oven', 'mouse', 'toilet seat cover dispenser', 'furniture', 'cart', 'scale', 'tissue box', 'light switch', 'crate', 'power outlet', 'decoration', 'sign', 'projector', 'closet door', 'vacuum cleaner', 'plunger', 'stuffed animal', 'headphones', 'dish rack', 'broom', 'range hood', 'dustpan', 'hair dryer', 'water bottle', 'handicap bar', 'vent', 'shower floor', 'water pitcher', 'mailbox', 'bowl', 'paper bag', 'projector screen', 'divider', 'laundry detergent', 'bathroom counter', 'object', 'bathroom vanity', 'closet wall', 'laundry hamper', 'bathroom stall door', 'ceiling light', 'trash bin', 'dumbbell', 'stair rail', 'tube', 'bathroom cabinet', 'closet rod', 'coffee kettle', 'shower head', 'keyboard piano', 'case of water bottles', 'coat rack', 'folded chair', 'fire alarm', 'power strip', 'calendar', 'poster', 'potted plant', 'mattress','unknown')

CLASS_LABELS_200_VALIDATION = (1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 13, 14, 15, 16, 17, 18, 19, 21, 22, 23, 24, 26, 27, 28, 29, 31, 32, 33, 34, 35, 36, 38, 39, 40, 41, 42, 44, 45, 46, 47, 48, 49, 50, 51, 52, 54, 55, 56, 57, 58, 59, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71, 72, 73, 74, 75, 76, 77, 78, 79, 80, 82, 84, 86, 87, 88, 89, 90, 93, 95, 96, 97, 98, 99, 100, 101, 102, 103, 104, 105, 106, 107, 110, 112, 115, 116, 118, 120, 122, 125, 128, 130, 131, 132, 134, 136, 138, 139, 140, 141, 145, 148, 154, 155, 156, 157, 159, 161, 163, 165, 166, 168, 169, 170, 177, 180, 185, 188, 191, 193, 195, 202, 208, 213, 214, 229, 230, 232, 233, 242, 250, 261, 264, 276, 283, 300, 304, 312, 323, 325, 342, 356, 370, 392, 395, 408, 417, 488, 540, 562, 570, 609, 748, 776, 1156, 1163, 1164, 1165, 1166, 1167, 1168, 1169, 1170, 1171, 1172, 1173, 1175, 1176, 1179, 1180, 1181, 1182, 1184, 1185, 1186, 1187, 1188, 1189, 1191,3000)

VALID_CLASS_IDS_200_TRAIN_ONLY = ('bicycle', 'storage container', 'candle', 'guitar case', 'purse', 'alarm clock', 'music stand', 'cd case', 'structure', 'storage organizer', 'luggage')

CLASS_LABELS_200_TRAIN_ONLY = (121,  221,  286,  331,  399,  572,  581, 1174, 1178, 1183, 1190)
 
HEAD_PRED_IDS = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 32, 33, 34, 39, 42, 43, 45, 46, 47, 48, 50, 57, 58, 59, 60, 62, 63, 66, 71, 74, 78, 79, 81, 90, 98, 99, 115, 116, 161, 170, 171, 172, 178]
COMMON_PRED_IDS = [31, 35, 36, 37, 38, 40, 41, 44, 49, 51, 52, 53, 54, 55, 56, 61, 64, 65, 68, 69, 70, 72, 75, 76, 77, 80, 83, 86, 87, 88, 89, 93, 95, 96, 97, 101, 102, 103, 104, 106, 107, 108, 111, 113, 118, 123, 126, 127, 128, 132, 133, 140, 143, 144, 149, 152, 166, 167, 169, 173, 174, 175, 176, 180, 184, 189, 195, 197]
TAIL_PRED_IDS = [67, 73, 82, 84, 85, 91, 92, 94, 100, 105, 109, 110, 112, 114, 117, 119, 120, 121, 122, 124, 125, 129, 130, 131, 134, 135, 136, 137, 138, 139, 141, 142, 145, 146, 147, 148, 150, 151, 153, 154, 155, 156, 157, 158, 159, 160, 162, 163, 164, 165, 168, 177, 179, 181, 182, 183, 185, 186, 187, 188, 190, 191, 192, 193, 194, 196]

datasets/scannet200/scannet200_splits.py
aminebdj/3D-OWIS
import logging
from itertools import product
from pathlib import Path
from random import random, sample, uniform
from typing import List, Optional, Tuple, Union
from random import choice
from copy import deepcopy
from random import randrange


import numpy
import torch
from datasets.random_cuboid import RandomCuboid

import albumentations as A
import numpy as np
import scipy
import volumentations as V
import yaml
# from yaml import CLoader as Loader
from torch.utils.data import Dataset
from datasets.scannet200.scannet200_constants import SCANNET_COLOR_MAP_200, SCANNET_COLOR_MAP_200, CLASS_LABELS_200
logger = logging.getLogger(__name__)


class SemanticSegmentationDataset(Dataset):
    """Docstring for SemanticSegmentationDataset. """

    def __init__(
        self,
        dataset_name="scannet",
        data_dir: Optional[Union[str, Tuple[str]]] = "data/processed/scannet",
        label_db_filepath: Optional[
            str
        ] = "configs/scannet_preprocessing/label_database.yaml",
        # mean std values from scannet
        color_mean_std: Optional[Union[str, Tuple[Tuple[float]]]] = (
            (0.47793125906962, 0.4303257521323044, 0.3749598901421883),
            (0.2834475483823543, 0.27566157565723015, 0.27018971370874995),
        ),
        mode: Optional[str] = "train",
        add_colors: Optional[bool] = True,
        add_normals: Optional[bool] = True,
        add_raw_coordinates: Optional[bool] = False,
        add_instance: Optional[bool] = False,
        num_labels: Optional[int] = -1,
        data_percent: Optional[float] = 1.0,
        ignore_label: Optional[Union[int, Tuple[int]]] = 255,
        volume_augmentations_path: Optional[str] = None,
        image_augmentations_path: Optional[str] = None,
        instance_oversampling=0,
        place_around_existing=False,
        max_cut_region=0,
        point_per_cut=100,
        flip_in_center=False,
        noise_rate=0.0,
        resample_points=0.0,
        cache_data=False,
        add_unlabeled_pc=False,
        task="instance_segmentation",
        cropping=False,
        cropping_args=None,
        is_tta=False,
        crop_min_size=20000,
        crop_length=6.0,
        cropping_v1=True,
        reps_per_epoch=1,
        area=-1,
        on_crops=False,
        eval_inner_core=-1,
        filter_out_classes=[],
        label_offset=0,
        add_clip=False,
        is_elastic_distortion=True,
        color_drop=0.0,
        finetune=False,
        Num_examplars=10,
        OW_task='task1',
        split='A'
    ):
        assert task in ["instance_segmentation", "semantic_segmentation"], "unknown task"

        self.add_clip = add_clip
        self.dataset_name = dataset_name
        self.is_elastic_distortion = is_elastic_distortion
        self.color_drop = color_drop

        if self.dataset_name == "scannet":
            self.color_map = SCANNET_COLOR_MAP_20
            self.color_map[255] = (255, 255, 255)
        elif self.dataset_name == "stpls3d":
            self.color_map = {
                0:[0, 255, 0],  # Ground
                1:[0, 0, 255],  # Build
                2:[0, 255, 255],  # LowVeg
                3:[255, 255, 0],  # MediumVeg
                4:[255, 0, 255],  # HiVeg
                5:[100, 100, 255],  # Vehicle
                6:[200, 200, 100],  # Truck
                7:[170, 120, 200],  # Aircraft
                8:[255, 0, 0],  # MilitaryVec
                9:[200, 100, 100],  # Bike
                10:[10, 200, 100],  # Motorcycle
                11:[200, 200, 200],  # LightPole
                12:[50, 50, 50],  # StreetSign
                13:[60, 130, 60],  # Clutter
                14:[130, 30, 60]}  # Fence
        elif self.dataset_name == "scannet200":
            self.color_map = SCANNET_COLOR_MAP_200
        elif self.dataset_name == "s3dis":
            self.color_map = {
            0: [0, 255, 0],        # ceiling
            1: [0, 0, 255],        # floor
            2: [0, 255, 255],      # wall
            3: [255, 255, 0],      # beam
            4: [255, 0, 255],      # column
            5: [100, 100, 255],    # window
            6: [200, 200, 100],    # door
            7: [170, 120, 200],    # table
            8: [255, 0, 0],        # chair
            9: [200, 100, 100],    # sofa
            10: [10, 200, 100],     # bookcase
            11: [200, 200, 200],    # board
            12: [50, 50, 50]      # clutter
            }
        else:
            assert False, "dataset not known"

        self.task = task

        self.filter_out_classes = filter_out_classes
        self.label_offset = label_offset

        self.area = area
        self.eval_inner_core = eval_inner_core

        self.reps_per_epoch = reps_per_epoch

        self.cropping = cropping
        self.cropping_args = cropping_args
        self.is_tta = is_tta
        self.on_crops = on_crops

        self.crop_min_size = crop_min_size
        self.crop_length = crop_length

        self.version1 = cropping_v1

        self.random_cuboid = RandomCuboid(self.crop_min_size,
                                          crop_length=self.crop_length,
                                          version1=self.version1)

        self.mode = mode
        self.data_dir = data_dir
        self.add_unlabeled_pc = add_unlabeled_pc
        if add_unlabeled_pc:
            self.other_database = self._load_yaml(
                Path(data_dir).parent / "matterport" / "train_database.yaml"
            )
        if type(data_dir) == str:
            self.data_dir = [self.data_dir]
        self.ignore_label = ignore_label
        self.add_colors = add_colors
        self.add_normals = add_normals
        self.add_instance = add_instance
        self.add_raw_coordinates = add_raw_coordinates
        self.instance_oversampling = instance_oversampling
        self.place_around_existing = place_around_existing
        self.max_cut_region = max_cut_region
        self.point_per_cut = point_per_cut
        self.flip_in_center = flip_in_center
        self.noise_rate = noise_rate
        self.resample_points = resample_points
        # loading database files
        self._data = []
        for database_path in self.data_dir:
            database_path = Path(database_path)
            if self.dataset_name != "s3dis":
                if not finetune or mode == 'validation':
                    if not (database_path / f"{mode}_database.yaml").exists():
                        print(f"generate {database_path}/{mode}_database.yaml first")
                        exit()
                    
                    self._data.extend(self._load_yaml(database_path / f"{mode}_database.yaml"))
                else:
                    if not (database_path/ f"exemplars/{split}/{OW_task}/train_database_{Num_examplars}_exemplars.yaml").exists():
                        print(f"generate {database_path}/exemplars/{split}/{OW_task}/train_database_{Num_examplars}_exemplars.yaml first")
                        exit()
                    
                    self._data.extend(self._load_yaml(database_path/ f"exemplars/{split}/{OW_task}/train_database_{Num_examplars}_exemplars.yaml"))
            else:
                mode_s3dis = f"Area_{self.area}"
                if self.mode == "train":
                    mode_s3dis = "train_" + mode_s3dis
                if not (database_path / f"{mode_s3dis}_database.yaml").exists():
                    print(f"generate {database_path}/{mode_s3dis}_database.yaml first")
                    exit()
                self._data.extend(self._load_yaml(database_path / f"{mode_s3dis}_database.yaml"))
        if data_percent < 1.0:
            self._data = sample(self._data, int(len(self._data) * data_percent))
        labels = self._load_yaml(Path(label_db_filepath))

        # if working only on classes for validation - discard others
        self._labels = self._select_correct_labels(labels, num_labels)

        if instance_oversampling > 0:
            self.instance_data = self._load_yaml(
                Path(label_db_filepath).parent / "instance_database.yaml"
            )

        # normalize color channels
        if self.dataset_name == "s3dis":
            color_mean_std = color_mean_std.replace("color_mean_std.yaml", f"Area_{self.area}_color_mean_std.yaml")

        if Path(str(color_mean_std)).exists():
            color_mean_std = self._load_yaml(color_mean_std)
            color_mean, color_std = (
                tuple(color_mean_std["mean"]),
                tuple(color_mean_std["std"]),
            )
        elif len(color_mean_std[0]) == 3 and len(color_mean_std[1]) == 3:
            color_mean, color_std = color_mean_std[0], color_mean_std[1]
        else:
            logger.error("pass mean and std as tuple of tuples, or as an .yaml file")

        # augmentations
        self.volume_augmentations = V.NoOp()
        if (volume_augmentations_path is not None) and (
            volume_augmentations_path != "none"
        ):
            self.volume_augmentations = V.load(
                Path(volume_augmentations_path), data_format="yaml"
            )
        self.image_augmentations = A.NoOp()
        if (image_augmentations_path is not None) and (
            image_augmentations_path != "none"
        ):
            self.image_augmentations = A.load(
                Path(image_augmentations_path), data_format="yaml"
            )
        # mandatory color augmentation
        if add_colors:
            self.normalize_color = A.Normalize(mean=color_mean, std=color_std)

        self.cache_data = cache_data
        # new_data = []
        if self.cache_data:
            new_data = []
            for i in range(len(self._data)):
                self._data[i]['data'] = np.load(self.data[i]["filepath"].replace("../../", ""))
                if self.on_crops:
                    if self.eval_inner_core == -1:
                        for block_id, block in enumerate(self.splitPointCloud(self._data[i]['data'])):
                            if len(block) > 10000:
                                new_data.append({
                                    'instance_gt_filepath': self._data[i]['instance_gt_filepath'][block_id] \
                                        if len(self._data[i]['instance_gt_filepath']) > 0 else list(),
                                    'scene': f"{self._data[i]['scene'].replace('.txt', '')}_{block_id}.txt",
                                    'raw_filepath': f"{self.data[i]['filepath'].replace('.npy', '')}_{block_id}",
                                    'data': block
                                })
                            else:
                                assert False
                    else:
                        conds_inner, blocks_outer = self.splitPointCloud(self._data[i]['data'],
                                                                         size=self.crop_length,
                                                                         inner_core=self.eval_inner_core)

                        for block_id in range(len(conds_inner)):
                            cond_inner = conds_inner[block_id]
                            block_outer = blocks_outer[block_id]

                            if cond_inner.sum() > 10000:
                                new_data.append({
                                    'instance_gt_filepath': self._data[i]['instance_gt_filepath'][block_id] \
                                        if len(self._data[i]['instance_gt_filepath']) > 0 else list(),
                                    'scene': f"{self._data[i]['scene'].replace('.txt', '')}_{block_id}.txt",
                                    'raw_filepath': f"{self.data[i]['filepath'].replace('.npy', '')}_{block_id}",
                                    'data': block_outer,
                                    'cond_inner': cond_inner
                                })
                            else:
                                assert False

            if self.on_crops:
                self._data = new_data
                #new_data.append(np.load(self.data[i]["filepath"].replace("../../", "")))
            #self._data = new_data

    def splitPointCloud(self, cloud, size=50.0, stride=50, inner_core=-1):
        if inner_core == -1:
            limitMax = np.amax(cloud[:, 0:3], axis=0)
            width = int(np.ceil((limitMax[0] - size) / stride)) + 1
            depth = int(np.ceil((limitMax[1] - size) / stride)) + 1
            cells = [(x * stride, y * stride) for x in range(width) for y in range(depth)]
            blocks = []
            for (x, y) in cells:
                xcond = (cloud[:, 0] <= x + size) & (cloud[:, 0] >= x)
                ycond = (cloud[:, 1] <= y + size) & (cloud[:, 1] >= y)
                cond = xcond & ycond
                block = cloud[cond, :]
                blocks.append(block)
            return blocks
        else:
            limitMax = np.amax(cloud[:, 0:3], axis=0)
            width = int(np.ceil((limitMax[0] - inner_core) / stride)) + 1
            depth = int(np.ceil((limitMax[1] - inner_core) / stride)) + 1
            cells = [(x * stride, y * stride) for x in range(width) for y in range(depth)]
            blocks_outer = []
            conds_inner = []
            for (x, y) in cells:
                xcond_outer = (cloud[:, 0] <= x + inner_core / 2. + size / 2) & (cloud[:, 0] >= x + inner_core / 2. - size / 2)
                ycond_outer = (cloud[:, 1] <= y + inner_core / 2. + size / 2) & (cloud[:, 1] >= y + inner_core / 2. - size / 2)

                cond_outer = xcond_outer & ycond_outer
                block_outer = cloud[cond_outer, :]

                xcond_inner = (block_outer[:, 0] <= x + inner_core) & (block_outer[:, 0] >= x)
                ycond_inner = (block_outer[:, 1] <= y + inner_core) & (block_outer[:, 1] >= y)

                cond_inner = xcond_inner & ycond_inner

                conds_inner.append(cond_inner)
                blocks_outer.append(block_outer)
            return conds_inner, blocks_outer

    def map2color(self, labels):
        output_colors = list()

        for label in labels:
            output_colors.append(self.color_map[label])

        return torch.tensor(output_colors)

    def __len__(self):
        if self.is_tta:
            return 5*len(self.data)
        else:
            return self.reps_per_epoch*len(self.data)

    def __getitem__(self, idx: int):
        idx = idx % len(self.data)
        if self.is_tta:
            idx = idx % len(self.data)

        if self.cache_data:
            points = self.data[idx]['data']
        else:
            assert not self.on_crops, "you need caching if on crops"
            points = np.load(self.data[idx]["filepath"].replace("../../", ""))

        if "train" in self.mode and self.dataset_name in ["s3dis", "stpls3d"]:
            inds = self.random_cuboid(points)
            points = points[inds]

        coordinates, color, normals, segments, labels = (
            points[:, :3],
            points[:, 3:6],
            points[:, 6:9],
            points[:, 9],
            points[:, 10:12],
        )

        raw_coordinates = coordinates.copy()
        raw_color = color
        raw_normals = normals

        if not self.add_colors:
            color = np.ones((len(color), 3))

        # volume and image augmentations for train
        if "train" in self.mode or self.is_tta:
            if self.cropping:
                new_idx = self.random_cuboid(coordinates, labels[:,1],
                                             self._remap_from_zero(labels[:, 0].copy()))

                coordinates = coordinates[new_idx]
                color = color[new_idx]
                labels = labels[new_idx]
                segments = segments[new_idx]
                raw_color = raw_color[new_idx]
                raw_normals = raw_normals[new_idx]
                normals = normals[new_idx]
                points = points[new_idx]


            coordinates -= coordinates.mean(0)

            try:
                coordinates += np.random.uniform(coordinates.min(0), coordinates.max(0)) / 2
            except OverflowError as err:
                print(coordinates)
                print(coordinates.shape)
                raise err

            if self.instance_oversampling > 0.0:
                coordinates, color, normals, labels = self.augment_individual_instance(
                    coordinates, color, normals, labels, self.instance_oversampling
                )

            if self.flip_in_center:
                coordinates = flip_in_center(coordinates)

            for i in (0, 1):
                if random() < 0.5:
                    coord_max = np.max(points[:, i])
                    coordinates[:, i] = coord_max - coordinates[:, i]

            if random() < 0.95:
                if self.is_elastic_distortion:
                    for granularity, magnitude in ((0.2, 0.4), (0.8, 1.6)):
                        coordinates = elastic_distortion(
                            coordinates, granularity, magnitude
                        )
            aug = self.volume_augmentations(
                points=coordinates, normals=normals, features=color, labels=labels,
            )
            coordinates, color, normals, labels = (
                aug["points"],
                aug["features"],
                aug["normals"],
                aug["labels"],
            )
            pseudo_image = color.astype(np.uint8)[np.newaxis, :, :]
            color = np.squeeze(self.image_augmentations(image=pseudo_image)["image"])

            if self.point_per_cut != 0:
                number_of_cuts = int(len(coordinates) / self.point_per_cut)
                for _ in range(number_of_cuts):
                    size_of_cut = np.random.uniform(0.05, self.max_cut_region)
                    # not wall, floor or empty
                    point = choice(coordinates)
                    x_min = point[0] - size_of_cut
                    x_max = x_min + size_of_cut
                    y_min = point[1] - size_of_cut
                    y_max = y_min + size_of_cut
                    z_min = point[2] - size_of_cut
                    z_max = z_min + size_of_cut
                    indexes = crop(
                        coordinates, x_min, y_min, z_min, x_max, y_max, z_max
                    )
                    coordinates, normals, color, labels = (
                        coordinates[~indexes],
                        normals[~indexes],
                        color[~indexes],
                        labels[~indexes],
                    )

            # if self.noise_rate > 0:
            #     coordinates, color, normals, labels = random_points(
            #         coordinates,
            #         color,
            #         normals,
            #         labels,
            #         self.noise_rate,
            #         self.ignore_label,
            #     )

            if (self.resample_points > 0) or (self.noise_rate > 0):
                coordinates, color, normals, labels = random_around_points(
                    coordinates,
                    color,
                    normals,
                    labels,
                    self.resample_points,
                    self.noise_rate,
                    self.ignore_label,
                )

            if self.add_unlabeled_pc:
                if random() < 0.8:
                    new_points = np.load(
                        self.other_database[
                            np.random.randint(0, len(self.other_database) - 1)
                        ]["filepath"]
                    )
                    (
                        unlabeled_coords,
                        unlabeled_color,
                        unlabeled_normals,
                        unlabeled_labels,
                    ) = (
                        new_points[:, :3],
                        new_points[:, 3:6],
                        new_points[:, 6:9],
                        new_points[:, 9:],
                    )
                    unlabeled_coords -= unlabeled_coords.mean(0)
                    unlabeled_coords += (
                        np.random.uniform(
                            unlabeled_coords.min(0), unlabeled_coords.max(0)
                        )
                        / 2
                    )

                    aug = self.volume_augmentations(
                        points=unlabeled_coords,
                        normals=unlabeled_normals,
                        features=unlabeled_color,
                        labels=unlabeled_labels,
                    )
                    (
                        unlabeled_coords,
                        unlabeled_color,
                        unlabeled_normals,
                        unlabeled_labels,
                    ) = (
                        aug["points"],
                        aug["features"],
                        aug["normals"],
                        aug["labels"],
                    )
                    pseudo_image = unlabeled_color.astype(np.uint8)[np.newaxis, :, :]
                    unlabeled_color = np.squeeze(
                        self.image_augmentations(image=pseudo_image)["image"]
                    )

                    coordinates = np.concatenate((coordinates, unlabeled_coords))
                    color = np.concatenate((color, unlabeled_color))
                    normals = np.concatenate((normals, unlabeled_normals))
                    labels = np.concatenate(
                        (labels, np.full_like(unlabeled_labels, self.ignore_label))
                    )

            if random() < self.color_drop:
                color[:] = 255

        # normalize color information
        pseudo_image = color.astype(np.uint8)[np.newaxis, :, :]
        color = np.squeeze(self.normalize_color(image=pseudo_image)["image"])

        # prepare labels and map from 0 to 20(40)
        labels = labels.astype(np.int32)
        if labels.size > 0:
            labels[:, 0] = self._remap_from_zero(labels[:, 0])
            if not self.add_instance:
                # taking only first column, which is segmentation label, not instance
                labels = labels[:, 0].flatten()[..., None]

        labels = np.hstack((labels, segments[..., None].astype(np.int32)))

        features = color
        if self.add_normals:
            features = np.hstack((features, normals))
        if self.add_raw_coordinates:
            if len(features.shape) == 1:
                features = np.hstack((features[None, ...], coordinates))
            else:
                features = np.hstack((features, coordinates))

        #if self.task != "semantic_segmentation":
        if self.data[idx]['raw_filepath'].split("/")[-2] in ['scene0636_00', 'scene0154_00']:
            return self.__getitem__(0)

        if self.dataset_name == "s3dis":
            return coordinates, features, labels, self.data[idx]['area'] + "_" + self.data[idx]['scene'],\
                   raw_color, raw_normals, raw_coordinates, idx
        if self.dataset_name == "stpls3d":
            if labels.shape[1] != 1: # only segments --> test set!
                if np.unique(labels[:, -2]).shape[0] < 2:
                    print("NO INSTANCES")
                    return self.__getitem__(0)
            return coordinates, features, labels, self.data[idx]['scene'], \
                   raw_color, raw_normals, raw_coordinates, idx
        else:
            return coordinates, features, labels, self.data[idx]['raw_filepath'].split("/")[-2], \
                   raw_color, raw_normals, raw_coordinates, idx

    @property
    def data(self):
        """ database file containing information about preproscessed dataset """
        return self._data

    @property
    def label_info(self):
        """ database file containing information labels used by dataset """
        return self._labels

    @staticmethod
    def _load_yaml(filepath):
        with open(filepath) as f:
            # file = yaml.load(f, Loader=Loader)
            file = yaml.load(f)
        return file

    def _select_correct_labels(self, labels, num_labels):
        number_of_validation_labels = 0
        number_of_all_labels = 0
        for k, v, in labels.items():
            number_of_all_labels += 1
            if v["validation"]:
                number_of_validation_labels += 1

        if num_labels == number_of_all_labels:
            return labels
        elif num_labels == number_of_validation_labels:
            valid_labels = dict()
            for k, v, in labels.items():
                if v["validation"]:
                    valid_labels.update({k: v})
            return valid_labels
        else:
            msg = f"""not available number labels, select from:
            {number_of_validation_labels}, {number_of_all_labels}"""
            raise ValueError(msg)

    def _remap_from_zero(self, labels):
        labels[~np.isin(labels, list(self.label_info.keys()))] = self.ignore_label
        # remap to the range from 0
        for i, k in enumerate(self.label_info.keys()):
            labels[labels == k] = i
        return labels

    def _remap_model_output(self, output):
        output = np.array(output)
        output[output==202] = 200 #setting the unknown class to common 200 label
        output_remapped = output.copy()
        for i, k in enumerate(self.label_info.keys()):
            output_remapped[output == i] = k
        return output_remapped

    def augment_individual_instance(
        self, coordinates, color, normals, labels, oversampling=1.0
    ):
        max_instance = int(len(np.unique(labels[:, 1])))
        # randomly selecting half of non-zero instances
        for instance in range(0, int(max_instance * oversampling)):
            if self.place_around_existing:
                center = choice(
                    coordinates[labels[:, 1] == choice(np.unique(labels[:, 1]))]
                )
            else:
                center = np.array([uniform(-5, 5), uniform(-5, 5), uniform(-0.5, 2)])
            instance = choice(choice(self.instance_data))
            instance = np.load(instance["instance_filepath"])
            # centering two objects
            instance[:, :3] = instance[:, :3] - instance[:, :3].mean(axis=0) + center
            max_instance = max_instance + 1
            instance[:, -1] = max_instance
            aug = V.Compose(
                [
                    V.Scale3d(),
                    V.RotateAroundAxis3d(rotation_limit=np.pi / 24, axis=(1, 0, 0)),
                    V.RotateAroundAxis3d(rotation_limit=np.pi / 24, axis=(0, 1, 0)),
                    V.RotateAroundAxis3d(rotation_limit=np.pi, axis=(0, 0, 1))
                ]
            )(
                points=instance[:, :3],
                features=instance[:, 3:6],
                normals=instance[:, 6:9],
                labels=instance[:, 9:],
            )
            coordinates = np.concatenate((coordinates, aug["points"]))
            color = np.concatenate((color, aug["features"]))
            normals = np.concatenate((normals, aug["normals"]))
            labels = np.concatenate((labels, aug["labels"]))

        return coordinates, color, normals, labels


def elastic_distortion(pointcloud, granularity, magnitude):
    """Apply elastic distortion on sparse coordinate space.

    pointcloud: numpy array of (number of points, at least 3 spatial dims)
    granularity: size of the noise grid (in same scale[m/cm] as the voxel grid)
    magnitude: noise multiplier
  """
    blurx = np.ones((3, 1, 1, 1)).astype("float32") / 3
    blury = np.ones((1, 3, 1, 1)).astype("float32") / 3
    blurz = np.ones((1, 1, 3, 1)).astype("float32") / 3
    coords = pointcloud[:, :3]
    coords_min = coords.min(0)

    # Create Gaussian noise tensor of the size given by granularity.
    noise_dim = ((coords - coords_min).max(0) // granularity).astype(int) + 3
    noise = np.random.randn(*noise_dim, 3).astype(np.float32)

    # Smoothing.
    for _ in range(2):
        noise = scipy.ndimage.filters.convolve(noise, blurx, mode="constant", cval=0)
        noise = scipy.ndimage.filters.convolve(noise, blury, mode="constant", cval=0)
        noise = scipy.ndimage.filters.convolve(noise, blurz, mode="constant", cval=0)

    # Trilinear interpolate noise filters for each spatial dimensions.
    ax = [
        np.linspace(d_min, d_max, d)
        for d_min, d_max, d in zip(
            coords_min - granularity,
            coords_min + granularity * (noise_dim - 2),
            noise_dim,
        )
    ]
    interp = scipy.interpolate.RegularGridInterpolator(
        ax, noise, bounds_error=0, fill_value=0
    )
    pointcloud[:, :3] = coords + interp(coords) * magnitude
    return pointcloud


def crop(points, x_min, y_min, z_min, x_max, y_max, z_max):
    if x_max <= x_min or y_max <= y_min or z_max <= z_min:
        raise ValueError(
            "We should have x_min < x_max and y_min < y_max and z_min < z_max. But we got"
            " (x_min = {x_min}, y_min = {y_min}, z_min = {z_min},"
            " x_max = {x_max}, y_max = {y_max}, z_max = {z_max})".format(
                x_min=x_min,
                x_max=x_max,
                y_min=y_min,
                y_max=y_max,
                z_min=z_min,
                z_max=z_max,
            )
        )
    inds = np.all(
        [
            (points[:, 0] >= x_min),
            (points[:, 0] < x_max),
            (points[:, 1] >= y_min),
            (points[:, 1] < y_max),
            (points[:, 2] >= z_min),
            (points[:, 2] < z_max),
        ],
        axis=0,
    )
    return inds


def flip_in_center(coordinates):
    # moving coordinates to center
    coordinates -= coordinates.mean(0)
    aug = V.Compose(
        [
            V.Flip3d(axis=(0, 1, 0), always_apply=True),
            V.Flip3d(axis=(1, 0, 0), always_apply=True),
        ]
    )

    first_crop = coordinates[:, 0] > 0
    first_crop &= coordinates[:, 1] > 0
    # x -y
    second_crop = coordinates[:, 0] > 0
    second_crop &= coordinates[:, 1] < 0
    # -x y
    third_crop = coordinates[:, 0] < 0
    third_crop &= coordinates[:, 1] > 0
    # -x -y
    fourth_crop = coordinates[:, 0] < 0
    fourth_crop &= coordinates[:, 1] < 0

    if first_crop.size > 1:
        coordinates[first_crop] = aug(points=coordinates[first_crop])["points"]
    if second_crop.size > 1:
        minimum = coordinates[second_crop].min(0)
        minimum[2] = 0
        minimum[0] = 0
        coordinates[second_crop] = aug(points=coordinates[second_crop])["points"]
        coordinates[second_crop] += minimum
    if third_crop.size > 1:
        minimum = coordinates[third_crop].min(0)
        minimum[2] = 0
        minimum[1] = 0
        coordinates[third_crop] = aug(points=coordinates[third_crop])["points"]
        coordinates[third_crop] += minimum
    if fourth_crop.size > 1:
        minimum = coordinates[fourth_crop].min(0)
        minimum[2] = 0
        coordinates[fourth_crop] = aug(points=coordinates[fourth_crop])["points"]
        coordinates[fourth_crop] += minimum

    return coordinates


def random_around_points(
    coordinates, color, normals, labels, rate=0.2, noise_rate=0, ignore_label=255
):
    coord_indexes = sample(
        list(range(len(coordinates))), k=int(len(coordinates) * rate)
    )
    noisy_coordinates = deepcopy(coordinates[coord_indexes])
    noisy_coordinates += np.random.uniform(
        -0.2 - noise_rate, 0.2 + noise_rate, size=noisy_coordinates.shape
    )

    if noise_rate > 0:
        noisy_color = np.random.randint(0, 255, size=noisy_coordinates.shape)
        noisy_normals = np.random.rand(*noisy_coordinates.shape) * 2 - 1
        noisy_labels = np.full(labels[coord_indexes].shape, ignore_label)

        coordinates = np.vstack((coordinates, noisy_coordinates))
        color = np.vstack((color, noisy_color))
        normals = np.vstack((normals, noisy_normals))
        labels = np.vstack((labels, noisy_labels))
    else:
        noisy_color = deepcopy(color[coord_indexes])
        noisy_normals = deepcopy(normals[coord_indexes])
        noisy_labels = deepcopy(labels[coord_indexes])

        coordinates = np.vstack((coordinates, noisy_coordinates))
        color = np.vstack((color, noisy_color))
        normals = np.vstack((normals, noisy_normals))
        labels = np.vstack((labels, noisy_labels))

    return coordinates, color, normals, labels


def random_points(
    coordinates, color, normals, labels, noise_rate=0.6, ignore_label=255
):
    max_boundary = coordinates.max(0) + 0.1
    min_boundary = coordinates.min(0) - 0.1

    noisy_coordinates = int((max(max_boundary) - min(min_boundary)) / noise_rate)

    noisy_coordinates = np.array(
        list(
            product(
                np.linspace(min_boundary[0], max_boundary[0], noisy_coordinates),
                np.linspace(min_boundary[1], max_boundary[1], noisy_coordinates),
                np.linspace(min_boundary[2], max_boundary[2], noisy_coordinates),
            )
        )
    )
    noisy_coordinates += np.random.uniform(
        -noise_rate, noise_rate, size=noisy_coordinates.shape
    )

    noisy_color = np.random.randint(0, 255, size=noisy_coordinates.shape)
    noisy_normals = np.random.rand(*noisy_coordinates.shape) * 2 - 1
    noisy_labels = np.full((noisy_coordinates.shape[0], labels.shape[1]), ignore_label)

    coordinates = np.vstack((coordinates, noisy_coordinates))
    color = np.vstack((color, noisy_color))
    normals = np.vstack((normals, noisy_normals))
    labels = np.vstack((labels, noisy_labels))
    return coordinates, color, normals, labels

datasets/semseg.py
aminebdj/3D-OWIS
import MinkowskiEngine as ME
import numpy as np
import torch
from random import random

##########################
from datasets.scannet200.scannet200_constants import MAP_STRING_TO_ID
from datasets.scannet200.scannet200_splits import HEAD_CATS_SCANNET_200, COMMON_CATS_SCANNET_200, TAIL_CATS_SCANNET_200
############################


class VoxelizeCollate:
    def __init__(
            self,
            ignore_label=255,
            voxel_size=1,
            mode="test",
            small_crops=False,
            very_small_crops=False,
            batch_instance=False,
            probing=False,
            task="instance_segmentation",
            ignore_class_threshold=100,
            filter_out_classes=[],
            label_offset=0,
            num_queries=None
    ):
        assert task in ["instance_segmentation", "semantic_segmentation"], "task not known"
        self.task = task
        self.filter_out_classes = filter_out_classes
        self.label_offset = label_offset
        self.voxel_size = voxel_size
        self.ignore_label = ignore_label
        self.mode = mode
        self.batch_instance = batch_instance
        self.small_crops = small_crops
        self.very_small_crops = very_small_crops
        self.probing = probing
        self.ignore_class_threshold = ignore_class_threshold

        self.num_queries = num_queries

    def __call__(self, batch):
        if ("train" in self.mode) and (self.small_crops or self.very_small_crops):
            batch = make_crops(batch)
        if ("train" in self.mode) and self.very_small_crops:
            batch = make_crops(batch)
        return voxelize(batch, self.ignore_label, self.voxel_size, self.probing, self.mode,
                        task=self.task, ignore_class_threshold=self.ignore_class_threshold,
                        filter_out_classes=self.filter_out_classes, label_offset=self.label_offset,
                        num_queries=self.num_queries)


class VoxelizeCollateMerge:
    def __init__(
            self,
            ignore_label=255,
            voxel_size=1,
            mode="test",
            scenes=2,
            small_crops=False,
            very_small_crops=False,
            batch_instance=False,
            make_one_pc_noise=False,
            place_nearby=False,
            place_far=False,
            proba=1,
            probing=False,
            task="instance_segmentation"
    ):
        assert task in ["instance_segmentation", "semantic_segmentation"], "task not known"
        self.task = task
        self.mode = mode
        self.scenes = scenes
        self.small_crops = small_crops
        self.very_small_crops = very_small_crops
        self.ignore_label = ignore_label
        self.voxel_size = voxel_size
        self.batch_instance = batch_instance
        self.make_one_pc_noise = make_one_pc_noise
        self.place_nearby = place_nearby
        self.place_far = place_far
        self.proba = proba
        self.probing = probing

    def __call__(self, batch):
        if (
                ("train" in self.mode)
                and (not self.make_one_pc_noise)
                and (self.proba > random())
        ):
            if self.small_crops or self.very_small_crops:
                batch = make_crops(batch)
            if self.very_small_crops:
                batch = make_crops(batch)
            if self.batch_instance:
                batch = batch_instances(batch)
            new_batch = []
            for i in range(0, len(batch), self.scenes):
                batch_coordinates = []
                batch_features = []
                batch_labels = []

                batch_filenames = ""
                batch_raw_color = []
                batch_raw_normals = []

                offset_instance_id = 0
                offset_segment_id = 0

                for j in range(min(len(batch[i:]), self.scenes)):
                    batch_coordinates.append(batch[i + j][0])
                    batch_features.append(batch[i + j][1])

                    if j == 0:
                        batch_filenames = batch[i + j][3]
                    else:
                        batch_filenames = batch_filenames + f"+{batch[i + j][3]}"

                    batch_raw_color.append(batch[i + j][4])
                    batch_raw_normals.append(batch[i + j][5])

                    # make instance ids and segment ids unique
                    # take care that -1 instances stay at -1
                    batch_labels.append(batch[i + j][2] + [0, offset_instance_id, offset_segment_id])
                    batch_labels[-1][batch[i + j][2][:, 1] == -1, 1] = -1

                    max_instance_id, max_segment_id = batch[i + j][2].max(axis=0)[1:]
                    offset_segment_id = offset_segment_id + max_segment_id + 1
                    offset_instance_id = offset_instance_id + max_instance_id + 1

                if (len(batch_coordinates) == 2) and self.place_nearby:
                    border = batch_coordinates[0][:, 0].max()
                    border -= batch_coordinates[1][:, 0].min()
                    batch_coordinates[1][:, 0] += border
                elif (len(batch_coordinates) == 2) and self.place_far:
                    batch_coordinates[1] += (
                            np.random.uniform((-10, -10, -10), (10, 10, 10)) * 200
                    )
                new_batch.append(
                    (
                        np.vstack(batch_coordinates),
                        np.vstack(batch_features),
                        np.concatenate(batch_labels),
                        batch_filenames,
                        np.vstack(batch_raw_color),
                        np.vstack(batch_raw_normals)
                    )
                )
            # TODO WHAT ABOUT POINT2SEGMENT AND SO ON ...
            batch = new_batch
        elif ("train" in self.mode) and self.make_one_pc_noise:
            new_batch = []
            for i in range(0, len(batch), 2):
                if (i + 1) < len(batch):
                    new_batch.append(
                        [
                            np.vstack((batch[i][0], batch[i + 1][0])),
                            np.vstack((batch[i][1], batch[i + 1][1])),
                            np.concatenate(
                                (
                                    batch[i][2],
                                    np.full_like(batch[i + 1][2], self.ignore_label),
                                )
                            ),
                        ]
                    )
                    new_batch.append(
                        [
                            np.vstack((batch[i][0], batch[i + 1][0])),
                            np.vstack((batch[i][1], batch[i + 1][1])),
                            np.concatenate(
                                (
                                    np.full_like(batch[i][2], self.ignore_label),
                                    batch[i + 1][2],
                                )
                            ),
                        ]
                    )
                else:
                    new_batch.append([batch[i][0], batch[i][1], batch[i][2]])
            batch = new_batch
        # return voxelize(batch, self.ignore_label, self.voxel_size, self.probing, self.mode)
        return voxelize(batch, self.ignore_label, self.voxel_size, self.probing, self.mode, task=self.task)


def batch_instances(batch):
    new_batch = []
    for sample in batch:
        for instance_id in np.unique(sample[2][:, 1]):
            new_batch.append(
                (
                    sample[0][sample[2][:, 1] == instance_id],
                    sample[1][sample[2][:, 1] == instance_id],
                    sample[2][sample[2][:, 1] == instance_id][:, 0],
                ),
            )
    return new_batch


def voxelize(batch, ignore_label, voxel_size, probing, mode, task,
             ignore_class_threshold, filter_out_classes, label_offset, num_queries):
    (coordinates, features, labels, original_labels, inverse_maps, original_colors, original_normals,
     original_coordinates, idx) = (
        [],
        [],
        [],
        [],
        [],
        [],
        [],
        [],
        []
    )
    voxelization_dict = {
        "ignore_label": ignore_label,
        # "quantization_size": self.voxel_size,
        "return_index": True,
        "return_inverse": True,
    }

    full_res_coords = []

    for sample in batch:
        idx.append(sample[7])
        original_coordinates.append(sample[6])
        original_labels.append(sample[2])
        full_res_coords.append(sample[0])
        original_colors.append(sample[4])
        original_normals.append(sample[5])

        coords = np.floor(sample[0] / voxel_size)
        voxelization_dict.update({"coordinates": torch.from_numpy(coords).to("cpu").contiguous(), "features": sample[1]})

        # maybe this change (_, _, ...) is not necessary and we can directly get out
        # the sample coordinates?
        _, _, unique_map, inverse_map = ME.utils.sparse_quantize(**voxelization_dict)
        inverse_maps.append(inverse_map)

        sample_coordinates = coords[unique_map]
        coordinates.append(torch.from_numpy(sample_coordinates).int())
        sample_features = sample[1][unique_map]
        features.append(torch.from_numpy(sample_features).float())
        if len(sample[2]) > 0:
            sample_labels = sample[2][unique_map]
            labels.append(torch.from_numpy(sample_labels).long())

    # Concatenate all lists
    input_dict = {"coords": coordinates, "feats": features}
    if len(labels) > 0:
        input_dict["labels"] = labels
        coordinates, features, labels = ME.utils.sparse_collate(**input_dict)
    else:
        coordinates, features = ME.utils.sparse_collate(**input_dict)
        labels = torch.Tensor([])

    if probing:
        return (
            NoGpu(coordinates, features, original_labels, inverse_maps, ),
            labels,
        )

    if mode == "test":
        for i in range(len(input_dict["labels"])):
            _, ret_index, ret_inv = np.unique(input_dict["labels"][i][:, 0], return_index=True, return_inverse=True)
            input_dict["labels"][i][:, 0] = torch.from_numpy(ret_inv)
            # input_dict["segment2label"].append(input_dict["labels"][i][ret_index][:, :-1])
    else:
        input_dict["segment2label"] = []

        if "labels" in input_dict:
            for i in range(len(input_dict["labels"])):
                # TODO BIGGER CHANGE CHECK!!!
                _, ret_index, ret_inv = np.unique(input_dict["labels"][i][:, -1], return_index=True, return_inverse=True)
                input_dict["labels"][i][:, -1] = torch.from_numpy(ret_inv)
                input_dict["segment2label"].append(input_dict["labels"][i][ret_index][:, :-1])

    if "labels" in input_dict:
        list_labels = input_dict["labels"]

        target = []
        target_full = []

        if len(list_labels[0].shape) == 1:
            for batch_id in range(len(list_labels)):
                label_ids = list_labels[batch_id].unique()
                if 255 in label_ids:
                    label_ids = label_ids[:-1]

                target.append({
                    'labels': label_ids,
                    'masks': list_labels[batch_id] == label_ids.unsqueeze(1)
                })
        else:
            if mode == "test":
                for i in range(len(input_dict["labels"])):
                    target.append({
                        "point2segment": input_dict["labels"][i][:, 0]
                    })
                    target_full.append({
                        "point2segment": torch.from_numpy(original_labels[i][:, 0]).long()
                    })
            else:
                target = get_instance_masks(list_labels,
                                            list_segments=input_dict["segment2label"],
                                            task=task,
                                            ignore_class_threshold=ignore_class_threshold,
                                            filter_out_classes=filter_out_classes,
                                            label_offset=label_offset)
                for i in range(len(target)):
                    target[i]["point2segment"] = input_dict["labels"][i][:, 2]
                if "train" not in mode:
                    target_full = get_instance_masks([torch.from_numpy(l) for l in original_labels],
                                                     task=task,
                                                     ignore_class_threshold=ignore_class_threshold,
                                                     filter_out_classes=filter_out_classes,
                                                     label_offset=label_offset)
                    for i in range(len(target_full)):
                        target_full[i]["point2segment"] = torch.from_numpy(original_labels[i][:, 2]).long()
    else:
        target = []
        target_full = []
        coordinates = []
        features = []

    if "train" not in mode:
        return (
            NoGpu(coordinates, features, original_labels, inverse_maps, full_res_coords,
                  target_full, original_colors, original_normals, original_coordinates, idx), target,
            [sample[3] for sample in batch]
        )
    else:
        return (
            NoGpu(coordinates, features, original_labels, inverse_maps, full_res_coords), target,
            [sample[3] for sample in batch]
        )


def get_instance_masks(list_labels, task, list_segments=None, ignore_class_threshold=100, 
                       filter_out_classes=[], label_offset=0): #<==============================================================
    target = []

    for batch_id in range(len(list_labels)):
        # #####################################################
        # COMMON_TAIL_CATS_SCANNET_200_IDS = list(map(lambda k:MAP_STRING_TO_ID[k], COMMON_CATS_SCANNET_200+TAIL_CATS_SCANNET_200))
    
        # for k in COMMON_TAIL_CATS_SCANNET_200_IDS:
        #     try:
        #         list_labels[batch_id][:, 0][list_labels[batch_id][:, 0]==k]=255
        #         list_segments[batch_id][:, 0][list_segments[batch_id][:, 0]==k]=255
        #     except:
        #         print('exception occured')
        # #####################################################
        # # check if any TAIL COMMON classes are there l = sum(list(map(lambda k: torch.mean((list_labels[batch_id][:, 0]==k).float()).item(),COMMON_TAIL_CATS_SCANNET_200_IDS)))
        label_ids = []
        masks = []
        segment_masks = []
        instance_ids = list_labels[batch_id][:, 1].unique()

        for instance_id in instance_ids:
            if instance_id == -1:
                continue

            # TODO is it possible that a ignore class (255) is an instance???
            # instance == -1 ???
            tmp = list_labels[batch_id][list_labels[batch_id][:, 1] == instance_id]
            label_id = tmp[0, 0]

            if label_id in filter_out_classes:  # floor, wall, undefined==255 is not included
                continue

            if 255 in filter_out_classes and label_id.item() == 255 and tmp.shape[0] < ignore_class_threshold:
                continue

            label_ids.append(label_id)
            masks.append(list_labels[batch_id][:, 1] == instance_id)

            if list_segments:
                segment_mask = torch.zeros(list_segments[batch_id].shape[0]).bool()
                segment_mask[list_labels[batch_id][list_labels[batch_id][:, 1] == instance_id][:, 2].unique()] = True
                segment_masks.append(segment_mask)

        if len(label_ids) == 0:
            return list()

        label_ids = torch.stack(label_ids)
        masks = torch.stack(masks)
        if list_segments:
            segment_masks = torch.stack(segment_masks)

        if task == "semantic_segmentation":
            new_label_ids = []
            new_masks = []
            new_segment_masks = []
            for label_id in label_ids.unique():
                masking = (label_ids == label_id)

                new_label_ids.append(label_id)
                new_masks.append(masks[masking, :].sum(dim=0).bool())

                if list_segments:
                    new_segment_masks.append(segment_masks[masking, :].sum(dim=0).bool())

            label_ids = torch.stack(new_label_ids)
            masks = torch.stack(new_masks)

            if list_segments:
                segment_masks = torch.stack(new_segment_masks)

                target.append({
                    'labels': label_ids,
                    'masks': masks,
                    'segment_mask': segment_masks
                })
            else:
                target.append({
                    'labels': label_ids,
                    'masks': masks
                })
        else:
            l = torch.clamp(label_ids - label_offset, min=0)

            if list_segments:
                target.append({
                    'labels': l,
                    'masks': masks,
                    'segment_mask': segment_masks
                })
            else:
                target.append({
                    'labels': l,
                    'masks': masks
                })
    return target


def make_crops(batch):
    new_batch = []
    # detupling
    for scene in batch:
        new_batch.append([scene[0], scene[1], scene[2]])
    batch = new_batch
    new_batch = []
    for scene in batch:
        # move to center for better quadrant split
        scene[0][:, :3] -= scene[0][:, :3].mean(0)

        # BUGFIX - there always would be a point in every quadrant
        scene[0] = np.vstack(
            (
                scene[0],
                np.array(
                    [
                        [0.1, 0.1, 0.1],
                        [0.1, -0.1, 0.1],
                        [-0.1, 0.1, 0.1],
                        [-0.1, -0.1, 0.1],
                    ]
                ),
            )
        )
        scene[1] = np.vstack((scene[1], np.zeros((4, scene[1].shape[1]))))
        scene[2] = np.concatenate((scene[2], np.full_like((scene[2]), 255)[:4]))

        crop = scene[0][:, 0] > 0
        crop &= scene[0][:, 1] > 0
        if crop.size > 1:
            new_batch.append([scene[0][crop], scene[1][crop], scene[2][crop]])

        crop = scene[0][:, 0] > 0
        crop &= scene[0][:, 1] < 0
        if crop.size > 1:
            new_batch.append([scene[0][crop], scene[1][crop], scene[2][crop]])

        crop = scene[0][:, 0] < 0
        crop &= scene[0][:, 1] > 0
        if crop.size > 1:
            new_batch.append([scene[0][crop], scene[1][crop], scene[2][crop]])

        crop = scene[0][:, 0] < 0
        crop &= scene[0][:, 1] < 0
        if crop.size > 1:
            new_batch.append([scene[0][crop], scene[1][crop], scene[2][crop]])

    # moving all of them to center
    for i in range(len(new_batch)):
        new_batch[i][0][:, :3] -= new_batch[i][0][:, :3].mean(0)
    return new_batch


class NoGpu:
    def __init__(
            self, coordinates, features, original_labels=None, inverse_maps=None, full_res_coords=None,
            target_full=None, original_colors=None, original_normals=None, original_coordinates=None,
            idx=None
    ):
        """ helper class to prevent gpu loading on lightning """
        self.coordinates = coordinates
        self.features = features
        self.original_labels = original_labels
        self.inverse_maps = inverse_maps
        self.full_res_coords = full_res_coords
        self.target_full = target_full
        self.original_colors = original_colors
        self.original_normals = original_normals
        self.original_coordinates = original_coordinates
        self.idx = idx


class NoGpuMask:
    def __init__(
            self, coordinates, features, original_labels=None, inverse_maps=None, masks=None, labels=None
    ):
        """ helper class to prevent gpu loading on lightning """
        self.coordinates = coordinates
        self.features = features
        self.original_labels = original_labels
        self.inverse_maps = inverse_maps

        self.masks = masks
        self.labels = labels
datasets/utils.py
aminebdj/3D-OWIS
import logging
import os
from hashlib import md5
from uuid import uuid4
import hydra 
from dotenv import load_dotenv
from trainer.trainer import InstanceSegmentation, RegularCheckpointing
from omegaconf import DictConfig, OmegaConf
from pytorch_lightning.callbacks import ModelCheckpoint
from utils.utils import (
    flatten_dict,
    load_baseline_model, 
    load_checkpoint_with_missing_or_exsessive_keys,
    load_backbone_checkpoint_with_missing_or_exsessive_keys
)
from pytorch_lightning import Trainer, seed_everything
from omegaconf import OmegaConf,open_dict
from tqdm import tqdm
import json 
import torch



 
####################################################################################################################################
import numpy as np
from datasets.scannet200.scannet200_constants import CLASS_LABELS_200
from datasets.scannet200.scannet200_splits import HEAD_CATS_SCANNET_200, COMMON_CATS_SCANNET_200, TAIL_CATS_SCANNET_200
####################################################################################################################################
############################
CLASS_LABELS_200 = list(CLASS_LABELS_200)
CLASS_LABELS_200.remove('floor')
CLASS_LABELS_200.remove('wall')
# HEAD_CATS_SCANNET_200.remove('floor')
# HEAD_CATS_SCANNET_200.remove('wall')
MAP_STRING_TO_ID = {CLASS_LABELS_200[i] : i for i in range(len(CLASS_LABELS_200))}
MAP_STRING_TO_ID['background'] = 253
############################

def get_parameters(cfg: DictConfig):
    logger = logging.getLogger(__name__)
    load_dotenv(".env")

    # parsing input parameters
    seed_everything(cfg.general.seed)

    # getting basic configuration 
    if cfg.general.get("gpus", None) is None:
        cfg.general.gpus = os.environ.get("CUDA_VISIBLE_DEVICES", None)
    loggers = []

    # cfg.general.experiment_id = "0" # str(Repo("./").commit())[:8]
    # params = flatten_dict(OmegaConf.to_container(cfg, resolve=True))

    # create unique id for experiments that are run locally
    # unique_id = "_" + str(uuid4())[:4]
    # cfg.general.version = md5(str(params).encode("utf-8")).hexdigest()[:8] + unique_id

    if not os.path.exists(cfg.general.save_dir):
        os.makedirs(cfg.general.save_dir)
    else:
        print("EXPERIMENT ALREADY EXIST")
        cfg['trainer']['resume_from_checkpoint'] = f"{cfg.general.save_dir}/last-epoch.ckpt"

    for log in cfg.logging:
        print(log)
        loggers.append(hydra.utils.instantiate(log))
        loggers[-1].log_hyperparams(
            flatten_dict(OmegaConf.to_container(cfg, resolve=True))
        )

    model = InstanceSegmentation(cfg)
    if cfg.general.backbone_checkpoint is not None:
        cfg, model = load_backbone_checkpoint_with_missing_or_exsessive_keys(cfg, model)
    if cfg.general.checkpoint is not None:
        cfg, model = load_checkpoint_with_missing_or_exsessive_keys(cfg, model)

    logger.info(flatten_dict(OmegaConf.to_container(cfg, resolve=True)))
    return cfg, model, loggers


@hydra.main(config_path="conf", config_name="config_base_instance_segmentation.yaml")
def train(cfg: DictConfig):
    if cfg.general.OW_task != "task1" and cfg.general.finetune:
        cfg.general.save_dir = cfg.general.save_dir+"_finetune"
        cfg.general.logg_suffix = cfg.general.logg_suffix+"_finetune"
        
    os.chdir(hydra.utils.get_original_cwd())
    cfg, model, loggers = get_parameters(cfg)
    callbacks = [] 
    for cb in cfg.callbacks:
        callbacks.append(hydra.utils.instantiate(cb))

    callbacks.append(RegularCheckpointing())
    # callbacks = [callbacks[0],callbacks[2]] #callbacks{1}: lr
    runner = Trainer(
        logger=loggers,
        gpus=cfg.general.gpus,
        callbacks=callbacks,
        weights_save_path=str(cfg.general.save_dir),
        **cfg.trainer
    )
    # resume_from = "/l/users/mohamed.boudjoghra/Research/Mask3D/alvis_exp/multiscale_10_AL_CC_unkn_dist_top20/task1/epoch=549-val_mean_ap_50=0.527.ckpt"
    # checkpoint = torch.load(resume_from)
    # model.load_state_dict(checkpoint["state_dict"])
    if not cfg.general.train_oracle and cfg.general.train_mode:
        if cfg.general.OW_task != "task1" and (not cfg.general.finetune):
            task = cfg.general.save_dir.split('/')[-1]
            prev_task = task.replace(task[-1], str(int(task[-1])-1))
            if os.path.exists(cfg.general.save_dir.replace(task, prev_task+"_finetune")):
                list_dir = os.listdir(cfg.general.save_dir.replace(task, prev_task+"_finetune"))
                for file in list_dir:
                    file_s = file.split("_")
                    if "ap" in file_s:
                        path_to_prev_task = cfg.general.save_dir.replace(task, prev_task+"_finetune")+"/"+file
                        break
                    else:
                        path_to_prev_task = cfg.general.save_dir.replace(task, prev_task)+"/"+"last-epoch.ckpt"
            else:
                path_to_prev_task = cfg.general.save_dir.replace(task, prev_task)+"/"+"last-epoch.ckpt"   
                  
            checkpoint = torch.load(path_to_prev_task)
            if not os.path.exists(cfg.general.save_dir+"/"+"last-epoch.ckpt"):
                model.load_state_dict(checkpoint["state_dict"]) 
        elif cfg.general.OW_task != "task1" and cfg.general.finetune:
            resume_from = os.path.join(cfg.general.save_dir.replace("_finetune",""),"last.ckpt")
            checkpoint = torch.load(resume_from)
            if not os.path.exists(cfg.general.save_dir+"/"+"last.ckpt"):
                model.load_state_dict(checkpoint["state_dict"]) 
    runner.fit(model)
    
@hydra.main(config_path="conf", config_name="config_base_instance_segmentation.yaml")
def test(cfg: DictConfig):
    # because hydra wants to change dir for some reasonimage.png
    os.chdir(hydra.utils.get_original_cwd())
    cfg, model, loggers = get_parameters(cfg)
    runner = Trainer(
        gpus=cfg.general.gpus,
        logger=loggers,
        weights_save_path=str(cfg.general.save_dir),
        **cfg.trainer
    )
    runner.test(model)

@hydra.main(config_path="conf", config_name="config_base_instance_segmentation.yaml")
def main(cfg: DictConfig):
    if cfg['general']['train_mode']:
        train(cfg)
    else:
        test(cfg)
        

if __name__ == "__main__":    main()
 
main_instance_segmentation.py
aminebdj/3D-OWIS
import models.resunet as resunet
import models.res16unet as res16unet
from models.res16unet import Res16UNet34C, Res16UNet34A, Res16UNet14A, Res16UNet34D, Res16UNet18D, Res16UNet18B, Custom30M
from models.mask3d import Mask3D

MODELS = []


def add_models(module):
    MODELS.extend([getattr(module, a) for a in dir(module) if "Net" in a])


add_models(resunet)
add_models(res16unet)
add_models(mask3d)


def get_models():
    """Returns a tuple of sample models."""
    return MODELS


def load_model(name):
    """Creates and returns an instance of the model given its class name."""
    # Find the model class from its name
    all_models = get_models()
    mdict = {model.__name__: model for model in all_models}
    if name not in mdict:
        print("Invalid model index. Options are:")
        # Display a list of valid model names
        for model in all_models:
            print(f"\t* {model.__name__}")
        return None
    NetClass = mdict[name]

    return NetClass

models/__init__.py
aminebdj/3D-OWIS
# Copyright (c) Facebook, Inc. and its affiliates.
# Modified by Bowen Cheng from https://github.com/facebookresearch/detr/blob/master/models/detr.py
# Modified for Mask3D
"""
MaskFormer criterion.
"""

import torch 
import torch.nn.functional as F
from torch import nn
import os
from detectron2.utils.comm import get_world_size
from detectron2.projects.point_rend.point_features import (
    get_uncertain_point_coords_with_randomness,
    point_sample,
)

from collections import deque
from models.misc import is_dist_avail_and_initialized, nested_tensor_from_tensor_list
# from torch.autograd.function import Function
# import os
# from datasets.scannet200.scannet200_splits import HEAD_PRED_IDS, COMMON_PRED_IDS, TAIL_PRED_IDS
from datasets.scannet200.owis_splits import KNOWN_CLASSES_IDS, PREV_KNOWN_CLASSES_IDS

def dice_loss(
        inputs: torch.Tensor,
        targets: torch.Tensor,
        num_masks: float,
    ):
    """
    Compute the DICE loss, similar to generalized IOU for masks
    Args:
        inputs: A float tensor of arbitrary shape.
                The predictions for each example.
        targets: A float tensor with the same shape as inputs. Stores the binary
                 classification label for each element in inputs
                (0 for the negative class and 1 for the positive class).
    """
    inputs = inputs.sigmoid()
    inputs = inputs.flatten(1)
    numerator = 2 * (inputs * targets).sum(-1)
    denominator = inputs.sum(-1) + targets.sum(-1)
    loss = 1 - (numerator + 1) / (denominator + 1)
    return loss.sum() / num_masks


dice_loss_jit = torch.jit.script(
    dice_loss
)  # type: torch.jit.ScriptModule


def sigmoid_ce_loss(
        inputs: torch.Tensor,
        targets: torch.Tensor,
        num_masks: float,
    ):
    """
    Args:
        inputs: A float tensor of arbitrary shape.
                The predictions for each example.
        targets: A float tensor with the same shape as inputs. Stores the binary
                 classification label for each element in inputs
                (0 for the negative class and 1 for the positive class).
    Returns:
        Loss tensor
    """
    loss = F.binary_cross_entropy_with_logits(inputs, targets, reduction="none")

    return loss.mean(1).sum() / num_masks


sigmoid_ce_loss_jit = torch.jit.script(
    sigmoid_ce_loss
)  # type: torch.jit.ScriptModule


def calculate_uncertainty(logits):
    """
    We estimate uncerainty as L1 distance between 0.0 and the logit prediction in 'logits' for the
        foreground class in `classes`.
    Args:
        logits (Tensor): A tensor of shape (R, 1, ...) for class-specific or
            class-agnostic, where R is the total number of predicted masks in all images and C is
            the number of foreground classes. The values are logits.
    Returns:
        scores (Tensor): A tensor of shape (R, 1, ...) that contains uncertainty scores with
            the most uncertain locations having the highest uncertainty score.
    """
    assert logits.shape[1] == 1
    gt_class_logits = logits.clone()
    return -(torch.abs(gt_class_logits))


class SetCriterion(nn.Module):
    """This class computes the loss for DETR.
    The process happens in two steps:
        1) we compute hungarian assignment between ground truth boxes and the outputs of the model
        2) we supervise each pair of matched ground-truth / prediction (supervise class and box)
    """

    def __init__(self, train_is_true,num_classes, 
                 matcher, 
                 weight_dict, 
                 eos_coef, 
                 losses,
                 num_points, 
                 oversample_ratio, 
                 importance_sample_ratio,
                 class_weights, 
                 num_querries,
                 hidden_dim, 
                 store_path,
                 clustering_start_iter, 
                 clustering_update_mu_iter, 
                 enable_baseline_clustering, 
                 enable_multiple_stage_clustering,
                 clustering_momentum, 
                 store_size, 
                 OW_task,
                 split,
                 train_oracle,
                 save_dir,
                 margin,
                 finetune,
                 baseline_finetune,
                 use_obj_loss):
        """Create the criterion.
        Parameters:
            num_classes: number of object categories, omitting the special no-object category
            matcher: module able to compute a matching between targets and proposals
            weight_dict: dict containing as key the names of the losses and as values their relative weight.
            eos_coef: relative classification weight applied to the no-object category
            losses: list of all the losses to be applied. See get_loss for list of available losses.
        """
        super().__init__()
        
        # global HEAD_PRED_IDS
        # global COMMON_PRED_IDS
        # global TAIL_PRED_IDS
        global KNOWN_CLASSES_IDS
        

        
        self.train_is_true = train_is_true
        self.num_classes = num_classes - 1
        self.class_weights = class_weights
        self.matcher = matcher
        self.weight_dict = weight_dict
        self.eos_coef = eos_coef
        self.losses = losses
        empty_weight = torch.ones(self.num_classes + 1)
        empty_weight[-1] = self.eos_coef
        if self.class_weights != -1:
            assert len(self.class_weights) == self.num_classes, "CLASS WEIGHTS DO NOT MATCH"
            empty_weight[:-1] = torch.tensor(self.class_weights)

        self.register_buffer("empty_weight", empty_weight)

        # pointwise mask loss parameters
        self.num_points = num_points
        self.oversample_ratio = oversample_ratio
        self.importance_sample_ratio = importance_sample_ratio
        self.OW_task = OW_task
        self.num_seen_cls = 200
        self.num_querries = num_querries
        self.store_size = store_size
        self.finetune = finetune
        self.baseline_finetune = baseline_finetune
        self.use_obj_loss = use_obj_loss
        
        self.PRED_IDS = KNOWN_CLASSES_IDS[split][OW_task]
        self.PREV_IDS = PREV_KNOWN_CLASSES_IDS[split][OW_task]
        if 200 not in self.PRED_IDS:
            self.PRED_IDS.append(200)
        
           
        self.store = Queue((self.num_seen_cls+1,self.store_size,hidden_dim), self.PRED_IDS,save_dir)
        # self.means =  self.store.get_means()
        self.clustering_start_iter = clustering_start_iter
        self.clustering_update_mu_iter = clustering_update_mu_iter
        self.enable_baseline_clustering = enable_baseline_clustering
        self.enable_multiple_stage_clustering = enable_multiple_stage_clustering
        self.save_dir = save_dir
        self.clustering_momentum = clustering_momentum
        self.train_oracle = train_oracle
        self.margin = margin
        if os.path.exists(self.save_dir+'/classes_means.pkl'):
            self.means = torch.load(self.save_dir+'/classes_means.pkl').detach().cpu()
        
        if self.finetune:
            self.means = torch.tensor([])
        self.kld_loss = nn.KLDivLoss(size_average=None, reduce=None, reduction='batchmean', log_target=False)
    
    #########################
    ####contrastive loss#####
    #########################
    def get_clustering_loss(self, L_queries,outputs, targets, indices, iter):
        device = targets[0]['labels'].device
        if (not self.enable_baseline_clustering) or self.finetune:
            return {"c_loss": torch.tensor(0.0).to(device)}
        c_loss = torch.tensor(0.0).to(device)
        if self.finetune:
            c_loss = self.clstr_loss_l2_cdist(L_queries, targets, indices)
        else:
            if (iter > self.clustering_start_iter//2) and (iter < self.clustering_start_iter):
                self.store.update_store(outputs, targets, indices)
            elif iter == self.clustering_start_iter:
                try:
                    self.means =  self.store.get_means().to(device)
                    c_loss = self.clstr_loss_l2_cdist(L_queries, targets, indices)
                except:
                    pass
                self.store.update_store(outputs, targets, indices)
                
            elif iter > self.clustering_start_iter:
                if iter % self.clustering_update_mu_iter == 0:
                    try:
                        self.means = self.clustering_momentum*self.means.to(device)+(1-self.clustering_momentum)*self.store.get_means().to(device)
                        torch.save(self.means, self.save_dir+'/classes_means.pkl')

                    except:
                        try:
                            self.means = self.store.get_means() #when resuming the training the and the classes_means.pkl file is deleted the sel.means will be initilized from the current features 
                        except:
                            pass
        
        
                try:
                    c_loss = self.clstr_loss_l2_cdist(L_queries, targets, indices)
                except:
                    pass
                self.store.update_store(outputs, targets, indices)
    

        return {"c_loss":  c_loss} 
            
    
    def clstr_loss_l2_cdist(self, L_queries, targets, indices):
        """
        Get the foreground input_features, generate distributions for the class,
        get probability of each feature from each distribution;
        Compute loss: if belonging to a class -> likelihood should be higher
                      else -> lower
        :param input_features:
        :param proposals:
        :return:
        """
        device = targets[0]['labels'].device

        for batch_id, (map_id, target_id) in enumerate(indices):
            tg_labels = targets[batch_id]['labels'][target_id] if batch_id == 0 else torch.cat([tg_labels,targets[batch_id]['labels'][target_id]])
            ref_qerries = L_queries[batch_id][map_id] if batch_id == 0 else torch.cat([ref_qerries,L_queries[batch_id][map_id]])
            
        
        if self.store.prev_means != None:
            distances = torch.cdist(ref_qerries, torch.cat([self.store.prev_means,self.means.to(device)]), p=2)
            LABELS_LIST = self.PREV_IDS+self.PRED_IDS
        else:
            distances = torch.cdist(ref_qerries, self.means.to(device), p=2)
            LABELS_LIST = self.PRED_IDS
            
        # distances[:,-1] = 100*distances[:,-1]
        cc_labels = []    
        for index in range(ref_qerries.shape[0]):
            for cls_index in LABELS_LIST:
                if  tg_labels[index] ==  cls_index:
                    cc_labels.append(1)
                else:
                    cc_labels.append(-1)
                            

        loss = F.hinge_embedding_loss(distances, torch.tensor(cc_labels).reshape((-1,len(LABELS_LIST))).to(device),reduction='mean', margin=self.margin)
        
        return loss    
    #########################
    ####contrastive loss#####
    #########################
    def loss_obj(self, outputs, targets, indices, num_masks, mask_type):
        if not self.use_obj_loss:
            losses = {"_loss_obj": torch.tensor(0.0)}
        else:    
            __loss_obj = []
            for batch_id, (map_id, target_id) in enumerate(indices):
                __loss_obj.append(self.kld_loss((outputs['obj'][batch_id])[map_id].log(),torch.ones_like((outputs['obj'][batch_id])[map_id])))

            _loss_obj = sum(__loss_obj)
            
            losses = {"_loss_obj": _loss_obj}
        
        return losses

    
    def loss_labels(self, outputs, targets, indices, num_masks, mask_type):
        """Classification loss (NLL)
        targets dicts must contain the key "labels" containing a tensor of dim [nb_target_boxes]
        """
        assert "pred_logits" in outputs
        src_logits = outputs["pred_logits"].float()

        idx = self._get_src_permutation_idx(indices)
        target_classes_o = torch.cat([t["labels"][J] for t, (_, J) in zip(targets, indices)])
        target_classes = torch.full(
            src_logits.shape[:2], self.num_classes, dtype=torch.int64, device=src_logits.device
        )
        target_classes[idx] = target_classes_o

        loss_ce = F.cross_entropy(src_logits.transpose(1, 2), target_classes, self.empty_weight, ignore_index=253)
        losses = {"loss_ce": loss_ce}
        return losses

    def loss_masks(self, outputs, targets, indices, num_masks, mask_type):
        """Compute the losses related to the masks: the focal loss and the dice loss.
        targets dicts must contain the key "masks" containing a tensor of dim [nb_target_boxes, h, w]
        """
        assert "pred_masks" in outputs

        loss_masks = []
        loss_dices = []

        for batch_id, (map_id, target_id) in enumerate(indices):
            map = outputs["pred_masks"][batch_id][:, map_id].T
            target_mask = targets[batch_id][mask_type][target_id]
            
            
            if self.num_points != -1:
                point_idx = torch.randperm(target_mask.shape[1], 
                                        device=target_mask.device)[:int(self.num_points*target_mask.shape[1])]
            else:
                # sample all points
                point_idx = torch.arange(target_mask.shape[1], device=target_mask.device)

            # num_masks = target_mask.shape[0]
            map = map[:, point_idx]
            device = map.device
            target_mask = target_mask[:, point_idx].float()
            #################################################################f
            #################################################################
            #################################################################
            # print(targets[batch_id]['labels'][target_id] )
            # print(mask_type)
            if self.train_is_true:
                if self.train_oracle:
                    ignore_masks = targets[batch_id]['labels'][target_id] != 253 
                else:
                    ignore_masks = (targets[batch_id]['labels'][target_id] != 253)*(targets[batch_id]['labels'][target_id] != 200)
                    
                target_mask = target_mask[ignore_masks]
                map = map[ignore_masks]
                num_masks = target_mask.shape[0]
                
            # print(map.shape)
            # print(target_mask.shape)
            # print(num_masks)
            #################################################################
            #################################################################
            #################################################################
            if num_masks!=0: #only for scenes which have the common class
                loss_masks.append(sigmoid_ce_loss_jit(map,
                                                    target_mask,
                                                    num_masks))
                loss_dices.append(dice_loss_jit(map,
                                                target_mask,
                                                num_masks))
        if loss_masks: 
        # del target_mask
            return {
                "loss_mask": torch.sum(torch.stack(loss_masks)),
                "loss_dice": torch.sum(torch.stack(loss_dices))
            }
        else: #if no scene in the batch has common or unknown
            return {
                "loss_mask": torch.tensor(0.0).to(device),
                "loss_dice": torch.tensor(0.0).to(device)
            }

        src_idx = self._get_src_permutation_idx(indices)
        tgt_idx = self._get_tgt_permutation_idx(indices)
        src_masks = outputs["pred_masks"]
        src_masks = src_masks[src_idx]
        masks = [t[mask_type] for t in targets]
        # TODO use valid to mask invalid areas due to padding in loss
        target_masks, valid = nested_tensor_from_tensor_list(masks).decompose()
        target_masks = target_masks.to(src_masks)
        target_masks = target_masks[tgt_idx]

        # No need to upsample predictions as we are using normalized coordinates :)
        # N x 1 x H x W
        src_masks = src_masks[:, None]
        target_masks = target_masks[:, None]

        with torch.no_grad():
            # sample point_coords
            point_coords = get_uncertain_point_coords_with_randomness(
                src_masks,
                lambda logits: calculate_uncertainty(logits),
                self.num_points,
                self.oversample_ratio,
                self.importance_sample_ratio,
            )
            # get gt labels
            point_labels = point_sample(
                target_masks,
                point_coords,
                align_corners=False,
            ).squeeze(1)

        point_logits = point_sample(
            src_masks,
            point_coords,
            align_corners=False,
        ).squeeze(1)

        losses = {
            "loss_mask": sigmoid_ce_loss_jit(point_logits, point_labels, num_masks, mask_type),
            "loss_dice": dice_loss_jit(point_logits, point_labels, num_masks, mask_type),
        }

        del src_masks
        del target_masks
        return losses

    def _get_src_permutation_idx(self, indices):
        # permute predictions following indices
        batch_idx = torch.cat([torch.full_like(src, i) for i, (src, _) in enumerate(indices)])
        src_idx = torch.cat([src for (src, _) in indices])
        return batch_idx, src_idx

    def _get_tgt_permutation_idx(self, indices):
        # permute targets following indices
        batch_idx = torch.cat([torch.full_like(tgt, i) for i, (_, tgt) in enumerate(indices)])
        tgt_idx = torch.cat([tgt for (_, tgt) in indices])
        return batch_idx, tgt_idx

    def get_loss(self, loss, outputs, targets, indices, num_masks, mask_type):
        loss_map = {
            'labels': self.loss_labels,
            'masks': self.loss_masks,
            'objectness': self.loss_obj
        }
        assert loss in loss_map, f"do you really want to compute {loss} loss?"
        return loss_map[loss](outputs, targets, indices, num_masks, mask_type)

    def forward(self, outputs, targets, mask_type, iteration):
        """This performs the loss computation.
        Parameters:
             outputs: dict of tensors, see the output specification of the model for the format
             targets: list of dicts, such that len(targets) == batch_size.
                      The expected keys in each dict depends on the losses applied, see each loss' doc
        """
        outputs_without_aux = {k: v for k, v in outputs.items() if k != "aux_outputs"}

        # Retrieve the matching between the outputs of the last layer and the targets
        
        indices = self.matcher(outputs_without_aux, targets, mask_type)

        # Compute the average number of target boxes accross all nodes, for normalization purposes
        num_masks = sum(len(t["labels"]) for t in targets)
        num_masks = torch.as_tensor(
            [num_masks], dtype=torch.float, device=next(iter(outputs.values())).device
        )
        if is_dist_avail_and_initialized():
            torch.distributed.all_reduce(num_masks)
        num_masks = torch.clamp(num_masks / get_world_size(), min=1).item()

        # Compute all the requested losses
        losses = {}
        for loss in self.losses:
            losses.update(self.get_loss(loss, outputs, targets, indices, num_masks, mask_type))
            
        if not self.enable_multiple_stage_clustering:
            losses.update(self.get_clustering_loss(outputs["refin_queries"], outputs, targets, indices, iteration))
        # In case of auxiliary losses, we repeat this process with the output of each intermediate layer.
        if "aux_outputs" in outputs:
            for i, aux_outputs in enumerate(outputs["aux_outputs"]):
                indices = self.matcher(aux_outputs, targets, mask_type)
                for loss in self.losses:
                    l_dict = self.get_loss(loss, aux_outputs, targets, indices, num_masks, mask_type)
                    l_dict = {k + f"_{i}": v for k, v in l_dict.items()}
                    losses.update(l_dict)
                
                if self.enable_multiple_stage_clustering: 
                    if (i+1)%3==0:   
                        l_dict = {"c_loss" + f"_{i}": self.get_clustering_loss(outputs["mul_queries"][:,i,:,:],outputs, targets, indices, iteration)["c_loss"]}
                        losses.update(l_dict)
                
        return losses
           

    def __repr__(self):
        head = "Criterion " + self.__class__.__name__
        body = [
            "matcher: {}".format(self.matcher.__repr__(_repr_indent=8)),
            "losses: {}".format(self.losses),
            "weight_dict: {}".format(self.weight_dict),
            "num_classes: {}".format(self.num_classes),
            "eos_coef: {}".format(self.eos_coef),
            "num_points: {}".format(self.num_points),
            "oversample_ratio: {}".format(self.oversample_ratio),
            "importance_sample_ratio: {}".format(self.importance_sample_ratio),
        ]
        _repr_indent = 4
        lines = [head] + [" " * _repr_indent + line for line in body]
        return "\n".join(lines)
    
    def get_indices(self, outputs, targets, mask_type):
        outputs_without_aux = {k: v for k, v in outputs.items() if k != "aux_outputs"}
        return self.matcher(outputs_without_aux, targets, mask_type)


class Queue:
    def __init__(self, store_cap, seen_cls, save_dir):
        
        self.seen_cls = seen_cls
        self.num_seen_cls = store_cap[0]
        # self.size_per_cls = (store_cap[1], store_cap[2])
        self.size_per_cls = store_cap[1]
        
        # self.store = [torch.zeros(size = self.size_per_cls) for _ in range(self.num_seen_cls)]
        self.store = [deque(maxlen=self.size_per_cls) for _ in range(self.num_seen_cls)]
        # self.store_path = store_path
        # root = store_path.replace('store.pt', '')
        # if not os.path.exists(root):
        #     os.makedirs(root)
        # torch.save(self.store, self.store_path)
        self.save_dir = save_dir
        self.task = save_dir.split('/')[-1]
        finetune_exists = "finetune" in self.task.split('_')
        if not finetune_exists:
            prev_task = self.task.replace(self.task[-1], str(int(self.task[-1])-1))
        else:
            prev_task = self.task.replace("_finetune", "")
        self.path_to_prev_means = save_dir.replace(self.task, prev_task)+"/"+prev_task+"_classes_means.pkl"
        if os.path.exists(self.path_to_prev_means):
            if not finetune_exists:
                self.prev_means = torch.load(self.path_to_prev_means)[:-1]
            else:
                self.prev_means = torch.load(self.path_to_prev_means)
        else:
            self.prev_means = None
            

            
    def update_store(self, outputs, targets , indices):
        # oracle = False
        # self.store = torch.load(self.store_path)
        # self.store = [t for t in self.store]
        self.device = outputs['pred_logits'].device
        # pred_logits = outputs['pred_logits']
        # pred_logits = torch.functional.F.softmax(
        #     pred_logits ,
        #     dim=-1)[..., :-1]
        # pred_labels = torch.argmax(pred_logits, dim=-1).detach().cpu()
        # unkn_lbs = pred_labels==200
        for batch_id, (map_id, target_id) in enumerate(indices): 
            #known
            tg_labels = targets[batch_id]['labels'][target_id].detach().cpu()
            ref_qerries = outputs['refin_queries'][batch_id][map_id].detach().cpu()
            all_ref_querries = outputs['refin_queries'][batch_id].detach().cpu()
            label_list = torch.unique(tg_labels).detach().cpu().tolist()
            if 198 in label_list:
                label_list.remove(198)
            if 253 in label_list:
                label_list.remove(253)
                
            for tg in label_list:
                # num_feats = (tg_labels == tg).sum()
                # if tg != 200:
                self.store[tg].append(ref_qerries[tg_labels==tg])
                # else:
                #     #oracle
                #     self.store[-1].append(ref_qerries[tg_labels==tg])
   
            #unknown
            # num_feats = (unkn_lbs[batch_id]).sum()
            # if num_feats != 0:
            #     self.store[200].append(all_ref_querries[unkn_lbs[batch_id]])

        # torch.save(self.store, self.store_path)
    def get_store(self):
        return self.store
    def get_means(self):
        # for i in len(self.store):
        means = [torch.mean(torch.cat(list(self.store[i])).float(), dim=0) for i in self.seen_cls]  
        curr_means = torch.stack(means).to(self.device)
        
        if self.prev_means != None:
            torch.save(torch.cat([self.prev_means, curr_means]), self.save_dir+"/"+self.task+"_classes_means.pkl")
        else:
            torch.save(curr_means, self.save_dir+"/"+self.task+"_classes_means.pkl")
            
        return curr_means

models/criterion.py
aminebdj/3D-OWIS
import torch
import hydra
import torch.nn as nn
import MinkowskiEngine.MinkowskiOps as me
from MinkowskiEngine.MinkowskiPooling import MinkowskiAvgPooling
import numpy as np
from torch.nn import functional as F
from models.modules.common import conv
from models.position_embedding import PositionEmbeddingCoordsSine
from third_party.pointnet2.pointnet2_utils import furthest_point_sample
from models.modules.helpers_3detr import GenericMLP
from torch_scatter import scatter_mean, scatter_max, scatter_min
from torch.cuda.amp import autocast


class Mask3D(nn.Module):
    def __init__(self, config, hidden_dim, num_queries, num_heads, dim_feedforward,
                 sample_sizes, shared_decoder, num_classes,
                 num_decoders, dropout, pre_norm,
                 positional_encoding_type, non_parametric_queries, train_on_segments, normalize_pos_enc,
                 use_level_embed, scatter_type, hlevels,
                 use_np_features,
                 voxel_size,
                 max_sample_size,
                 random_queries,
                 gauss_scale,
                 random_query_both,
                 random_normal
                 ):
        super().__init__()

        self.random_normal = random_normal
        self.random_query_both = random_query_both
        self.random_queries = random_queries
        self.max_sample_size = max_sample_size
        self.gauss_scale = gauss_scale
        self.voxel_size = voxel_size
        self.scatter_type = scatter_type
        self.hlevels = hlevels
        self.use_level_embed = use_level_embed
        self.train_on_segments = train_on_segments
        self.normalize_pos_enc = normalize_pos_enc
        self.num_decoders = num_decoders
        self.num_classes = num_classes
        self.dropout = dropout
        self.pre_norm = pre_norm
        self.shared_decoder = shared_decoder
        self.sample_sizes = sample_sizes
        self.non_parametric_queries = non_parametric_queries
        self.use_np_features = use_np_features
        self.mask_dim = hidden_dim
        self.num_heads = num_heads
        self.num_queries = num_queries
        self.pos_enc_type = positional_encoding_type

        self.backbone = hydra.utils.instantiate(config.backbone)
        self.num_levels = len(self.hlevels)
        sizes = self.backbone.PLANES[-5:]

        self.mask_features_head = conv(
            self.backbone.PLANES[7], self.mask_dim, kernel_size=1, stride=1, bias=True, D=3
        )

        if self.scatter_type == "mean":
            self.scatter_fn = scatter_mean
        elif self.scatter_type == "max":
            self.scatter_fn = lambda mask, p2s, dim: scatter_max(mask, p2s, dim=dim)[0]
        else:
            assert False, "Scatter function not known"

        assert (not use_np_features) or non_parametric_queries, "np features only with np queries"

        if self.non_parametric_queries:
            self.query_projection = GenericMLP(
                input_dim=self.mask_dim,
                hidden_dims=[self.mask_dim],
                output_dim=self.mask_dim,
                use_conv=True,
                output_use_activation=True,
                hidden_use_bias=True,
            )

            if self.use_np_features:
                self.np_feature_projection = nn.Sequential(
                    nn.Linear(sizes[-1], hidden_dim),
                    nn.ReLU(),
                    nn.Linear(hidden_dim, hidden_dim)
                )
        elif self.random_query_both:
            self.query_projection = GenericMLP(
                input_dim=2*self.mask_dim,
                hidden_dims=[2*self.mask_dim],
                output_dim=2*self.mask_dim,
                use_conv=True,
                output_use_activation=True,
                hidden_use_bias=True
            )
        else:
            # PARAMETRIC QUERIES
            # learnable query features
            self.query_feat = nn.Embedding(num_queries, hidden_dim)
            # learnable query p.e.
            self.query_pos = nn.Embedding(num_queries, hidden_dim)

        if self.use_level_embed:
            # learnable scale-level embedding
            self.level_embed = nn.Embedding(self.num_levels, hidden_dim)

        self.mask_embed_head = nn.Sequential(
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim)
        )

        self.class_embed_head = nn.Linear(hidden_dim, self.num_classes)

        if self.pos_enc_type == "legacy":
            self.pos_enc = PositionalEncoding3D(channels=self.mask_dim)
        elif self.pos_enc_type == "fourier":
            self.pos_enc = PositionEmbeddingCoordsSine(pos_type="fourier",
                                                       d_pos=self.mask_dim,
                                                       gauss_scale=self.gauss_scale,
                                                       normalize=self.normalize_pos_enc)
        elif self.pos_enc_type == "sine":
            self.pos_enc = PositionEmbeddingCoordsSine(pos_type="sine",
                                                       d_pos=self.mask_dim,
                                                       normalize=self.normalize_pos_enc)
        else:
            assert False, 'pos enc type not known'

        self.pooling = MinkowskiAvgPooling(kernel_size=2, stride=2, dimension=3)

        self.masked_transformer_decoder = nn.ModuleList()
        self.cross_attention = nn.ModuleList()
        self.self_attention = nn.ModuleList()
        self.ffn_attention = nn.ModuleList()
        self.lin_squeeze = nn.ModuleList()
        self.softmax = nn.Softmax(dim=2)
        num_shared = self.num_decoders if not self.shared_decoder else 1

        for _ in range(num_shared):
            tmp_cross_attention = nn.ModuleList()
            tmp_self_attention = nn.ModuleList()
            tmp_ffn_attention = nn.ModuleList()
            tmp_squeeze_attention = nn.ModuleList()
            for i, hlevel in enumerate(self.hlevels):
                tmp_cross_attention.append(
                    CrossAttentionLayer(
                        d_model=self.mask_dim,
                        nhead=self.num_heads,
                        dropout=self.dropout,
                        normalize_before=self.pre_norm
                    )
                )

                tmp_squeeze_attention.append(nn.Linear(sizes[hlevel], self.mask_dim))

                tmp_self_attention.append(
                    SelfAttentionLayer(
                        d_model=self.mask_dim,
                        nhead=self.num_heads,
                        dropout=self.dropout,
                        normalize_before=self.pre_norm
                    )
                )

                tmp_ffn_attention.append(
                    FFNLayer(
                        d_model=self.mask_dim,
                        dim_feedforward=dim_feedforward,
                        dropout=self.dropout,
                        normalize_before=self.pre_norm
                    )
                )

            self.cross_attention.append(tmp_cross_attention)
            self.self_attention.append(tmp_self_attention)
            self.ffn_attention.append(tmp_ffn_attention)
            self.lin_squeeze.append(tmp_squeeze_attention)

        self.decoder_norm = nn.LayerNorm(hidden_dim)

    def get_pos_encs(self, coords):
        pos_encodings_pcd = []

        for i in range(len(coords)):
            pos_encodings_pcd.append([[]])
            for coords_batch in coords[i].decomposed_features:
                scene_min = coords_batch.min(dim=0)[0][None, ...]
                scene_max = coords_batch.max(dim=0)[0][None, ...]

                with autocast(enabled=False):
                    tmp = self.pos_enc(coords_batch[None, ...].float(),
                                       input_range=[scene_min, scene_max])

                pos_encodings_pcd[-1][0].append(tmp.squeeze(0).permute((1, 0)))

        return pos_encodings_pcd

    def forward(self, x, point2segment=None, raw_coordinates=None, is_eval=False):
        pcd_features, aux = self.backbone(x)

        batch_size = len(x.decomposed_coordinates) 

        with torch.no_grad():
            coordinates = me.SparseTensor(features=raw_coordinates,
                                          coordinate_manager=aux[-1].coordinate_manager,
                                          coordinate_map_key=aux[-1].coordinate_map_key,
                                          device=aux[-1].device)

            coords = [coordinates]
            for _ in reversed(range(len(aux)-1)):
                coords.append(self.pooling(coords[-1]))

            coords.reverse()

        pos_encodings_pcd = self.get_pos_encs(coords)
        mask_features = self.mask_features_head(pcd_features)

        if self.train_on_segments:
            mask_segments = []
            for i, mask_feature in enumerate(mask_features.decomposed_features):
                mask_segments.append(self.scatter_fn(mask_feature, point2segment[i], dim=0))

        sampled_coords = None

        if self.non_parametric_queries:
            fps_idx = [furthest_point_sample(x.decomposed_coordinates[i][None, ...].float(),
                                             self.num_queries).squeeze(0).long()
                       for i in range(len(x.decomposed_coordinates))]

            sampled_coords = torch.stack([coordinates.decomposed_features[i][fps_idx[i].long(), :]
                                          for i in range(len(fps_idx))])

            mins = torch.stack([coordinates.decomposed_features[i].min(dim=0)[0] for i in range(len(coordinates.decomposed_features))])
            maxs = torch.stack([coordinates.decomposed_features[i].max(dim=0)[0] for i in range(len(coordinates.decomposed_features))])

            query_pos = self.pos_enc(sampled_coords.float(),
                                     input_range=[mins, maxs]
                                     )  # Batch, Dim, queries
            query_pos = self.query_projection(query_pos)
            # query_pos = torch.zeros_like(query_pos)
            if not self.use_np_features:
                queries = torch.zeros_like(query_pos).permute((0, 2, 1))
            else:
                queries = torch.stack([pcd_features.decomposed_features[i][fps_idx[i].long(), :]
                                       for i in range(len(fps_idx))])
                queries = self.np_feature_projection(queries)
            query_pos = query_pos.permute((2, 0, 1))
        elif self.random_queries:
            query_pos = torch.rand(batch_size, self.mask_dim, self.num_queries, device=x.device) - 0.5

            queries = torch.zeros_like(query_pos).permute((0, 2, 1))
            query_pos = query_pos.permute((2, 0, 1))
        elif self.random_query_both:
            if not self.random_normal:
                query_pos_feat = torch.rand(batch_size, 2*self.mask_dim, self.num_queries, device=x.device) - 0.5
            else:
                query_pos_feat = torch.randn(batch_size, 2 * self.mask_dim, self.num_queries, device=x.device)

            queries = query_pos_feat[:, :self.mask_dim, :].permute((0, 2, 1))
            query_pos = query_pos_feat[:, self.mask_dim:, :].permute((2, 0, 1))
        else:
            # PARAMETRIC QUERIES
            queries = self.query_feat.weight.unsqueeze(0).repeat(batch_size, 1, 1)
            query_pos = self.query_pos.weight.unsqueeze(1).repeat(1, batch_size, 1)

        predictions_class = []
        predictions_mask = []
        mul_queries = []
        for decoder_counter in range(self.num_decoders):
            if self.shared_decoder:
                decoder_counter = 0
            for i, hlevel in enumerate(self.hlevels):
                if self.train_on_segments:
                    output_class, outputs_mask, attn_mask = self.mask_module(queries,
                                                          mask_features,
                                                          mask_segments,
                                                          len(aux) - hlevel - 1,
                                                          ret_attn_mask=True,
                                                          point2segment=point2segment,
                                                          coords=coords)
                else:
                    output_class, outputs_mask, attn_mask = self.mask_module(queries,
                                                          mask_features,
                                                          None,
                                                          len(aux) - hlevel - 1,
                                                          ret_attn_mask=True,
                                                          point2segment=None,
                                                          coords=coords)

                decomposed_aux = aux[hlevel].decomposed_features
                decomposed_attn = attn_mask.decomposed_features

                curr_sample_size = max([pcd.shape[0] for pcd in decomposed_aux])

                if min([pcd.shape[0] for pcd in decomposed_aux]) == 1:
                    raise RuntimeError("only a single point gives nans in cross-attention")

                if not (self.max_sample_size or is_eval):
                    curr_sample_size = min(curr_sample_size, self.sample_sizes[hlevel])

                rand_idx = []
                mask_idx = []
                for k in range(len(decomposed_aux)):
                    pcd_size = decomposed_aux[k].shape[0]
                    if pcd_size <= curr_sample_size:
                        # we do not need to sample
                        # take all points and pad the rest with zeroes and mask it
                        idx = torch.zeros(curr_sample_size,
                                          dtype=torch.long,
                                          device=queries.device)

                        midx = torch.ones(curr_sample_size,
                                          dtype=torch.bool,
                                          device=queries.device)

                        idx[:pcd_size] = torch.arange(pcd_size,
                                                      device=queries.device)

                        midx[:pcd_size] = False  # attend to first points
                    else:
                        # we have more points in pcd as we like to sample
                        # take a subset (no padding or masking needed)
                        idx = torch.randperm(decomposed_aux[k].shape[0],
                                             device=queries.device)[:curr_sample_size]
                        midx = torch.zeros(curr_sample_size,
                                           dtype=torch.bool,
                                           device=queries.device)  # attend to all

                    rand_idx.append(idx)
                    mask_idx.append(midx)

                batched_aux = torch.stack([
                    decomposed_aux[k][rand_idx[k], :] for k in range(len(rand_idx))
                ])

                batched_attn = torch.stack([
                    decomposed_attn[k][rand_idx[k], :] for k in range(len(rand_idx))
                ])

                batched_pos_enc = torch.stack([
                    pos_encodings_pcd[hlevel][0][k][rand_idx[k], :] for k in range(len(rand_idx))
                ])

                batched_attn.permute((0, 2, 1))[batched_attn.sum(1) == rand_idx[0].shape[0]] = False

                m = torch.stack(mask_idx)
                batched_attn = torch.logical_or(batched_attn, m[..., None])

                src_pcd = self.lin_squeeze[decoder_counter][i](batched_aux.permute((1, 0, 2)))
                if self.use_level_embed:
                    src_pcd += self.level_embed.weight[i]

                output = self.cross_attention[decoder_counter][i](
                    queries.permute((1, 0, 2)),
                    src_pcd,
                    memory_mask=batched_attn.repeat_interleave(self.num_heads, dim=0).permute((0, 2, 1)),
                    memory_key_padding_mask=None,  # here we do not apply masking on padded region
                    pos=batched_pos_enc.permute((1, 0, 2)),
                    query_pos=query_pos
                )

                output = self.self_attention[decoder_counter][i](
                    output, tgt_mask=None,
                    tgt_key_padding_mask=None,
                    query_pos=query_pos
                )

                # FFN
                queries = self.ffn_attention[decoder_counter][i](
                    output
                ).permute((1, 0, 2))

                predictions_class.append(output_class)
                predictions_mask.append(outputs_mask)
                mul_queries.append(queries)
                
        if self.train_on_segments:
            output_class, outputs_mask = self.mask_module(queries,
                                                          mask_features,
                                                          mask_segments,
                                                          0,
                                                          ret_attn_mask=False,
                                                          point2segment=point2segment,
                                                          coords=coords)
        else:
            output_class, outputs_mask = self.mask_module(queries,
                                                          mask_features,
                                                          None,
                                                          0,
                                                          ret_attn_mask=False,
                                                          point2segment=None,
                                                          coords=coords)
        predictions_class.append(output_class) 
        predictions_mask.append(outputs_mask)

        return {
            'pred_logits': predictions_class[-1],
            'pred_masks': predictions_mask[-1],
            'aux_outputs': self._set_aux_loss(
                predictions_class, predictions_mask
            ),
            'sampled_coords': sampled_coords.detach().cpu().numpy() if sampled_coords is not None else None,
            'backbone_features': pcd_features,
            'refin_queries': queries,
            'mul_queries': torch.stack(mul_queries).permute(1,0,2,3),
            # 'obj': (torch.max(self.softmax(predictions_class[-1]),dim=-1).values)*torch.stack([torch.sum(_.float().sigmoid(),dim=0)/(torch.sum((_.float().sigmoid()>0.5).float(),dim=0)+0.0001) for _ in predictions_mask[-1]])
            'obj': (torch.max(self.softmax(predictions_class[-1]),dim=-1).values)*(torch.stack([torch.mean(torch.topk(_.float(),k=10,dim=0).values,dim=0) for _ in predictions_mask[-1]]).sigmoid())
        }

    def mask_module(self, query_feat, mask_features, mask_segments, num_pooling_steps, ret_attn_mask=True,
                                 point2segment=None, coords=None):
        query_feat = self.decoder_norm(query_feat)
        mask_embed = self.mask_embed_head(query_feat)
        outputs_class = self.class_embed_head(query_feat)

        output_masks = []

        if point2segment is not None:
            output_segments = []
            for i in range(len(mask_segments)):
                output_segments.append(mask_segments[i] @ mask_embed[i].T)
                output_masks.append(output_segments[-1][point2segment[i]])
        else:
            for i in range(mask_features.C[-1, 0] + 1):
                output_masks.append(mask_features.decomposed_features[i] @ mask_embed[i].T)

        output_masks = torch.cat(output_masks)
        outputs_mask = me.SparseTensor(features=output_masks,
                                       coordinate_manager=mask_features.coordinate_manager,
                                       coordinate_map_key=mask_features.coordinate_map_key)

        if ret_attn_mask:
            attn_mask = outputs_mask
            for _ in range(num_pooling_steps):
                attn_mask = self.pooling(attn_mask.float())

            attn_mask = me.SparseTensor(features=(attn_mask.F.detach().sigmoid() < 0.5),
                                        coordinate_manager=attn_mask.coordinate_manager,
                                        coordinate_map_key=attn_mask.coordinate_map_key)

            if point2segment is not None:
                return outputs_class, output_segments, attn_mask
            else:
                return outputs_class, outputs_mask.decomposed_features, attn_mask

        if point2segment is not None:
            return outputs_class, output_segments
        else:
            return outputs_class, outputs_mask.decomposed_features

    @torch.jit.unused
    def _set_aux_loss(self, outputs_class, outputs_seg_masks):
        # this is a workaround to make torchscript happy, as torchscript
        # doesn't support dictionary with non-homogeneous values, such
        # as a dict having both a Tensor and a list. 
        return [
            # {"pred_logits": a, "pred_masks": b, "obj":(torch.max(self.softmax(a),dim=-1).values)*torch.stack([torch.sum(_.float().sigmoid(),dim=0)/(torch.sum((_.float().sigmoid()>0.5).float(),dim=0)+0.0001) for _ in b])}
            {"pred_logits": a, "pred_masks": b, "obj":(torch.max(self.softmax(a),dim=-1).values)*(torch.stack([torch.mean(torch.topk(_.float(),k=10,dim=0).values,dim=0) for _ in b]).sigmoid())}
            for a, b in zip(outputs_class[:-1], outputs_seg_masks[:-1])
        ]


class PositionalEncoding3D(nn.Module):
    def __init__(self, channels):
        """
        :param channels: The last dimension of the tensor you want to apply pos emb to.
        """
        self.orig_ch = channels
        super(PositionalEncoding3D, self).__init__()
        channels = int(np.ceil(channels / 6) * 2)
        if channels % 2:
            channels += 1
        self.channels = channels
        inv_freq = 1.0 / (10000 ** (torch.arange(0, channels, 2).float() / channels))
        self.register_buffer("inv_freq", inv_freq)

    def forward(self, tensor, input_range=None):
        """
        :param tensor: A 5d tensor of size (batch_size, x, y, z, ch)
        :return: Positional Encoding Matrix of size (batch_size, x, y, z, ch)
        """
        pos_x, pos_y, pos_z = tensor[:, :, 0], tensor[:, :, 1], tensor[:, :, 2]
        sin_inp_x = torch.einsum("bi,j->bij", pos_x, self.inv_freq)
        sin_inp_y = torch.einsum("bi,j->bij", pos_y, self.inv_freq)
        sin_inp_z = torch.einsum("bi,j->bij", pos_z, self.inv_freq)
        emb_x = torch.cat((sin_inp_x.sin(), sin_inp_x.cos()), dim=-1)

        emb_y = torch.cat((sin_inp_y.sin(), sin_inp_y.cos()), dim=-1)
        emb_z = torch.cat((sin_inp_z.sin(), sin_inp_z.cos()), dim=-1)

        emb = torch.cat((emb_x, emb_y, emb_z), dim=-1)
        return emb[:, :, :self.orig_ch].permute((0, 2, 1))


class SelfAttentionLayer(nn.Module):

    def __init__(self, d_model, nhead, dropout=0.0,
                 activation="relu", normalize_before=False):
        super().__init__()
        self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout)

        self.norm = nn.LayerNorm(d_model)
        self.dropout = nn.Dropout(dropout)

        self.activation = _get_activation_fn(activation)
        self.normalize_before = normalize_before

        self._reset_parameters()

    def _reset_parameters(self):
        for p in self.parameters():
            if p.dim() > 1:
                nn.init.xavier_uniform_(p)

    def with_pos_embed(self, tensor, pos):
        return tensor if pos is None else tensor + pos

    def forward_post(self, tgt,
                     tgt_mask = None,
                     tgt_key_padding_mask= None,
                     query_pos = None):
        q = k = self.with_pos_embed(tgt, query_pos)
        tgt2 = self.self_attn(q, k, value=tgt, attn_mask=tgt_mask,
                              key_padding_mask=tgt_key_padding_mask)[0]
        tgt = tgt + self.dropout(tgt2)
        tgt = self.norm(tgt)

        return tgt

    def forward_pre(self, tgt,
                    tgt_mask= None,
                    tgt_key_padding_mask = None,
                    query_pos = None):
        tgt2 = self.norm(tgt)
        q = k = self.with_pos_embed(tgt2, query_pos)
        tgt2 = self.self_attn(q, k, value=tgt2, attn_mask=tgt_mask,
                              key_padding_mask=tgt_key_padding_mask)[0]
        tgt = tgt + self.dropout(tgt2)

        return tgt

    def forward(self, tgt,
                tgt_mask = None,
                tgt_key_padding_mask = None,
                query_pos = None):
        if self.normalize_before:
            return self.forward_pre(tgt, tgt_mask,
                                    tgt_key_padding_mask, query_pos)
        return self.forward_post(tgt, tgt_mask,
                                 tgt_key_padding_mask, query_pos)


class CrossAttentionLayer(nn.Module):

    def __init__(self, d_model, nhead, dropout=0.0,
                 activation="relu", normalize_before=False):
        super().__init__()
        self.multihead_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout)

        self.norm = nn.LayerNorm(d_model)
        self.dropout = nn.Dropout(dropout)

        self.activation = _get_activation_fn(activation)
        self.normalize_before = normalize_before

        self._reset_parameters()

    def _reset_parameters(self):
        for p in self.parameters():
            if p.dim() > 1:
                nn.init.xavier_uniform_(p)

    def with_pos_embed(self, tensor, pos):
        return tensor if pos is None else tensor + pos

    def forward_post(self, tgt, memory,
                     memory_mask = None,
                     memory_key_padding_mask = None,
                     pos = None,
                     query_pos = None):
        tgt2 = self.multihead_attn(query=self.with_pos_embed(tgt, query_pos),
                                   key=self.with_pos_embed(memory, pos),
                                   value=memory, attn_mask=memory_mask,
                                   key_padding_mask=memory_key_padding_mask)[0]
        tgt = tgt + self.dropout(tgt2)
        tgt = self.norm(tgt)

        return tgt

    def forward_pre(self, tgt, memory,
                    memory_mask = None,
                    memory_key_padding_mask = None,
                    pos = None,
                    query_pos = None):
        tgt2 = self.norm(tgt)

        tgt2 = self.multihead_attn(query=self.with_pos_embed(tgt2, query_pos),
                                   key=self.with_pos_embed(memory, pos),
                                   value=memory, attn_mask=memory_mask,
                                   key_padding_mask=memory_key_padding_mask)[0]
        tgt = tgt + self.dropout(tgt2)

        return tgt

    def forward(self, tgt, memory,
                memory_mask = None,
                memory_key_padding_mask = None,
                pos = None,
                query_pos = None):
        if self.normalize_before:
            return self.forward_pre(tgt, memory, memory_mask,
                                    memory_key_padding_mask, pos, query_pos)
        return self.forward_post(tgt, memory, memory_mask,
                                 memory_key_padding_mask, pos, query_pos)

class FFNLayer(nn.Module):

    def __init__(self, d_model, dim_feedforward=2048, dropout=0.0,
                 activation="relu", normalize_before=False):
        super().__init__()
        # Implementation of Feedforward model
        self.linear1 = nn.Linear(d_model, dim_feedforward)
        self.dropout = nn.Dropout(dropout)
        self.linear2 = nn.Linear(dim_feedforward, d_model)

        self.norm = nn.LayerNorm(d_model)

        self.activation = _get_activation_fn(activation)
        self.normalize_before = normalize_before

        self._reset_parameters()

    def _reset_parameters(self):
        for p in self.parameters():
            if p.dim() > 1:
                nn.init.xavier_uniform_(p)

    def with_pos_embed(self, tensor, pos):
        return tensor if pos is None else tensor + pos

    def forward_post(self, tgt):
        tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt))))
        tgt = tgt + self.dropout(tgt2)
        tgt = self.norm(tgt)
        return tgt

    def forward_pre(self, tgt):
        tgt2 = self.norm(tgt)
        tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt2))))
        tgt = tgt + self.dropout(tgt2)
        return tgt

    def forward(self, tgt):
        if self.normalize_before:
            return self.forward_pre(tgt)
        return self.forward_post(tgt)


def _get_activation_fn(activation):
    """Return an activation function given a string"""
    if activation == "relu":
        return F.relu
    if activation == "gelu":
        return F.gelu
    if activation == "glu":
        return F.glu
    raise RuntimeError(F"activation should be relu/gelu, not {activation}.")



models/mask3d.py
aminebdj/3D-OWIS
# Copyright (c) Facebook, Inc. and its affiliates.
# Modified by Bowen Cheng from https://github.com/facebookresearch/detr/blob/master/models/matcher.py
"""
Modules to compute the matching cost and solve the corresponding LSAP.
"""
import torch
import torch.nn.functional as F
from scipy.optimize import linear_sum_assignment
from torch import nn
from torch.cuda.amp import autocast
from detectron2.projects.point_rend.point_features import point_sample


#######################################################################################################################
#######################################################################################################################
#######################################################################################################################
##########################
from datasets.scannet200.scannet200_constants import CLASS_LABELS_200
from datasets.scannet200.scannet200_splits import HEAD_CATS_SCANNET_200, COMMON_CATS_SCANNET_200, TAIL_CATS_SCANNET_200
############################
CLASS_LABELS_200 = list(CLASS_LABELS_200)
CLASS_LABELS_200.remove('floor')
CLASS_LABELS_200.remove('wall')

MAP_ID_TO_STRING = {i : CLASS_LABELS_200[i] for i in range(len(CLASS_LABELS_200))}
MAP_ID_TO_STRING['background'] = 253

#######################################################################################################################
#######################################################################################################################
#######################################################################################################################


def batch_dice_loss(inputs: torch.Tensor, targets: torch.Tensor):
    """
    Compute the DICE loss, similar to generalized IOU for masks
    Args:
        inputs: A float tensor of arbitrary shape.
                The predictions for each example.
        targets: A float tensor with the same shape as inputs. Stores the binary
                 classification label for each element in inputs
                (0 for the negative class and 1 for the positive class).
    """
    inputs = inputs.sigmoid()
    inputs = inputs.flatten(1)
    numerator = 2 * torch.einsum("nc,mc->nm", inputs, targets)
    denominator = inputs.sum(-1)[:, None] + targets.sum(-1)[None, :]
    loss = 1 - (numerator + 1) / (denominator + 1)
    return loss


batch_dice_loss_jit = torch.jit.script(
    batch_dice_loss
)  # type: torch.jit.ScriptModule


def batch_sigmoid_ce_loss(inputs: torch.Tensor, targets: torch.Tensor):
    """
    Args:
        inputs: A float tensor of arbitrary shape.
                The predictions for each example.
        targets: A float tensor with the same shape as inputs. Stores the binary
                 classification label for each element in inputs
                (0 for the negative class and 1 for the positive class).
    Returns:
        Loss tensor
    """
    hw = inputs.shape[1]

    pos = F.binary_cross_entropy_with_logits(
        inputs, torch.ones_like(inputs), reduction="none"
    )
    neg = F.binary_cross_entropy_with_logits(
        inputs, torch.zeros_like(inputs), reduction="none"
    )

    loss = torch.einsum("nc,mc->nm", pos, targets) + torch.einsum(
        "nc,mc->nm", neg, (1 - targets)
    )

    return loss / hw


batch_sigmoid_ce_loss_jit = torch.jit.script(
    batch_sigmoid_ce_loss
)  # type: torch.jit.ScriptModule


class HungarianMatcher(nn.Module):
    """This class computes an assignment between the targets and the predictions of the network

    For efficiency reasons, the targets don't include the no_object. Because of this, in general,
    there are more predictions than targets. In this case, we do a 1-to-1 matching of the best predictions,
    while the others are un-matched (and thus treated as non-objects).
    """

    def __init__(self,cost_class: float = 1, cost_mask: float = 1, cost_dice: float = 1, num_points: int = 0, train_oracle: bool = False):
        """Creates the matcher

        Params:
            cost_class: This is the relative weight of the classification error in the matching cost
            cost_mask: This is the relative weight of the focal loss of the binary mask in the matching cost
            cost_dice: This is the relative weight of the dice loss of the binary mask in the matching cost
        """
        super().__init__()
        self.cost_class = cost_class
        self.cost_mask = cost_mask
        self.cost_dice = cost_dice
        assert cost_class != 0 or cost_mask != 0 or cost_dice != 0, "all costs cant be 0"

        self.num_points = num_points
        self.train_oracle = train_oracle
    @torch.no_grad()
    def memory_efficient_forward(self, outputs, targets, mask_type):
        """More memory-friendly matching"""
        bs, num_queries = outputs["pred_logits"].shape[:2]

        indices = []

        # Iterate through batch size
        for b in range(bs):
            out_prob = outputs["pred_logits"][b].softmax(-1)  # [num_queries, num_classes]
            tgt_ids = targets[b]["labels"].clone()
            
            # Compute the classification cost. Contrary to the loss, we don't use the NLL,
            # but approximate it in 1 - proba[target class].
            # The 1 is a constant that doesn't change the matching, it can be ommitted.
            # if not self.train_oracle:
            #     filter_ignore = tgt_ids == 253 
            # else:    
            filter_ignore = tgt_ids == 253
                
            tgt_ids[filter_ignore] = 0
            cost_class = -out_prob[:, tgt_ids]
            cost_class[:, filter_ignore] = -1.  # for ignore classes pretend perfect match ;) TODO better worst class match?

            out_mask = outputs['pred_masks'][b].T  # [num_queries, H_pred, W_pred]
            # gt masks are already padded when preparing target
            tgt_mask = targets[b][mask_type].to(out_mask)

            if self.num_points != -1:
                point_idx = torch.randperm(tgt_mask.shape[1],
                                           device=tgt_mask.device)[:int(self.num_points*tgt_mask.shape[1])]
                #point_idx = torch.randint(0, tgt_mask.shape[1], size=(self.num_points,), device=tgt_mask.device)
            else:
                # sample all points
                point_idx = torch.arange(tgt_mask.shape[1], device=tgt_mask.device)

            # out_mask = out_mask[:, None]
            # tgt_mask = tgt_mask[:, None]
            # all masks share the same set of points for efficient matching!
            # point_coords = torch.rand(1, self.num_points, 2, device=out_mask.device)
            # get gt labels
            # tgt_mask = point_sample(
            #     tgt_mask,
            #     point_coords.repeat(tgt_mask.shape[0], 1, 1),
            #     align_corners=False,
            # ).squeeze(1)

            # out_mask = point_sample(
            #     out_mask,
            #     point_coords.repeat(out_mask.shape[0], 1, 1),
            #     align_corners=False,
            # ).squeeze(1)
            
            
            
        #    #======================================================================================================================================
        #    #======================================================================================================================================
        #    #======================================================================================================================================
        #    #====================================================================================================================================== 
        #     if self.save_recall:
        #         threshold = 0.5
        #         intersection = targets[b]['segment_mask'].float()@(torch.nn.Sigmoid()(outputs['pred_masks'][b])>threshold).float()
        #         union = torch.stack(tuple(torch.sum((targets[b]['segment_mask'][inst_id,:].float()+(torch.nn.Sigmoid()(outputs['pred_masks'][b])>threshold).float().T),dim=1)-intersection[inst_id,:] for inst_id in range(targets[b]['segment_mask'].shape[0])), dim = 0)
        #         iou = intersection/union
        #         for i, label_itr in enumerate(targets[b]['labels']): 
        #             gt_iou["gt_id"].append(label_itr.item())
        #             gt_iou["gt_label"].append(MAP_ID_TO_STRING[label_itr.item()])
        #             gt_iou["iou"].append(iou[i,:])
                
        #    #======================================================================================================================================
        #    #======================================================================================================================================
        #    #======================================================================================================================================
        #    #======================================================================================================================================
            
            
            
            
            with autocast(enabled=False):
                out_mask = out_mask.float()
                tgt_mask = tgt_mask.float()
                # Compute the focal loss between masks
                cost_mask = batch_sigmoid_ce_loss_jit(out_mask[:, point_idx], tgt_mask[:, point_idx])

                # Compute the dice loss betwen masks
                cost_dice = batch_dice_loss_jit(out_mask[:, point_idx], tgt_mask[:, point_idx])
            
            # Final cost matrix
            C = (
                self.cost_mask * cost_mask
                + self.cost_class * cost_class
                + self.cost_dice * cost_dice
            )
            C = C.reshape(num_queries, -1).cpu()

            indices.append(linear_sum_assignment(C))

        return [
            (torch.as_tensor(i, dtype=torch.int64), torch.as_tensor(j, dtype=torch.int64))
            for i, j in indices
        ]

    @torch.no_grad()
    def forward(self, outputs, targets, mask_type):
        """Performs the matching

        Params:
            outputs: This is a dict that contains at least these entries:
                 "pred_logits": Tensor of dim [batch_size, num_queries, num_classes] with the classification logits
                 "pred_masks": Tensor of dim [batch_size, num_queries, H_pred, W_pred] with the predicted masks

            targets: This is a list of targets (len(targets) = batch_size), where each target is a dict containing:
                 "labels": Tensor of dim [num_target_boxes] (where num_target_boxes is the number of ground-truth
                           objects in the target) containing the class labels
                 "masks": Tensor of dim [num_target_boxes, H_gt, W_gt] containing the target masks

        Returns:
            A list of size batch_size, containing tuples of (index_i, index_j) where:
                - index_i is the indices of the selected predictions (in order)
                - index_j is the indices of the corresponding selected targets (in order)
            For each batch element, it holds:
                len(index_i) = len(index_j) = min(num_queries, num_target_boxes)
        """
        return self.memory_efficient_forward(outputs, targets, mask_type)

    def __repr__(self, _repr_indent=4):
        head = "Matcher " + self.__class__.__name__
        body = [
            "cost_class: {}".format(self.cost_class),
            "cost_mask: {}".format(self.cost_mask),
            "cost_dice: {}".format(self.cost_dice),
        ]
        lines = [head] + [" " * _repr_indent + line for line in body]
        return "\n".join(lines)

models/matcher.py
aminebdj/3D-OWIS
# Copyright (c) Facebook, Inc. and its affiliates.
# Modified by Bowen Cheng from https://github.com/facebookresearch/detr/blob/master/util/misc.py
"""
Misc functions, including distributed helpers.

Mostly copy-paste from torchvision references.
"""
from typing import List, Optional

import torch
import torch.distributed as dist
import torchvision
from torch import Tensor


def _max_by_axis(the_list):
    # type: (List[List[int]]) -> List[int]
    maxes = the_list[0]
    for sublist in the_list[1:]:
        for index, item in enumerate(sublist):
            maxes[index] = max(maxes[index], item)
    return maxes


class NestedTensor(object):
    def __init__(self, tensors, mask: Optional[Tensor]):
        self.tensors = tensors
        self.mask = mask

    def to(self, device):
        # type: (Device) -> NestedTensor # noqa
        cast_tensor = self.tensors.to(device)
        mask = self.mask
        if mask is not None:
            assert mask is not None
            cast_mask = mask.to(device)
        else:
            cast_mask = None
        return NestedTensor(cast_tensor, cast_mask)

    def decompose(self):
        return self.tensors, self.mask

    def __repr__(self):
        return str(self.tensors)


def nested_tensor_from_tensor_list(tensor_list: List[Tensor]):
    # TODO make this more general
    if tensor_list[0].ndim == 3:
        if torchvision._is_tracing():
            # nested_tensor_from_tensor_list() does not export well to ONNX
            # call _onnx_nested_tensor_from_tensor_list() instead
            return _onnx_nested_tensor_from_tensor_list(tensor_list)

        # TODO make it support different-sized images
        max_size = _max_by_axis([list(img.shape) for img in tensor_list])
        # min_size = tuple(min(s) for s in zip(*[img.shape for img in tensor_list]))
        batch_shape = [len(tensor_list)] + max_size
        b, c, h, w = batch_shape
        dtype = tensor_list[0].dtype
        device = tensor_list[0].device
        tensor = torch.zeros(batch_shape, dtype=dtype, device=device)
        mask = torch.ones((b, h, w), dtype=torch.bool, device=device)
        for img, pad_img, m in zip(tensor_list, tensor, mask):
            pad_img[: img.shape[0], : img.shape[1], : img.shape[2]].copy_(img)
            m[: img.shape[1], : img.shape[2]] = False
    else:
        raise ValueError("not supported")
    return NestedTensor(tensor, mask)


# _onnx_nested_tensor_from_tensor_list() is an implementation of
# nested_tensor_from_tensor_list() that is supported by ONNX tracing.
@torch.jit.unused
def _onnx_nested_tensor_from_tensor_list(tensor_list: List[Tensor]) -> NestedTensor:
    max_size = []
    for i in range(tensor_list[0].dim()):
        max_size_i = torch.max(
            torch.stack([img.shape[i] for img in tensor_list]).to(torch.float32)
        ).to(torch.int64)
        max_size.append(max_size_i)
    max_size = tuple(max_size)

    # work around for
    # pad_img[: img.shape[0], : img.shape[1], : img.shape[2]].copy_(img)
    # m[: img.shape[1], :img.shape[2]] = False
    # which is not yet supported in onnx
    padded_imgs = []
    padded_masks = []
    for img in tensor_list:
        padding = [(s1 - s2) for s1, s2 in zip(max_size, tuple(img.shape))]
        padded_img = torch.nn.functional.pad(img, (0, padding[2], 0, padding[1], 0, padding[0]))
        padded_imgs.append(padded_img)

        m = torch.zeros_like(img[0], dtype=torch.int, device=img.device)
        padded_mask = torch.nn.functional.pad(m, (0, padding[2], 0, padding[1]), "constant", 1)
        padded_masks.append(padded_mask.to(torch.bool))

    tensor = torch.stack(padded_imgs)
    mask = torch.stack(padded_masks)

    return NestedTensor(tensor, mask=mask)


def is_dist_avail_and_initialized():
    if not dist.is_available():
        return False
    if not dist.is_initialized():
        return False
    return True

models/misc.py
aminebdj/3D-OWIS
import sys

if sys.version_info[:2] >= (3, 8):
    from collections.abc import Sequence
else:
    from collections import Sequence

from enum import Enum

import torch.nn as nn
import MinkowskiEngine as ME


class NormType(Enum):
    BATCH_NORM = 0
    INSTANCE_NORM = 1
    INSTANCE_BATCH_NORM = 2


def get_norm(norm_type, n_channels, D, bn_momentum=0.1):
    if norm_type == NormType.BATCH_NORM:
        return ME.MinkowskiBatchNorm(n_channels, momentum=bn_momentum)
    elif norm_type == NormType.INSTANCE_NORM:
        return ME.MinkowskiInstanceNorm(n_channels)
    elif norm_type == NormType.INSTANCE_BATCH_NORM:
        return nn.Sequential(
            ME.MinkowskiInstanceNorm(n_channels),
            ME.MinkowskiBatchNorm(n_channels, momentum=bn_momentum),
        )
    else:
        raise ValueError(f"Norm type: {norm_type} not supported")


class ConvType(Enum):
    """
    Define the kernel region type
    """

    HYPERCUBE = 0, "HYPERCUBE"
    SPATIAL_HYPERCUBE = 1, "SPATIAL_HYPERCUBE"
    SPATIO_TEMPORAL_HYPERCUBE = 2, "SPATIO_TEMPORAL_HYPERCUBE"
    HYPERCROSS = 3, "HYPERCROSS"
    SPATIAL_HYPERCROSS = 4, "SPATIAL_HYPERCROSS"
    SPATIO_TEMPORAL_HYPERCROSS = 5, "SPATIO_TEMPORAL_HYPERCROSS"
    SPATIAL_HYPERCUBE_TEMPORAL_HYPERCROSS = 6, "SPATIAL_HYPERCUBE_TEMPORAL_HYPERCROSS "

    def __new__(cls, value, name):
        member = object.__new__(cls)
        member._value_ = value
        member.fullname = name
        return member

    def __int__(self):
        return self.value


# Covert the ConvType var to a RegionType var
conv_to_region_type = {
    # kernel_size = [k, k, k, 1]
    ConvType.HYPERCUBE: ME.RegionType.HYPER_CUBE,
    ConvType.SPATIAL_HYPERCUBE: ME.RegionType.HYPER_CUBE,
    ConvType.SPATIO_TEMPORAL_HYPERCUBE: ME.RegionType.HYPER_CUBE,
    ConvType.HYPERCROSS: ME.RegionType.HYPER_CROSS,
    ConvType.SPATIAL_HYPERCROSS: ME.RegionType.HYPER_CROSS,
    ConvType.SPATIO_TEMPORAL_HYPERCROSS: ME.RegionType.HYPER_CROSS,
    ConvType.SPATIAL_HYPERCUBE_TEMPORAL_HYPERCROSS: ME.RegionType.HYPER_CUBE,  # JONAS CHANGE from HYBRID
}

# int_to_region_type = {m.value: m for m in ME.RegionType}
int_to_region_type = {m: ME.RegionType(m) for m in range(3)}


def convert_region_type(region_type):
    """
    Convert the integer region_type to the corresponding RegionType enum object.
    """
    return int_to_region_type[region_type]


def convert_conv_type(conv_type, kernel_size, D):
    assert isinstance(conv_type, ConvType), "conv_type must be of ConvType"
    region_type = conv_to_region_type[conv_type]
    axis_types = None
    if conv_type == ConvType.SPATIAL_HYPERCUBE:
        # No temporal convolution
        if isinstance(kernel_size, Sequence):
            kernel_size = kernel_size[:3]
        else:
            kernel_size = [
                kernel_size,
            ] * 3
        if D == 4:
            kernel_size.append(1)
    elif conv_type == ConvType.SPATIO_TEMPORAL_HYPERCUBE:
        # conv_type conversion already handled
        assert D == 4
    elif conv_type == ConvType.HYPERCUBE:
        # conv_type conversion already handled
        pass
    elif conv_type == ConvType.SPATIAL_HYPERCROSS:
        if isinstance(kernel_size, Sequence):
            kernel_size = kernel_size[:3]
        else:
            kernel_size = [
                kernel_size,
            ] * 3
        if D == 4:
            kernel_size.append(1)
    elif conv_type == ConvType.HYPERCROSS:
        # conv_type conversion already handled
        pass
    elif conv_type == ConvType.SPATIO_TEMPORAL_HYPERCROSS:
        # conv_type conversion already handled
        assert D == 4
    elif conv_type == ConvType.SPATIAL_HYPERCUBE_TEMPORAL_HYPERCROSS:
        # Define the CUBIC conv kernel for spatial dims and CROSS conv for temp dim
        axis_types = [
            ME.RegionType.HYPER_CUBE,
        ] * 3
        if D == 4:
            axis_types.append(ME.RegionType.HYPER_CROSS)
    return region_type, axis_types, kernel_size


def conv(
    in_planes,
    out_planes,
    kernel_size,
    stride=1,
    dilation=1,
    bias=False,
    conv_type=ConvType.HYPERCUBE,
    D=-1,
):
    assert D > 0, "Dimension must be a positive integer"
    region_type, axis_types, kernel_size = convert_conv_type(conv_type, kernel_size, D)
    kernel_generator = ME.KernelGenerator(
        kernel_size,
        stride,
        dilation,
        region_type=region_type,
        axis_types=None, # axis_types JONAS
        dimension=D,
    )

    return ME.MinkowskiConvolution(
        in_channels=in_planes,
        out_channels=out_planes,
        kernel_size=kernel_size,
        stride=stride,
        dilation=dilation,
        bias=bias,
        kernel_generator=kernel_generator,
        dimension=D,
    )


def conv_tr(
    in_planes,
    out_planes,
    kernel_size,
    upsample_stride=1,
    dilation=1,
    bias=False,
    conv_type=ConvType.HYPERCUBE,
    D=-1,
):
    assert D > 0, "Dimension must be a positive integer"
    region_type, axis_types, kernel_size = convert_conv_type(conv_type, kernel_size, D)
    kernel_generator = ME.KernelGenerator(
        kernel_size,
        upsample_stride,
        dilation,
        region_type=region_type,
        axis_types=axis_types,
        dimension=D,
    )

    return ME.MinkowskiConvolutionTranspose(
        in_channels=in_planes,
        out_channels=out_planes,
        kernel_size=kernel_size,
        stride=upsample_stride,
        dilation=dilation,
        bias=bias,
        kernel_generator=kernel_generator,
        dimension=D,
    )


def avg_pool(
    kernel_size,
    stride=1,
    dilation=1,
    conv_type=ConvType.HYPERCUBE,
    in_coords_key=None,
    D=-1,
):
    assert D > 0, "Dimension must be a positive integer"
    region_type, axis_types, kernel_size = convert_conv_type(conv_type, kernel_size, D)
    kernel_generator = ME.KernelGenerator(
        kernel_size,
        stride,
        dilation,
        region_type=region_type,
        axis_types=axis_types,
        dimension=D,
    )

    return ME.MinkowskiAvgPooling(
        kernel_size=kernel_size,
        stride=stride,
        dilation=dilation,
        kernel_generator=kernel_generator,
        dimension=D,
    )


def avg_unpool(kernel_size, stride=1, dilation=1, conv_type=ConvType.HYPERCUBE, D=-1):
    assert D > 0, "Dimension must be a positive integer"
    region_type, axis_types, kernel_size = convert_conv_type(conv_type, kernel_size, D)
    kernel_generator = ME.KernelGenerator(
        kernel_size,
        stride,
        dilation,
        region_type=region_type,
        axis_types=axis_types,
        dimension=D,
    )

    return ME.MinkowskiAvgUnpooling(
        kernel_size=kernel_size,
        stride=stride,
        dilation=dilation,
        kernel_generator=kernel_generator,
        dimension=D,
    )


def sum_pool(kernel_size, stride=1, dilation=1, conv_type=ConvType.HYPERCUBE, D=-1):
    assert D > 0, "Dimension must be a positive integer"
    region_type, axis_types, kernel_size = convert_conv_type(conv_type, kernel_size, D)
    kernel_generator = ME.KernelGenerator(
        kernel_size,
        stride,
        dilation,
        region_type=region_type,
        axis_types=axis_types,
        dimension=D,
    )

    return ME.MinkowskiSumPooling(
        kernel_size=kernel_size,
        stride=stride,
        dilation=dilation,
        kernel_generator=kernel_generator,
        dimension=D,
    )

models/modules/common.py
aminebdj/3D-OWIS
import torch.nn as nn
from MinkowskiEngine import MinkowskiReLU

from models.modules.common import ConvType, NormType, conv, get_norm


class BasicBlockBase(nn.Module):
    expansion = 1
    NORM_TYPE = NormType.BATCH_NORM

    def __init__(
        self,
        inplanes,
        planes,
        stride=1,
        dilation=1,
        downsample=None,
        conv_type=ConvType.HYPERCUBE,
        bn_momentum=0.1,
        D=3,
    ):
        super().__init__()

        self.conv1 = conv(
            inplanes,
            planes,
            kernel_size=3,
            stride=stride,
            dilation=dilation,
            conv_type=conv_type,
            D=D,
        )
        self.norm1 = get_norm(self.NORM_TYPE, planes, D, bn_momentum=bn_momentum)
        self.conv2 = conv(
            planes,
            planes,
            kernel_size=3,
            stride=1,
            dilation=dilation,
            bias=False,
            conv_type=conv_type,
            D=D,
        )
        self.norm2 = get_norm(self.NORM_TYPE, planes, D, bn_momentum=bn_momentum)
        self.relu = MinkowskiReLU(inplace=True)
        self.downsample = downsample

    def forward(self, x):
        residual = x

        out = self.conv1(x)
        out = self.norm1(out)
        out = self.relu(out)

        out = self.conv2(out)
        out = self.norm2(out)

        if self.downsample is not None:
            residual = self.downsample(x)

        out += residual
        out = self.relu(out)

        return out


class BasicBlock(BasicBlockBase):
    NORM_TYPE = NormType.BATCH_NORM


class BasicBlockIN(BasicBlockBase):
    NORM_TYPE = NormType.INSTANCE_NORM


class BasicBlockINBN(BasicBlockBase):
    NORM_TYPE = NormType.INSTANCE_BATCH_NORM


class BottleneckBase(nn.Module):
    expansion = 4
    NORM_TYPE = NormType.BATCH_NORM

    def __init__(
        self,
        inplanes,
        planes,
        stride=1,
        dilation=1,
        downsample=None,
        conv_type=ConvType.HYPERCUBE,
        bn_momentum=0.1,
        D=3,
    ):
        super().__init__()
        self.conv1 = conv(inplanes, planes, kernel_size=1, D=D)
        self.norm1 = get_norm(self.NORM_TYPE, planes, D, bn_momentum=bn_momentum)

        self.conv2 = conv(
            planes,
            planes,
            kernel_size=3,
            stride=stride,
            dilation=dilation,
            conv_type=conv_type,
            D=D,
        )
        self.norm2 = get_norm(self.NORM_TYPE, planes, D, bn_momentum=bn_momentum)

        self.conv3 = conv(planes, planes * self.expansion, kernel_size=1, D=D)
        self.norm3 = get_norm(
            self.NORM_TYPE, planes * self.expansion, D, bn_momentum=bn_momentum
        )

        self.relu = MinkowskiReLU(inplace=True)
        self.downsample = downsample

    def forward(self, x):
        residual = x

        out = self.conv1(x)
        out = self.norm1(out)
        out = self.relu(out)

        out = self.conv2(out)
        out = self.norm2(out)
        out = self.relu(out)

        out = self.conv3(out)
        out = self.norm3(out)

        if self.downsample is not None:
            residual = self.downsample(x)

        out += residual
        out = self.relu(out)

        return out


class Bottleneck(BottleneckBase):
    NORM_TYPE = NormType.BATCH_NORM


class BottleneckIN(BottleneckBase):
    NORM_TYPE = NormType.INSTANCE_NORM


class BottleneckINBN(BottleneckBase):
    NORM_TYPE = NormType.INSTANCE_BATCH_NORM

models/modules/resnet_block.py
aminebdj/3D-OWIS
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
"""
Various positional encodings for the transformer.
"""
import math
import torch
from torch import nn
import numpy as np
#from utils.pc_util import shift_scale_points


def shift_scale_points(pred_xyz, src_range, dst_range=None):
    """
    pred_xyz: B x N x 3
    src_range: [[B x 3], [B x 3]] - min and max XYZ coords
    dst_range: [[B x 3], [B x 3]] - min and max XYZ coords
    """
    if dst_range is None:
        dst_range = [
            torch.zeros((src_range[0].shape[0], 3), device=src_range[0].device),
            torch.ones((src_range[0].shape[0], 3), device=src_range[0].device),
        ]

    if pred_xyz.ndim == 4:
        src_range = [x[:, None] for x in src_range]
        dst_range = [x[:, None] for x in dst_range]

    assert src_range[0].shape[0] == pred_xyz.shape[0]
    assert dst_range[0].shape[0] == pred_xyz.shape[0]
    assert src_range[0].shape[-1] == pred_xyz.shape[-1]
    assert src_range[0].shape == src_range[1].shape
    assert dst_range[0].shape == dst_range[1].shape
    assert src_range[0].shape == dst_range[1].shape

    src_diff = src_range[1][:, None, :] - src_range[0][:, None, :]
    dst_diff = dst_range[1][:, None, :] - dst_range[0][:, None, :]
    prop_xyz = (
        ((pred_xyz - src_range[0][:, None, :]) * dst_diff) / src_diff
    ) + dst_range[0][:, None, :]
    return prop_xyz


class PositionEmbeddingCoordsSine(nn.Module):
    def __init__(
        self,
        temperature=10000,
        normalize=False,
        scale=None,
        pos_type="fourier",
        d_pos=None,
        d_in=3,
        gauss_scale=1.0
    ):
        super().__init__()
        self.d_pos = d_pos
        self.temperature = temperature
        self.normalize = normalize
        if scale is not None and normalize is False:
            raise ValueError("normalize should be True if scale is passed")
        if scale is None:
            scale = 2 * math.pi
        assert pos_type in ["sine", "fourier"]
        self.pos_type = pos_type
        self.scale = scale
        if pos_type == "fourier":
            assert d_pos is not None
            assert d_pos % 2 == 0
            # define a gaussian matrix input_ch -> output_ch
            B = torch.empty((d_in, d_pos // 2)).normal_()
            B *= gauss_scale
            self.register_buffer("gauss_B", B)
            self.d_pos = d_pos

    def get_sine_embeddings(self, xyz, num_channels, input_range):
        num_channels = self.d_pos
        # clone coords so that shift/scale operations do not affect original tensor
        orig_xyz = xyz
        xyz = orig_xyz.clone()

        ncoords = xyz.shape[1]
        if self.normalize:
            xyz = shift_scale_points(xyz, src_range=input_range)

        ndim = num_channels // xyz.shape[2]
        if ndim % 2 != 0:
            ndim -= 1
        # automatically handle remainder by assiging it to the first dim
        rems = num_channels - (ndim * xyz.shape[2])

        assert (
            ndim % 2 == 0
        ), f"Cannot handle odd sized ndim={ndim} where num_channels={num_channels} and xyz={xyz.shape}"

        final_embeds = []
        prev_dim = 0

        for d in range(xyz.shape[2]):
            cdim = ndim
            if rems > 0:
                # add remainder in increments of two to maintain even size
                cdim += 2
                rems -= 2

            if cdim != prev_dim:
                dim_t = torch.arange(cdim, dtype=torch.float32, device=xyz.device)
                dim_t = self.temperature ** (2 * (dim_t // 2) / cdim)

            # create batch x cdim x nccords embedding
            raw_pos = xyz[:, :, d]
            if self.scale:
                raw_pos *= self.scale
            pos = raw_pos[:, :, None] / dim_t
            pos = torch.stack(
                (pos[:, :, 0::2].sin(), pos[:, :, 1::2].cos()), dim=3
            ).flatten(2)
            final_embeds.append(pos)
            prev_dim = cdim

        final_embeds = torch.cat(final_embeds, dim=2).permute(0, 2, 1)
        return final_embeds

    def get_fourier_embeddings(self, xyz, num_channels=None, input_range=None):
        # Follows - https://people.eecs.berkeley.edu/~bmild/fourfeat/index.html

        if num_channels is None:
            num_channels = self.gauss_B.shape[1] * 2

        bsize, npoints = xyz.shape[0], xyz.shape[1]
        assert num_channels > 0 and num_channels % 2 == 0
        d_in, max_d_out = self.gauss_B.shape[0], self.gauss_B.shape[1]
        d_out = num_channels // 2
        assert d_out <= max_d_out
        assert d_in == xyz.shape[-1]

        # clone coords so that shift/scale operations do not affect original tensor
        orig_xyz = xyz
        xyz = orig_xyz.clone()

        ncoords = xyz.shape[1]
        if self.normalize:
            xyz = shift_scale_points(xyz, src_range=input_range)

        xyz *= 2 * np.pi
        xyz_proj = torch.mm(xyz.view(-1, d_in), self.gauss_B[:, :d_out]).view(
            bsize, npoints, d_out
        )
        final_embeds = [xyz_proj.sin(), xyz_proj.cos()]

        # return batch x d_pos x npoints embedding
        final_embeds = torch.cat(final_embeds, dim=2).permute(0, 2, 1)
        return final_embeds

    def forward(self, xyz, num_channels=None, input_range=None):
        assert isinstance(xyz, torch.Tensor)
        assert xyz.ndim == 3
        # xyz is batch x npoints x 3
        if self.pos_type == "sine":
            with torch.no_grad():
                out = self.get_sine_embeddings(xyz, num_channels, input_range)
        elif self.pos_type == "fourier":
            with torch.no_grad():
                out = self.get_fourier_embeddings(xyz, num_channels, input_range)
        else:
            raise ValueError(f"Unknown {self.pos_type}")

        return out

    def extra_repr(self):
        st = f"type={self.pos_type}, scale={self.scale}, normalize={self.normalize}"
        if hasattr(self, "gauss_B"):
            st += (
                f", gaussB={self.gauss_B.shape}, gaussBsum={self.gauss_B.sum().item()}"
            )
        return st
models/position_embedding.py
aminebdj/3D-OWIS
import MinkowskiEngine.MinkowskiOps as me
from MinkowskiEngine import MinkowskiReLU

from models.resnet import ResNetBase, get_norm
from models.modules.common import ConvType, NormType, conv, conv_tr
from models.modules.resnet_block import BasicBlock, Bottleneck


class Res16UNetBase(ResNetBase):
    BLOCK = None
    PLANES = (32, 64, 128, 256, 256, 256, 256, 256)
    DILATIONS = (1, 1, 1, 1, 1, 1, 1, 1)
    LAYERS = (2, 2, 2, 2, 2, 2, 2, 2)
    INIT_DIM = 32
    OUT_PIXEL_DIST = 1
    NORM_TYPE = NormType.BATCH_NORM
    NON_BLOCK_CONV_TYPE = ConvType.SPATIAL_HYPERCUBE
    CONV_TYPE = ConvType.SPATIAL_HYPERCUBE_TEMPORAL_HYPERCROSS

    # To use the model, must call initialize_coords before forward pass.
    # Once data is processed, call clear to reset the model before calling initialize_coords
    def __init__(self, in_channels, out_channels, config, D=3, out_fpn=False, **kwargs):
        super().__init__(in_channels, out_channels, config, D)
        self.out_fpn = out_fpn

    def network_initialization(self, in_channels, out_channels, config, D):
        # Setup net_metadata
        dilations = self.DILATIONS
        bn_momentum = config.bn_momentum

        def space_n_time_m(n, m):
            return n if D == 3 else [n, n, n, m]

        if D == 4:
            self.OUT_PIXEL_DIST = space_n_time_m(self.OUT_PIXEL_DIST, 1)

        # Output of the first conv concated to conv6
        self.inplanes = self.INIT_DIM
        self.conv0p1s1 = conv(
            in_channels,
            self.inplanes,
            kernel_size=space_n_time_m(config.conv1_kernel_size, 1),
            stride=1,
            dilation=1,
            conv_type=self.NON_BLOCK_CONV_TYPE,
            D=D,
        )

        self.bn0 = get_norm(self.NORM_TYPE, self.inplanes, D, bn_momentum=bn_momentum)

        self.conv1p1s2 = conv(
            self.inplanes,
            self.inplanes,
            kernel_size=space_n_time_m(2, 1),
            stride=space_n_time_m(2, 1),
            dilation=1,
            conv_type=self.NON_BLOCK_CONV_TYPE,
            D=D,
        )
        self.bn1 = get_norm(self.NORM_TYPE, self.inplanes, D, bn_momentum=bn_momentum)
        self.block1 = self._make_layer(
            self.BLOCK,
            self.PLANES[0],
            self.LAYERS[0],
            dilation=dilations[0],
            norm_type=self.NORM_TYPE,
            bn_momentum=bn_momentum,
        )

        self.conv2p2s2 = conv(
            self.inplanes,
            self.inplanes,
            kernel_size=space_n_time_m(2, 1),
            stride=space_n_time_m(2, 1),
            dilation=1,
            conv_type=self.NON_BLOCK_CONV_TYPE,
            D=D,
        )
        self.bn2 = get_norm(self.NORM_TYPE, self.inplanes, D, bn_momentum=bn_momentum)
        self.block2 = self._make_layer(
            self.BLOCK,
            self.PLANES[1],
            self.LAYERS[1],
            dilation=dilations[1],
            norm_type=self.NORM_TYPE,
            bn_momentum=bn_momentum,
        )

        self.conv3p4s2 = conv(
            self.inplanes,
            self.inplanes,
            kernel_size=space_n_time_m(2, 1),
            stride=space_n_time_m(2, 1),
            dilation=1,
            conv_type=self.NON_BLOCK_CONV_TYPE,
            D=D,
        )
        self.bn3 = get_norm(self.NORM_TYPE, self.inplanes, D, bn_momentum=bn_momentum)
        self.block3 = self._make_layer(
            self.BLOCK,
            self.PLANES[2],
            self.LAYERS[2],
            dilation=dilations[2],
            norm_type=self.NORM_TYPE,
            bn_momentum=bn_momentum,
        )

        self.conv4p8s2 = conv(
            self.inplanes,
            self.inplanes,
            kernel_size=space_n_time_m(2, 1),
            stride=space_n_time_m(2, 1),
            dilation=1,
            conv_type=self.NON_BLOCK_CONV_TYPE,
            D=D,
        )
        self.bn4 = get_norm(self.NORM_TYPE, self.inplanes, D, bn_momentum=bn_momentum)
        self.block4 = self._make_layer(
            self.BLOCK,
            self.PLANES[3],
            self.LAYERS[3],
            dilation=dilations[3],
            norm_type=self.NORM_TYPE,
            bn_momentum=bn_momentum,
        )
        self.convtr4p16s2 = conv_tr(
            self.inplanes,
            self.PLANES[4],
            kernel_size=space_n_time_m(2, 1),
            upsample_stride=space_n_time_m(2, 1),
            dilation=1,
            bias=False,
            conv_type=self.NON_BLOCK_CONV_TYPE,
            D=D,
        )
        self.bntr4 = get_norm(
            self.NORM_TYPE, self.PLANES[4], D, bn_momentum=bn_momentum
        )

        self.inplanes = self.PLANES[4] + self.PLANES[2] * self.BLOCK.expansion
        self.block5 = self._make_layer(
            self.BLOCK,
            self.PLANES[4],
            self.LAYERS[4],
            dilation=dilations[4],
            norm_type=self.NORM_TYPE,
            bn_momentum=bn_momentum,
        )
        self.convtr5p8s2 = conv_tr(
            self.inplanes,
            self.PLANES[5],
            kernel_size=space_n_time_m(2, 1),
            upsample_stride=space_n_time_m(2, 1),
            dilation=1,
            bias=False,
            conv_type=self.NON_BLOCK_CONV_TYPE,
            D=D,
        )
        self.bntr5 = get_norm(
            self.NORM_TYPE, self.PLANES[5], D, bn_momentum=bn_momentum
        )

        self.inplanes = self.PLANES[5] + self.PLANES[1] * self.BLOCK.expansion
        self.block6 = self._make_layer(
            self.BLOCK,
            self.PLANES[5],
            self.LAYERS[5],
            dilation=dilations[5],
            norm_type=self.NORM_TYPE,
            bn_momentum=bn_momentum,
        )
        self.convtr6p4s2 = conv_tr(
            self.inplanes,
            self.PLANES[6],
            kernel_size=space_n_time_m(2, 1),
            upsample_stride=space_n_time_m(2, 1),
            dilation=1,
            bias=False,
            conv_type=self.NON_BLOCK_CONV_TYPE,
            D=D,
        )
        self.bntr6 = get_norm(
            self.NORM_TYPE, self.PLANES[6], D, bn_momentum=bn_momentum
        )

        self.inplanes = self.PLANES[6] + self.PLANES[0] * self.BLOCK.expansion
        self.block7 = self._make_layer(
            self.BLOCK,
            self.PLANES[6],
            self.LAYERS[6],
            dilation=dilations[6],
            norm_type=self.NORM_TYPE,
            bn_momentum=bn_momentum,
        )
        self.convtr7p2s2 = conv_tr(
            self.inplanes,
            self.PLANES[7],
            kernel_size=space_n_time_m(2, 1),
            upsample_stride=space_n_time_m(2, 1),
            dilation=1,
            bias=False,
            conv_type=self.NON_BLOCK_CONV_TYPE,
            D=D,
        )
        self.bntr7 = get_norm(
            self.NORM_TYPE, self.PLANES[7], D, bn_momentum=bn_momentum
        )

        self.inplanes = self.PLANES[7] + self.INIT_DIM
        self.block8 = self._make_layer(
            self.BLOCK,
            self.PLANES[7],
            self.LAYERS[7],
            dilation=dilations[7],
            norm_type=self.NORM_TYPE,
            bn_momentum=bn_momentum,
        )

        self.final = conv(
            self.PLANES[7], out_channels, kernel_size=1, stride=1, bias=True, D=D
        )
        self.relu = MinkowskiReLU(inplace=True)

    def forward(self, x):
        feature_maps = []

        out = self.conv0p1s1(x)
        out = self.bn0(out)
        out_p1 = self.relu(out)

        out = self.conv1p1s2(out_p1)
        out = self.bn1(out)
        out = self.relu(out)
        out_b1p2 = self.block1(out)

        out = self.conv2p2s2(out_b1p2)
        out = self.bn2(out)
        out = self.relu(out)
        out_b2p4 = self.block2(out)

        out = self.conv3p4s2(out_b2p4)
        out = self.bn3(out)
        out = self.relu(out)
        out_b3p8 = self.block3(out)

        # pixel_dist=16
        out = self.conv4p8s2(out_b3p8)
        out = self.bn4(out)
        out = self.relu(out)
        out = self.block4(out)

        feature_maps.append(out)

        # pixel_dist=8
        out = self.convtr4p16s2(out)
        out = self.bntr4(out)
        out = self.relu(out)

        out = me.cat(out, out_b3p8)
        out = self.block5(out)

        feature_maps.append(out)

        # pixel_dist=4
        out = self.convtr5p8s2(out)
        out = self.bntr5(out)
        out = self.relu(out)

        out = me.cat(out, out_b2p4)
        out = self.block6(out)

        feature_maps.append(out)

        # pixel_dist=2
        out = self.convtr6p4s2(out)
        out = self.bntr6(out)
        out = self.relu(out)

        out = me.cat(out, out_b1p2)
        out = self.block7(out)

        feature_maps.append(out)

        # pixel_dist=1
        out = self.convtr7p2s2(out)
        out = self.bntr7(out)
        out = self.relu(out)

        out = me.cat(out, out_p1)
        out = self.block8(out)

        feature_maps.append(out)

        if not self.out_fpn:
            return out
        else:
            return out, feature_maps


class Res16UNet14(Res16UNetBase):
    BLOCK = BasicBlock
    LAYERS = (1, 1, 1, 1, 1, 1, 1, 1)


class Res16UNet18(Res16UNetBase):
    BLOCK = BasicBlock
    LAYERS = (2, 2, 2, 2, 2, 2, 2, 2)


class Res16UNet34(Res16UNetBase):
    BLOCK = BasicBlock
    LAYERS = (2, 3, 4, 6, 2, 2, 2, 2)


class Res16UNet50(Res16UNetBase):
    BLOCK = Bottleneck
    LAYERS = (2, 3, 4, 6, 2, 2, 2, 2)


class Res16UNet101(Res16UNetBase):
    BLOCK = Bottleneck
    LAYERS = (2, 3, 4, 23, 2, 2, 2, 2)


class Res16UNet14A(Res16UNet14):
    PLANES = (32, 64, 128, 256, 128, 128, 96, 96)


class Res16UNet14A2(Res16UNet14A):
    LAYERS = (1, 1, 1, 1, 2, 2, 2, 2)


class Res16UNet14B(Res16UNet14):
    PLANES = (32, 64, 128, 256, 128, 128, 128, 128)


class Res16UNet14B2(Res16UNet14B):
    LAYERS = (1, 1, 1, 1, 2, 2, 2, 2)


class Res16UNet14B3(Res16UNet14B):
    LAYERS = (2, 2, 2, 2, 1, 1, 1, 1)


class Res16UNet14C(Res16UNet14):
    PLANES = (32, 64, 128, 256, 192, 192, 128, 128)


class Res16UNet14D(Res16UNet14):
    PLANES = (32, 64, 128, 256, 384, 384, 384, 384)


class Res16UNet18A(Res16UNet18):
    PLANES = (32, 64, 128, 256, 128, 128, 96, 96)


class Res16UNet18B(Res16UNet18):
    PLANES = (32, 64, 128, 256, 128, 128, 128, 128)


class Res16UNet18D(Res16UNet18):
    PLANES = (32, 64, 128, 256, 384, 384, 384, 384)


class Res16UNet34A(Res16UNet34):
    PLANES = (32, 64, 128, 256, 256, 128, 64, 64)


class Res16UNet34B(Res16UNet34):
    PLANES = (32, 64, 128, 256, 256, 128, 64, 32)


class Res16UNet34C(Res16UNet34):
    PLANES = (32, 64, 128, 256, 256, 128, 96, 96)

class Custom30M(Res16UNet34):
    PLANES = (32, 64, 128, 256, 128, 64, 64, 32)

class Res16UNet34D(Res16UNet34):
    PLANES = (32, 64, 128, 256, 256, 128, 96, 128)


class STRes16UNetBase(Res16UNetBase):

    CONV_TYPE = ConvType.SPATIAL_HYPERCUBE_TEMPORAL_HYPERCROSS

    def __init__(self, in_channels, out_channels, config, D=4, **kwargs):
        super().__init__(in_channels, out_channels, config, D, **kwargs)


class STRes16UNet14(STRes16UNetBase, Res16UNet14):
    pass


class STRes16UNet14A(STRes16UNetBase, Res16UNet14A):
    pass


class STRes16UNet18(STRes16UNetBase, Res16UNet18):
    pass


class STRes16UNet34(STRes16UNetBase, Res16UNet34):
    pass


class STRes16UNet50(STRes16UNetBase, Res16UNet50):
    pass


class STRes16UNet101(STRes16UNetBase, Res16UNet101):
    pass


class STRes16UNet18A(STRes16UNet18):
    PLANES = (32, 64, 128, 256, 128, 128, 96, 96)


class STResTesseract16UNetBase(STRes16UNetBase):
    pass
    #CONV_TYPE = ConvType.HYPERCUBE


class STResTesseract16UNet18A(STRes16UNet18A, STResTesseract16UNetBase):
    pass

models/res16unet.py
aminebdj/3D-OWIS
import torch.nn as nn
import MinkowskiEngine as ME

from models.model import Model
from models.modules.common import ConvType, NormType, conv, get_norm, sum_pool
from models.modules.resnet_block import BasicBlock, Bottleneck


class ResNetBase(Model):
    BLOCK = None
    LAYERS = ()
    INIT_DIM = 64
    PLANES = (64, 128, 256, 512)
    OUT_PIXEL_DIST = 32
    HAS_LAST_BLOCK = False
    CONV_TYPE = ConvType.HYPERCUBE

    def __init__(self, in_channels, out_channels, config, D=3, **kwargs):
        assert self.BLOCK is not None
        assert self.OUT_PIXEL_DIST > 0

        super().__init__(in_channels, out_channels, config, D, **kwargs)

        self.network_initialization(in_channels, out_channels, config, D)
        self.weight_initialization()

    def network_initialization(self, in_channels, out_channels, config, D):
        def space_n_time_m(n, m):
            return n if D == 3 else [n, n, n, m]

        if D == 4:
            self.OUT_PIXEL_DIST = space_n_time_m(self.OUT_PIXEL_DIST, 1)

        dilations = config.dilations
        bn_momentum = config.bn_momentum
        self.inplanes = self.INIT_DIM
        self.conv1 = conv(
            in_channels,
            self.inplanes,
            kernel_size=space_n_time_m(config.conv1_kernel_size, 1),
            stride=1,
            D=D,
        )

        self.bn1 = get_norm(
            NormType.BATCH_NORM, self.inplanes, D=self.D, bn_momentum=bn_momentum
        )
        self.relu = ME.MinkowskiReLU(inplace=True)
        self.pool = sum_pool(
            kernel_size=space_n_time_m(2, 1), stride=space_n_time_m(2, 1), D=D
        )

        self.layer1 = self._make_layer(
            self.BLOCK,
            self.PLANES[0],
            self.LAYERS[0],
            stride=space_n_time_m(2, 1),
            dilation=space_n_time_m(dilations[0], 1),
        )
        self.layer2 = self._make_layer(
            self.BLOCK,
            self.PLANES[1],
            self.LAYERS[1],
            stride=space_n_time_m(2, 1),
            dilation=space_n_time_m(dilations[1], 1),
        )
        self.layer3 = self._make_layer(
            self.BLOCK,
            self.PLANES[2],
            self.LAYERS[2],
            stride=space_n_time_m(2, 1),
            dilation=space_n_time_m(dilations[2], 1),
        )
        self.layer4 = self._make_layer(
            self.BLOCK,
            self.PLANES[3],
            self.LAYERS[3],
            stride=space_n_time_m(2, 1),
            dilation=space_n_time_m(dilations[3], 1),
        )

        self.final = conv(
            self.PLANES[3] * self.BLOCK.expansion,
            out_channels,
            kernel_size=1,
            bias=True,
            D=D,
        )

    def weight_initialization(self):
        for m in self.modules():
            if isinstance(m, ME.MinkowskiBatchNorm):
                nn.init.constant_(m.bn.weight, 1)
                nn.init.constant_(m.bn.bias, 0)

    def _make_layer(
        self,
        block,
        planes,
        blocks,
        stride=1,
        dilation=1,
        norm_type=NormType.BATCH_NORM,
        bn_momentum=0.1,
    ):
        downsample = None
        if stride != 1 or self.inplanes != planes * block.expansion:
            downsample = nn.Sequential(
                conv(
                    self.inplanes,
                    planes * block.expansion,
                    kernel_size=1,
                    stride=stride,
                    bias=False,
                    D=self.D,
                ),
                get_norm(
                    norm_type,
                    planes * block.expansion,
                    D=self.D,
                    bn_momentum=bn_momentum,
                ),
            )
        layers = []
        layers.append(
            block(
                self.inplanes,
                planes,
                stride=stride,
                dilation=dilation,
                downsample=downsample,
                conv_type=self.CONV_TYPE,
                D=self.D,
            )
        )
        self.inplanes = planes * block.expansion
        for i in range(1, blocks):
            layers.append(
                block(
                    self.inplanes,
                    planes,
                    stride=1,
                    dilation=dilation,
                    conv_type=self.CONV_TYPE,
                    D=self.D,
                )
            )

        return nn.Sequential(*layers)

    def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)
        x = self.pool(x)

        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)

        x = self.final(x)
        return x


class ResNet14(ResNetBase):
    BLOCK = BasicBlock
    LAYERS = (1, 1, 1, 1)


class ResNet18(ResNetBase):
    BLOCK = BasicBlock
    LAYERS = (2, 2, 2, 2)


class ResNet34(ResNetBase):
    BLOCK = BasicBlock
    LAYERS = (3, 4, 6, 3)


class ResNet50(ResNetBase):
    BLOCK = Bottleneck
    LAYERS = (3, 4, 6, 3)


class ResNet101(ResNetBase):
    BLOCK = Bottleneck
    LAYERS = (3, 4, 23, 3)


class STResNetBase(ResNetBase):

    CONV_TYPE = ConvType.SPATIAL_HYPERCUBE_TEMPORAL_HYPERCROSS

    def __init__(self, in_channels, out_channels, config, D=4, **kwargs):
        super().__init__(in_channels, out_channels, config, D, **kwargs)


class STResNet14(STResNetBase, ResNet14):
    pass


class STResNet18(STResNetBase, ResNet18):
    pass


class STResNet34(STResNetBase, ResNet34):
    pass


class STResNet50(STResNetBase, ResNet50):
    pass


class STResNet101(STResNetBase, ResNet101):
    pass


class STResTesseractNetBase(STResNetBase):
    CONV_TYPE = ConvType.HYPERCUBE


class STResTesseractNet14(STResTesseractNetBase, STResNet14):
    pass


class STResTesseractNet18(STResTesseractNetBase, STResNet18):
    pass


class STResTesseractNet34(STResTesseractNetBase, STResNet34):
    pass


class STResTesseractNet50(STResTesseractNetBase, STResNet50):
    pass


class STResTesseractNet101(STResTesseractNetBase, STResNet101):
    pass

models/resnet.py
aminebdj/3D-OWIS
import torch.nn as nn
import MinkowskiEngine as ME
import MinkowskiEngine.MinkowskiOps as me
from MinkowskiEngine import MinkowskiReLU

from models.resnet import ResNetBase, get_norm
from models.modules.common import ConvType, NormType, conv, conv_tr
from models.modules.resnet_block import BasicBlock, Bottleneck, BasicBlockINBN


class MinkUNetBase(ResNetBase):
    BLOCK = None
    PLANES = (64, 128, 256, 512, 256, 128, 128)
    DILATIONS = (1, 1, 1, 1, 1, 1)
    LAYERS = (2, 2, 2, 2, 2, 2)
    INIT_DIM = 64
    OUT_PIXEL_DIST = 1
    NORM_TYPE = NormType.BATCH_NORM
    NON_BLOCK_CONV_TYPE = ConvType.SPATIAL_HYPERCUBE
    CONV_TYPE = ConvType.SPATIAL_HYPERCUBE_TEMPORAL_HYPERCROSS

    # To use the model, must call initialize_coords before forward pass.
    # Once data is processed, call clear to reset the model before calling initialize_coords
    def __init__(self, in_channels, out_channels, config, D=3, **kwargs):
        super().__init__(in_channels, out_channels, config, D)

    def network_initialization(self, in_channels, out_channels, config, D):
        # Setup net_metadata
        dilations = self.DILATIONS
        bn_momentum = config.bn_momentum

        def space_n_time_m(n, m):
            return n if D == 3 else [n, n, n, m]

        if D == 4:
            self.OUT_PIXEL_DIST = space_n_time_m(self.OUT_PIXEL_DIST, 1)

        # Output of the first conv concated to conv6
        self.inplanes = self.INIT_DIM
        self.conv1p1s1 = conv(
            in_channels,
            self.inplanes,
            kernel_size=space_n_time_m(config.conv1_kernel_size, 1),
            stride=1,
            dilation=1,
            conv_type=self.NON_BLOCK_CONV_TYPE,
            D=D,
        )

        self.bn1 = get_norm(self.NORM_TYPE, self.PLANES[0], D, bn_momentum=bn_momentum)
        self.block1 = self._make_layer(
            self.BLOCK,
            self.PLANES[0],
            self.LAYERS[0],
            dilation=dilations[0],
            norm_type=self.NORM_TYPE,
            bn_momentum=bn_momentum,
        )

        self.conv2p1s2 = conv(
            self.inplanes,
            self.inplanes,
            kernel_size=space_n_time_m(2, 1),
            stride=space_n_time_m(2, 1),
            dilation=1,
            conv_type=self.NON_BLOCK_CONV_TYPE,
            D=D,
        )
        self.bn2 = get_norm(self.NORM_TYPE, self.inplanes, D, bn_momentum=bn_momentum)
        self.block2 = self._make_layer(
            self.BLOCK,
            self.PLANES[1],
            self.LAYERS[1],
            dilation=dilations[1],
            norm_type=self.NORM_TYPE,
            bn_momentum=bn_momentum,
        )

        self.conv3p2s2 = conv(
            self.inplanes,
            self.inplanes,
            kernel_size=space_n_time_m(2, 1),
            stride=space_n_time_m(2, 1),
            dilation=1,
            conv_type=self.NON_BLOCK_CONV_TYPE,
            D=D,
        )
        self.bn3 = get_norm(self.NORM_TYPE, self.inplanes, D, bn_momentum=bn_momentum)
        self.block3 = self._make_layer(
            self.BLOCK,
            self.PLANES[2],
            self.LAYERS[2],
            dilation=dilations[2],
            norm_type=self.NORM_TYPE,
            bn_momentum=bn_momentum,
        )

        self.conv4p4s2 = conv(
            self.inplanes,
            self.inplanes,
            kernel_size=space_n_time_m(2, 1),
            stride=space_n_time_m(2, 1),
            dilation=1,
            conv_type=self.NON_BLOCK_CONV_TYPE,
            D=D,
        )
        self.bn4 = get_norm(self.NORM_TYPE, self.inplanes, D, bn_momentum=bn_momentum)
        self.block4 = self._make_layer(
            self.BLOCK,
            self.PLANES[3],
            self.LAYERS[3],
            dilation=dilations[3],
            norm_type=self.NORM_TYPE,
            bn_momentum=bn_momentum,
        )
        self.convtr4p8s2 = conv_tr(
            self.inplanes,
            self.PLANES[4],
            kernel_size=space_n_time_m(2, 1),
            upsample_stride=space_n_time_m(2, 1),
            dilation=1,
            bias=False,
            conv_type=self.NON_BLOCK_CONV_TYPE,
            D=D,
        )
        self.bntr4 = get_norm(
            self.NORM_TYPE, self.PLANES[4], D, bn_momentum=bn_momentum
        )

        self.inplanes = self.PLANES[4] + self.PLANES[2] * self.BLOCK.expansion
        self.block5 = self._make_layer(
            self.BLOCK,
            self.PLANES[4],
            self.LAYERS[4],
            dilation=dilations[4],
            norm_type=self.NORM_TYPE,
            bn_momentum=bn_momentum,
        )
        self.convtr5p4s2 = conv_tr(
            self.inplanes,
            self.PLANES[5],
            kernel_size=space_n_time_m(2, 1),
            upsample_stride=space_n_time_m(2, 1),
            dilation=1,
            bias=False,
            conv_type=self.NON_BLOCK_CONV_TYPE,
            D=D,
        )
        self.bntr5 = get_norm(
            self.NORM_TYPE, self.PLANES[5], D, bn_momentum=bn_momentum
        )

        self.inplanes = self.PLANES[5] + self.PLANES[1] * self.BLOCK.expansion
        self.block6 = self._make_layer(
            self.BLOCK,
            self.PLANES[5],
            self.LAYERS[5],
            dilation=dilations[5],
            norm_type=self.NORM_TYPE,
            bn_momentum=bn_momentum,
        )
        self.convtr6p2s2 = conv_tr(
            self.inplanes,
            self.PLANES[6],
            kernel_size=space_n_time_m(2, 1),
            upsample_stride=space_n_time_m(2, 1),
            dilation=1,
            bias=False,
            conv_type=self.NON_BLOCK_CONV_TYPE,
            D=D,
        )
        self.bntr6 = get_norm(
            self.NORM_TYPE, self.PLANES[6], D, bn_momentum=bn_momentum
        )
        self.relu = MinkowskiReLU(inplace=True)

        self.final = nn.Sequential(
            conv(
                self.PLANES[6] + self.PLANES[0] * self.BLOCK.expansion,
                512,
                kernel_size=1,
                stride=1,
                dilation=1,
                bias=False,
                D=D,
            ),
            ME.MinkowskiBatchNorm(512),
            ME.MinkowskiReLU(),
            conv(
                512, out_channels, kernel_size=1, stride=1, dilation=1, bias=True, D=D
            ),
        )

    def forward(self, x):
        out = self.conv1p1s1(x)
        out = self.bn1(out)
        out = self.relu(out)

        out_b1p1 = self.block1(out)

        out = self.conv2p1s2(out_b1p1)
        out = self.bn2(out)
        out = self.relu(out)

        out_b2p2 = self.block2(out)

        out = self.conv3p2s2(out_b2p2)
        out = self.bn3(out)
        out = self.relu(out)

        out_b3p4 = self.block3(out)

        out = self.conv4p4s2(out_b3p4)
        out = self.bn4(out)
        out = self.relu(out)

        # pixel_dist=8
        out = self.block4(out)

        out = self.convtr4p8s2(out)
        out = self.bntr4(out)
        out = self.relu(out)

        out = me.cat(out, out_b3p4)
        out = self.block5(out)

        out = self.convtr5p4s2(out)
        out = self.bntr5(out)
        out = self.relu(out)

        out = me.cat(out, out_b2p2)
        out = self.block6(out)

        out = self.convtr6p2s2(out)
        out = self.bntr6(out)
        out = self.relu(out)

        out = me.cat(out, out_b1p1)
        return self.final(out)


class ResUNet14(MinkUNetBase):
    BLOCK = BasicBlock
    LAYERS = (1, 1, 1, 1, 1, 1)


class ResUNet18(MinkUNetBase):
    BLOCK = BasicBlock
    LAYERS = (2, 2, 2, 2, 2, 2)


class ResUNet18INBN(ResUNet18):
    NORM_TYPE = NormType.INSTANCE_BATCH_NORM
    BLOCK = BasicBlockINBN


class ResUNet34(MinkUNetBase):
    BLOCK = BasicBlock
    LAYERS = (3, 4, 6, 3, 2, 2)


class ResUNet50(MinkUNetBase):
    BLOCK = Bottleneck
    LAYERS = (3, 4, 6, 3, 2, 2)


class ResUNet101(MinkUNetBase):
    BLOCK = Bottleneck
    LAYERS = (3, 4, 23, 3, 2, 2)


class ResUNet14D(ResUNet14):
    PLANES = (64, 128, 256, 512, 512, 512, 512)


class ResUNet18D(ResUNet18):
    PLANES = (64, 128, 256, 512, 512, 512, 512)


class ResUNet34D(ResUNet34):
    PLANES = (64, 128, 256, 512, 512, 512, 512)


class ResUNet34E(ResUNet34):
    INIT_DIM = 32
    PLANES = (32, 64, 128, 256, 128, 64, 64)


class ResUNet34F(ResUNet34):
    INIT_DIM = 32
    PLANES = (32, 64, 128, 256, 128, 64, 32)


class MinkUNetHyper(MinkUNetBase):
    BLOCK = None
    PLANES = (64, 128, 256, 512, 256, 128, 128)
    DILATIONS = (1, 1, 1, 1, 1, 1)
    LAYERS = (2, 2, 2, 2, 2, 2)
    INIT_DIM = 64
    OUT_PIXEL_DIST = 1
    NORM_TYPE = NormType.BATCH_NORM
    NON_BLOCK_CONV_TYPE = ConvType.SPATIAL_HYPERCUBE
    CONV_TYPE = ConvType.SPATIAL_HYPERCUBE_TEMPORAL_HYPERCROSS

    # To use the model, must call initialize_coords before forward pass.
    # Once data is processed, call clear to reset the model before calling initialize_coords
    def __init__(self, in_channels, out_channels, config, D=3, **kwargs):
        super(MinkUNetBase, self).__init__(in_channels, out_channels, config, D)

    def network_initialization(self, in_channels, out_channels, config, D):
        # Setup net_metadata
        dilations = self.DILATIONS
        bn_momentum = config.bn_momentum

        def space_n_time_m(n, m):
            return n if D == 3 else [n, n, n, m]

        if D == 4:
            self.OUT_PIXEL_DIST = space_n_time_m(self.OUT_PIXEL_DIST, 1)

        # Output of the first conv concated to conv6
        self.inplanes = self.INIT_DIM
        self.conv1p1s1 = conv(
            in_channels,
            self.inplanes,
            kernel_size=space_n_time_m(config.conv1_kernel_size, 1),
            stride=1,
            dilation=1,
            conv_type=self.NON_BLOCK_CONV_TYPE,
            D=D,
        )

        self.bn1 = get_norm(self.NORM_TYPE, self.PLANES[0], D, bn_momentum=bn_momentum)
        self.block1 = self._make_layer(
            self.BLOCK,
            self.PLANES[0],
            self.LAYERS[0],
            dilation=dilations[0],
            norm_type=self.NORM_TYPE,
            bn_momentum=bn_momentum,
        )

        self.conv2p1s2 = conv(
            self.inplanes,
            self.inplanes,
            kernel_size=space_n_time_m(2, 1),
            stride=space_n_time_m(2, 1),
            dilation=1,
            conv_type=self.NON_BLOCK_CONV_TYPE,
            D=D,
        )
        self.bn2 = get_norm(self.NORM_TYPE, self.inplanes, D, bn_momentum=bn_momentum)
        self.block2 = self._make_layer(
            self.BLOCK,
            self.PLANES[1],
            self.LAYERS[1],
            dilation=dilations[1],
            norm_type=self.NORM_TYPE,
            bn_momentum=bn_momentum,
        )

        self.conv3p2s2 = conv(
            self.inplanes,
            self.inplanes,
            kernel_size=space_n_time_m(2, 1),
            stride=space_n_time_m(2, 1),
            dilation=1,
            conv_type=self.NON_BLOCK_CONV_TYPE,
            D=D,
        )
        self.bn3 = get_norm(self.NORM_TYPE, self.inplanes, D, bn_momentum=bn_momentum)
        self.block3 = self._make_layer(
            self.BLOCK,
            self.PLANES[2],
            self.LAYERS[2],
            dilation=dilations[2],
            norm_type=self.NORM_TYPE,
            bn_momentum=bn_momentum,
        )

        self.conv4p4s2 = conv(
            self.inplanes,
            self.inplanes,
            kernel_size=space_n_time_m(2, 1),
            stride=space_n_time_m(2, 1),
            dilation=1,
            conv_type=self.NON_BLOCK_CONV_TYPE,
            D=D,
        )
        self.bn4 = get_norm(self.NORM_TYPE, self.inplanes, D, bn_momentum=bn_momentum)
        self.block4 = self._make_layer(
            self.BLOCK,
            self.PLANES[3],
            self.LAYERS[3],
            dilation=dilations[3],
            norm_type=self.NORM_TYPE,
            bn_momentum=bn_momentum,
        )
        self.pool_tr4 = ME.MinkowskiPoolingTranspose(
            kernel_size=8, stride=8, dimension=D
        )
        _ = self.inplanes
        self.convtr4p8s2 = conv_tr(
            self.inplanes,
            self.PLANES[4],
            kernel_size=space_n_time_m(2, 1),
            upsample_stride=space_n_time_m(2, 1),
            dilation=1,
            bias=False,
            conv_type=self.NON_BLOCK_CONV_TYPE,
            D=D,
        )
        self.bntr4 = get_norm(
            self.NORM_TYPE, self.PLANES[4], D, bn_momentum=bn_momentum
        )

        self.inplanes = self.PLANES[4] + self.PLANES[2] * self.BLOCK.expansion
        self.block5 = self._make_layer(
            self.BLOCK,
            self.PLANES[4],
            self.LAYERS[4],
            dilation=dilations[4],
            norm_type=self.NORM_TYPE,
            bn_momentum=bn_momentum,
        )
        self.pool_tr5 = ME.MinkowskiPoolingTranspose(
            kernel_size=4, stride=4, dimension=D
        )
        out_pool5 = self.inplanes
        self.convtr5p4s2 = conv_tr(
            self.inplanes,
            self.PLANES[5],
            kernel_size=space_n_time_m(2, 1),
            upsample_stride=space_n_time_m(2, 1),
            dilation=1,
            bias=False,
            conv_type=self.NON_BLOCK_CONV_TYPE,
            D=D,
        )
        self.bntr5 = get_norm(
            self.NORM_TYPE, self.PLANES[5], D, bn_momentum=bn_momentum
        )

        self.inplanes = self.PLANES[5] + self.PLANES[1] * self.BLOCK.expansion
        self.block6 = self._make_layer(
            self.BLOCK,
            self.PLANES[5],
            self.LAYERS[5],
            dilation=dilations[5],
            norm_type=self.NORM_TYPE,
            bn_momentum=bn_momentum,
        )
        self.pool_tr6 = ME.MinkowskiPoolingTranspose(
            kernel_size=2, stride=2, dimension=D
        )
        out_pool6 = self.inplanes
        self.convtr6p2s2 = conv_tr(
            self.inplanes,
            self.PLANES[6],
            kernel_size=space_n_time_m(2, 1),
            upsample_stride=space_n_time_m(2, 1),
            dilation=1,
            bias=False,
            conv_type=self.NON_BLOCK_CONV_TYPE,
            D=D,
        )
        self.bntr6 = get_norm(
            self.NORM_TYPE, self.PLANES[6], D, bn_momentum=bn_momentum
        )

        self.relu = MinkowskiReLU(inplace=True)

        self.final = nn.Sequential(
            conv(
                out_pool5
                + out_pool6
                + self.PLANES[6]
                + self.PLANES[0] * self.BLOCK.expansion,
                512,
                kernel_size=1,
                bias=False,
                D=D,
            ),
            ME.MinkowskiBatchNorm(512),
            ME.MinkowskiReLU(),
            conv(512, out_channels, kernel_size=1, bias=True, D=D),
        )

    def forward(self, x):
        out = self.conv1p1s1(x)
        out = self.bn1(out)
        out = self.relu(out)

        out_b1p1 = self.block1(out)

        out = self.conv2p1s2(out_b1p1)
        out = self.bn2(out)
        out = self.relu(out)

        out_b2p2 = self.block2(out)

        out = self.conv3p2s2(out_b2p2)
        out = self.bn3(out)
        out = self.relu(out)

        out_b3p4 = self.block3(out)

        out = self.conv4p4s2(out_b3p4)
        out = self.bn4(out)
        out = self.relu(out)

        # pixel_dist=8
        out = self.block4(out)

        out = self.convtr4p8s2(out)
        out = self.bntr4(out)
        out = self.relu(out)

        out = me.cat(out, out_b3p4)
        out = self.block5(out)
        out_5 = self.pool_tr5(out)

        out = self.convtr5p4s2(out)
        out = self.bntr5(out)
        out = self.relu(out)

        out = me.cat(out, out_b2p2)
        out = self.block6(out)
        out_6 = self.pool_tr6(out)

        out = self.convtr6p2s2(out)
        out = self.bntr6(out)
        out = self.relu(out)

        out = me.cat(out, out_b1p1, out_6, out_5)
        return self.final(out)


class MinkUNetHyper14INBN(MinkUNetHyper):
    NORM_TYPE = NormType.INSTANCE_BATCH_NORM
    BLOCK = BasicBlockINBN


class STMinkUNetBase(MinkUNetBase):

    CONV_TYPE = ConvType.SPATIAL_HYPERCUBE_TEMPORAL_HYPERCROSS

    def __init__(self, in_channels, out_channels, config, D=4, **kwargs):
        super().__init__(in_channels, out_channels, config, D, **kwargs)


class STResUNet14(STMinkUNetBase, ResUNet14):
    pass


class STResUNet18(STMinkUNetBase, ResUNet18):
    pass


class STResUNet34(STMinkUNetBase, ResUNet34):
    pass


class STResUNet50(STMinkUNetBase, ResUNet50):
    pass


class STResUNet101(STMinkUNetBase, ResUNet101):
    pass


class STResTesseractUNetBase(STMinkUNetBase):
    CONV_TYPE = ConvType.HYPERCUBE


class STResTesseractUNet14(STResTesseractUNetBase, ResUNet14):
    pass


class STResTesseractUNet18(STResTesseractUNetBase, ResUNet18):
    pass


class STResTesseractUNet34(STResTesseractUNetBase, ResUNet34):
    pass


class STResTesseractUNet50(STResTesseractUNetBase, ResUNet50):
    pass


class STResTesseractUNet101(STResTesseractUNetBase, ResUNet101):
    pass

models/resunet.py
aminebdj/3D-OWIS
# Evaluates semantic instance task
# Adapted from the CityScapes evaluation: https://github.com/mcordts/cityscapesScripts/tree/master/cityscapesscripts/evaluation
# Input:
#   - path to .txt prediction files
#   - path to .txt ground truth files
#   - output file to write results to
# Each .txt prediction file look like:
#    [(pred0) rel. path to pred. mask over verts as .txt] [(pred0) label id] [(pred0) confidence]
#    [(pred1) rel. path to pred. mask over verts as .txt] [(pred1) label id] [(pred1) confidence]
#    [(pred2) rel. path to pred. mask over verts as .txt] [(pred2) label id] [(pred2) confidence]
#    ...
#
# NOTE: The prediction files must live in the root of the given prediction path.
#       Predicted mask .txt files must live in a subfolder.
#       Additionally, filenames must not contain spaces.
# The relative paths to predicted masks must contain one integer per line,
# where each line corresponds to vertices in the *_vh_clean_2.ply (in that order).
# Non-zero integers indicate part of the predicted instance.
# The label ids specify the class of the corresponding mask.
# Confidence is a float confidence score of the mask.
#
# Note that only the valid classes are used for evaluation,
# i.e., any ground truth label not in the valid label set
# is ignored in the evaluation.
#
# example usage: evaluate_semantic_instance.py --scan_path [path to scan data] --output_file [output file]

# python imports

import math
import os, sys, argparse
import inspect
from copy import deepcopy
from uuid import uuid4
import yaml
import torch
import json
try:
    import numpy as np
except:
    print("Failed to import numpy package.")
    sys.exit(-1)

from scipy import stats
import csv
import util
import util_3d


from owis_splits import KNOWN_CLASSES_LABELS, PREV_KNOWN_CLASSES_LABELS, VALID_CLASS_IDS_200, CLASS_LABELS_200


# ---------- Label info ---------- #
ID_TO_LABEL_200 = {}
LABEL_TO_ID_200 = {}
for i in range(len(VALID_CLASS_IDS_200)):
    LABEL_TO_ID_200[CLASS_LABELS_200[i]] = VALID_CLASS_IDS_200[i]
    ID_TO_LABEL_200[VALID_CLASS_IDS_200[i]] = CLASS_LABELS_200[i]
# ---------- Evaluation params ---------- #


# overlaps for evaluation
opt = {}
opt['overlaps']             = np.append(np.arange(0.5,0.95,0.05), 0.25)
# minimum region size for evaluation [verts]
opt['min_region_sizes']     = np.array( [ 100 ] ) # 100 for s3dis, scannet
# distance thresholds [m]
opt['distance_threshes']    = np.array( [  float('inf') ] )
# distance confidences
opt['distance_confs']       = np.array( [ -float('inf') ] )


parser = argparse.ArgumentParser()
parser.add_argument('--pred_path', required=True, help='path to directory of predicted .txt files')
parser.add_argument('--gt_path', required=True, help='path to directory of gt .txt files')
parser.add_argument('--output_file', default='', help='output file [default: pred_path/semantic_instance_evaluation.txt]')
opt = parser.parse_args()

if opt.output_file == '':
    opt.output_file = os.path.join(opt.pred_path, 'semantic_instance_evaluation.txt')



def evaluate_matches(matches):
    overlaps = opt['overlaps']
    min_region_sizes = [ opt['min_region_sizes'][0] ]
    dist_threshes = [ opt['distance_threshes'][0] ]
    dist_confs = [ opt['distance_confs'][0] ]
    # results: class x overlap
    ap = np.zeros( (len(dist_threshes) , len(CLASS_LABELS) , len(overlaps)) , float )
    U_Recall = {}
    WI = dict(zip([round(ov,2) for ov in overlaps],len(overlaps)*[0.0]))
    A_OSE = dict(zip([round(ov,2) for ov in overlaps],len(overlaps)*[0.0]))
    FPo = dict(zip([round(ov,2) for ov in overlaps],len(overlaps)*[0.0]))
    FPc = dict(zip([round(ov,2) for ov in overlaps],len(overlaps)*[0.0]))
    TPc = dict(zip([round(ov,2) for ov in overlaps],len(overlaps)*[0.0]))
    
    for di, (min_region_size, distance_thresh, distance_conf) in enumerate(zip(min_region_sizes, dist_threshes, dist_confs)):
        for oi, overlap_th in enumerate(overlaps):
            pred_visited = {}
            for m in matches:
                for p in matches[m]['pred']:
                    for label_name in CLASS_LABELS:
                        for p in matches[m]['pred'][label_name]:
                            if 'uuid' in p:
                                pred_visited[p['uuid']] = False
            
            
            for li, label_name in enumerate(CLASS_LABELS):
                y_true = np.empty(0)
                y_score = np.empty(0)
                
                    
                hard_false_negatives = 0
                has_gt = False
                has_pred = False
                dic = {}
                
                for m in matches:
                    pred_instances = matches[m]['pred'][label_name]
                    gt_instances = matches[m]['gt'][label_name]
                    # filter groups in ground truth
                    gt_instances = [ gt for gt in gt_instances if gt['instance_id']>=1000 and gt['vert_count']>=min_region_size and gt['med_dist']<=distance_thresh and gt['dist_conf']>=distance_conf ]
                    if gt_instances:
                        has_gt = True
                    if pred_instances:
                        has_pred = True

                    cur_true  = np.ones ( len(gt_instances) )
                    cur_score = np.ones ( len(gt_instances) ) * (-float("inf"))
                    cur_match = np.zeros( len(gt_instances) , dtype=bool )
  
                    # collect matches
                    for (gti,gt) in enumerate(gt_instances):
                        found_match = False
                        num_pred = len(gt['matched_pred'])
                        for pred in gt['matched_pred']:
                            # greedy assignments
                            if pred_visited[pred['uuid']]:
                                continue
                            overlap = float(pred['intersection']) / (gt['vert_count']+pred['vert_count']-pred['intersection'])
                            if overlap > overlap_th:
                                confidence = pred['confidence']
                                # if already have a prediction for this gt,
                                # the prediction with the lower score is automatically a false positive
                                if cur_match[gti]:
                                    max_score = max( cur_score[gti] , confidence )
                                    min_score = min( cur_score[gti] , confidence )
                                    cur_score[gti] = max_score
                                    # append false positive with unknown
                                    cur_true  = np.append(cur_true,0)
                                    cur_score = np.append(cur_score,min_score)
                                    cur_match = np.append(cur_match,True)
                                    if pred['pred_id'] != 3000 and gt['label_id'] == 3000:
                                        A_OSE[round(overlap_th,2)] += 1
                                    if pred['pred_id'] != 3000 and gt['label_id'] == 3000 and pred['confidence']>0.5:
                                        FPo[round(overlap_th,2)] += 1
                                    if pred['pred_id'] != 3000 and gt['label_id'] != 3000 and gt['label_id'] != pred['pred_id'] and pred['confidence']>0.5:
                                        FPc[round(overlap_th,2)] += 1
                                    if pred['pred_id'] == gt['label_id'] and gt['label_id'] != 3000 and pred['confidence']>0.5:
                                        TPc[round(overlap_th,2)] += 1
  
                                # otherwise set score
                                else:
                                    found_match = True
                                    cur_match[gti] = True
                                    cur_score[gti] = confidence
                                    pred_visited[pred['uuid']] = True


                        if not found_match:
                            hard_false_negatives += 1
                    # remove non-matched ground truth instances
                    
                    cur_true  = cur_true [ cur_match==True ]
                    cur_score = cur_score[ cur_match==True ]
                    
                    
                    # collect non-matched predictions as false positive
                    for pred in pred_instances:
                        found_gt = False
                        for gt in pred['matched_gt']:
                            overlap = float(gt['intersection']) / (gt['vert_count']+pred['vert_count']-gt['intersection'])
                            if overlap > overlap_th:
                                found_gt = True
                                break
                        if not found_gt:
                            num_ignore = pred['void_intersection']
                            for gt in pred['matched_gt']:
                                # group?
                                if gt['instance_id'] < 1000:
                                    num_ignore += gt['intersection']
                                # small ground truth instances
                                if gt['vert_count'] < min_region_size or gt['med_dist']>distance_thresh or gt['dist_conf']<distance_conf:
                                    num_ignore += gt['intersection']
                            proportion_ignore = float(num_ignore)/pred['vert_count']
                            # if not ignored append false positive
                            if proportion_ignore <= overlap_th:
                                cur_true = np.append(cur_true,0)
                                confidence = pred["confidence"]
                                cur_score = np.append(cur_score,confidence)
                                


                    # append to overall results
                    y_true  = np.append(y_true,cur_true)
                    y_score = np.append(y_score,cur_score)


                # compute average precision
                if has_gt and has_pred:
                    # compute precision recall curve first

                    # sorting and cumsum
                    score_arg_sort      = np.argsort(y_score)
                    y_score_sorted      = y_score[score_arg_sort]
                    y_true_sorted       = y_true[score_arg_sort]
                    y_true_sorted_cumsum = np.cumsum(y_true_sorted)

                    # unique thresholds
                    (thresholds,unique_indices) = np.unique( y_score_sorted , return_index=True )
                    num_prec_recall = len(unique_indices) + 1
                    
                    

                    # prepare precision recall
                    num_examples      = len(y_score_sorted)
                    
                    
                    # https://github.com/ScanNet/ScanNet/pull/26
                    # all predictions are non-matched but also all of them are ignored and not counted as FP
                    # y_true_sorted_cumsum is empty
                    # num_true_examples = y_true_sorted_cumsum[-1]
                    num_true_examples = y_true_sorted_cumsum[-1] if len(y_true_sorted_cumsum) > 0 else 0
                    precision         = np.zeros(num_prec_recall)
                    recall            = np.zeros(num_prec_recall) 

                    # deal with the first point
                    y_true_sorted_cumsum = np.append( y_true_sorted_cumsum , 0 )
                    # deal with remaining
                    for idx_res,idx_scores in enumerate(unique_indices):
                        cumsum = y_true_sorted_cumsum[idx_scores-1]
                        tp = num_true_examples - cumsum
                        fp = num_examples      - idx_scores - tp
                        fn = cumsum + hard_false_negatives
                        p  = float(tp)/(tp+fp)
                        r  = float(tp)/(tp+fn)
                        precision[idx_res] = p
                        recall   [idx_res] = r
                    
                    #Getting the Urecall
                    if label_name == "unknown":
                        U_Recall[int(overlap_th*100)] = recall[0]            
                    # first point in curve is artificial
                    precision[-1] = 1.
                    recall   [-1] = 0.
                    
                        
                    # compute average of precision-recall curve
                    recall_for_conv = np.copy(recall)
                    recall_for_conv = np.append(recall_for_conv[0], recall_for_conv)
                    recall_for_conv = np.append(recall_for_conv, 0.)
                    
                    
                    stepWidths = np.convolve(recall_for_conv,[-0.5,0,0.5],'valid')
                    # integrate is now simply a dot product
                    ap_current = np.dot(precision, stepWidths)

                elif has_gt:
                    ap_current = 0.0
                else:
                    ap_current = float('nan')
                ap[di,li,oi] = ap_current
            
            try:
                WI[round(overlap_th,2)] = FPo[round(overlap_th,2)]/(TPc[round(overlap_th,2)]+FPc[round(overlap_th,2)])
            except:
                pass
            
    
    return U_Recall,WI,A_OSE, ap

def compute_averages(aps):
    d_inf = 0
    o50   = np.where(np.isclose(opt['overlaps'],0.5))
    o25   = np.where(np.isclose(opt['overlaps'],0.25))
    oAllBut25  = np.where(np.logical_not(np.isclose(opt['overlaps'],0.25)))
    avg_dict = {}
    #avg_dict['all_ap']     = np.nanmean(aps[ d_inf,:,:  ])
    avg_dict['all_ap']     = np.nanmean(aps[ d_inf,:,oAllBut25])
    avg_dict['all_ap_50%'] = np.nanmean(aps[ d_inf,:,o50])
    avg_dict['all_ap_25%'] = np.nanmean(aps[ d_inf,:,o25])
    avg_dict["classes"]  = {}
    for (li,label_name) in enumerate(CLASS_LABELS):
        avg_dict["classes"][label_name]             = {}
        #avg_dict["classes"][label_name]["ap"]       = np.average(aps[ d_inf,li,  :])
        avg_dict["classes"][label_name]["ap"]       = np.average(aps[ d_inf,li,oAllBut25])
        avg_dict["classes"][label_name]["ap50%"]    = np.average(aps[ d_inf,li,o50])
        avg_dict["classes"][label_name]["ap25%"]    = np.average(aps[ d_inf,li,o25])
    return avg_dict

def make_pred_info(pred: dict):
    # pred = {'pred_scores' = 100, 'pred_classes' = 100 'pred_masks' = Nx100}
    pred_info = {}
    assert(pred['pred_classes'].shape[0] == pred['pred_scores'].shape[0] == pred['pred_masks'].shape[1])
    for i in range(len(pred['pred_classes'])):
        info = {}
        info["label_id"] = pred['pred_classes'][i]
        info["conf"] = pred['pred_scores'][i]#<=================================
        info["mask"] = pred['pred_masks'][:,i]
        pred_info[uuid4()] = info # we later need to identify these objects
    return pred_info

def assign_instances_for_scan(split,task, pred: dict, gt_file: str):
    pred_info = make_pred_info(pred)
    try:
        gt_ids = util_3d.load_ids(gt_file)
    except Exception as e:
        util.print_error('unable to load ' + gt_file + ': ' + str(e))

    # get gt instances
    gt_instances = util_3d.get_instances(split,task, gt_ids, VALID_CLASS_IDS, CLASS_LABELS, ID_TO_LABEL)
    # associate
    gt2pred = deepcopy(gt_instances)
    for label in gt2pred: 
        for gt in gt2pred[label]:
            gt['matched_pred'] = []
    pred2gt = {}
    for label in CLASS_LABELS:
        pred2gt[label] = []
    num_pred_instances = 0
    # mask of void labels in the groundtruth
    bool_void = np.logical_not(np.in1d(gt_ids//1000, VALID_CLASS_IDS))
    # go thru all prediction masks
    for uuid in pred_info:
        label_id = int(pred_info[uuid]['label_id'])
        conf = pred_info[uuid]['conf']
        if not label_id in ID_TO_LABEL:
            continue
        label_name = ID_TO_LABEL[label_id]
        # read the mask
        pred_mask = pred_info[uuid]['mask']
        assert(len(pred_mask) == len(gt_ids))
        # convert to binary
        pred_mask = np.not_equal(pred_mask, 0)
        num = np.count_nonzero(pred_mask)
        if num < opt['min_region_sizes'][0]:
            continue  # skip if empty

        pred_instance = {}
        pred_instance['uuid'] = uuid
        pred_instance['pred_id'] = num_pred_instances
        pred_instance['label_id'] = label_id
        pred_instance['vert_count'] = num
        pred_instance['confidence'] = conf
        pred_instance['void_intersection'] = np.count_nonzero(np.logical_and(bool_void, pred_mask))

        # matched gt instances
        matched_gt = []
        # go thru all gt instances with matching label
        for (gt_num, gt_inst) in enumerate(gt2pred[label_name]):
            intersection = np.count_nonzero(np.logical_and(gt_ids == gt_inst['instance_id'], pred_mask))
            if intersection > 0:
                gt_copy = gt_inst.copy()
                pred_copy = pred_instance.copy()
                gt_copy['intersection']   = intersection
                pred_copy['intersection'] = intersection
                matched_gt.append(gt_copy)
                gt2pred[label_name][gt_num]['matched_pred'].append(pred_copy)
        pred_instance['matched_gt'] = matched_gt
        num_pred_instances += 1
        pred2gt[label_name].append(pred_instance)

    return gt2pred, pred2gt


def print_results(avgs,current_epoch,save_dir,split,task_ow, u_recall, wi, a_ose):
    prev_known_avg = []
    prev_known_avg_25 = []
    prev_known_avg_50 = []
    curr_known_avg = []
    curr_known_avg_25 = []
    curr_known_avg_50 = []
    sep     = ""
    col1    = ":"
    lineLen = 64

    print("")
    print("#"*lineLen)
    line  = ""
    line += "{:<15}".format("what"      ) + sep + col1
    line += "{:>15}".format("AP"        ) + sep
    line += "{:>15}".format("AP_50%"    ) + sep
    line += "{:>15}".format("AP_25%"    ) + sep
    print(line)
    print("#"*lineLen)

    for (li,label_name) in enumerate(CLASS_LABELS):
        ap_avg  = avgs["classes"][label_name]["ap"]
        ap_50o  = avgs["classes"][label_name]["ap50%"]
        ap_25o  = avgs["classes"][label_name]["ap25%"]
        line  = "{:<15}".format(label_name) + sep + col1
        line += sep + "{:>15.3f}".format(ap_avg ) + sep
        line += sep + "{:>15.3f}".format(ap_50o ) + sep
        line += sep + "{:>15.3f}".format(ap_25o ) + sep
        print(line)
        if PREV_KNOWN_CLASSES_LABELS[split][task_ow] != None and label_name in PREV_KNOWN_CLASSES_LABELS[split][task_ow]:
            prev_known_avg.append(avgs["classes"][label_name]["ap"])
            prev_known_avg_25.append(avgs["classes"][label_name]["ap50%"])
            prev_known_avg_50.append(avgs["classes"][label_name]["ap25%"])
        
        if label_name in KNOWN_CLASSES_LABELS[split][task_ow]:
            curr_known_avg.append(avgs["classes"][label_name]["ap"])
            curr_known_avg_25.append(avgs["classes"][label_name]["ap50%"])
            curr_known_avg_50.append(avgs["classes"][label_name]["ap25%"])

    all_ap_avg  = avgs["all_ap"]
    all_ap_50o  = avgs["all_ap_50%"]
    all_ap_25o  = avgs["all_ap_25%"]

    print("-"*lineLen)
    line  = "{:<15}".format("Average") + sep + col1
    line += "{:>15.3f}".format(all_ap_avg)  + sep
    line += "{:>15.3f}".format(all_ap_50o)  + sep
    line += "{:>15.3f}".format(all_ap_25o)  + sep
    print(line)
    print("")
    
    if PREV_KNOWN_CLASSES_LABELS[split][task_ow] != None:
        line  = "{:<15}".format("Previous known") + sep + col1
        line += "{:>15.3f}".format(np.nanmean(np.array(prev_known_avg)))  + sep
        line += "{:>15.3f}".format(np.nanmean(np.array(prev_known_avg_25)))  + sep
        line += "{:>15.3f}".format(np.nanmean(np.array(prev_known_avg_50)))  + sep
        print(line)
        print("")
    
    line  = "{:<15}".format("Current known") + sep + col1
    line += "{:>15.3f}".format(np.nanmean(np.array(curr_known_avg)))  + sep
    line += "{:>15.3f}".format(np.nanmean(np.array(curr_known_avg_25)))  + sep
    line += "{:>15.3f}".format(np.nanmean(np.array(curr_known_avg_50)))  + sep
    print(line)
    print("")
    
    print("")
    print("#"*lineLen)
    line  = ""
    line += "{:<15}".format(""      ) + sep
    line += "{:>30}".format("Unknown class evaluation metrics"        ) + sep
    line += "{:>15}".format(""    ) + sep
    print(line)
    print("#"*lineLen)
    
    line  = "{:<21}".format("U-Recall") + sep + col1
    line += "{:>13}".format("@0.5"        ) + sep
    line += "{:>13}".format("@0.8"    ) + sep
    line += "{:>13}".format("@0.9"    ) + sep
    print(line)
    if type(u_recall).__name__ != 'float':
        line  = "{:<21}".format("") + sep + col1
        line += sep + "{:>13.5f}".format(u_recall[50] ) + sep
        line += sep + "{:>13.5f}".format(u_recall[80]) + sep
        line += sep + "{:>13.5f}".format(u_recall[90]) + sep
    else: 
        line  = "{:<21}".format("") + sep + col1
        line += sep + "{:>13.5f}".format(u_recall) + sep
        line += sep + "{:>13.5f}".format(u_recall) + sep
        line += sep + "{:>13.5f}".format(u_recall) + sep
    
    print(line)
    # line += "{:>12.5f}".format(u_recall)  + sep
    # print(line)
    print("")
    print("-"*lineLen)
    line  = "{:<21}".format("Wilderness Impact") + sep + col1
    line += "{:>13.5f}".format(wi)  + sep
    print(line)
    print("")
    
    line  = "{:<21}".format("A-OSE") + sep + col1
    line += "{:>13d}".format(int(a_ose))  + sep
    print(line)
    print("")
    






def write_result_file(avgs, filename):
    _SPLITTER = ','
    with open(filename, 'w') as f:
        f.write(_SPLITTER.join(['class', 'class id', 'ap', 'ap50', 'ap25']) + '\n')
        for i in range(len(VALID_CLASS_IDS)):
            class_name = CLASS_LABELS[i]
            class_id = VALID_CLASS_IDS[i]
            ap = avgs["classes"][class_name]["ap"]
            ap50 = avgs["classes"][class_name]["ap50%"]
            ap25 = avgs["classes"][class_name]["ap25%"]
            f.write(_SPLITTER.join([str(x) for x in [class_name, class_id, ap, ap50, ap25]]) + '\n')    


def evaluate(current_epoch,config_general, preds: dict, gt_path: str, output_file: str, dataset: str = "scannet"):
    global CLASS_LABELS
    global VALID_CLASS_IDS
    global ID_TO_LABEL
    global LABEL_TO_ID
    global opt
    
    VALID_CLASS_IDS = list(VALID_CLASS_IDS_200) 
    CLASS_LABELS = list(CLASS_LABELS_200)
    CLASS_LABELS.remove('wall')
    CLASS_LABELS.remove('floor')
    VALID_CLASS_IDS.remove(1)
    VALID_CLASS_IDS.remove(3)
    VALID_CLASS_IDS = np.array(VALID_CLASS_IDS)
    CLASS_LABELS = np.array(CLASS_LABELS)
    
    ID_TO_LABEL = {}
    LABEL_TO_ID = {}
    for i in range(len(VALID_CLASS_IDS)):
        LABEL_TO_ID[CLASS_LABELS[i]] = VALID_CLASS_IDS[i]
        ID_TO_LABEL[VALID_CLASS_IDS[i]] = CLASS_LABELS[i]

    total_true = 0
    total_seen = 0
    NUM_CLASSES = len(VALID_CLASS_IDS)

    true_positive_classes = np.zeros(NUM_CLASSES)
    positive_classes = np.zeros(NUM_CLASSES)
    gt_classes = np.zeros(NUM_CLASSES)

    # precision & recall
    total_gt_ins = np.zeros(NUM_CLASSES)
    at = 0.5
    tpsins = [[] for _ in range(NUM_CLASSES)]
    fpsins = [[] for _ in range(NUM_CLASSES)]
    # mucov and mwcov
    all_mean_cov = [[] for _ in range(NUM_CLASSES)]
    all_mean_weighted_cov = [[] for _ in range(NUM_CLASSES)]

    print('evaluating', len(preds), 'scans...')
    matches = {}
    for i,(k,v) in enumerate(preds.items()):
        gt_file = os.path.join(gt_path, k + ".txt")
        if not os.path.isfile(gt_file):
            util.print_error('Scan {} does not match any gt file'.format(k), user_fault=True)

        if dataset == "s3dis":
            gt_ids = util_3d.load_ids(gt_file)
            gt_sem = (gt_ids // 1000) - 1
            gt_ins = gt_ids - (gt_ids // 1000) * 1000

            # pred_sem = v['pred_classes'] - 1
            pred_sem = np.zeros(v['pred_masks'].shape[0], dtype=np.int)
            # TODO CONTINUE HERE!!!!!!!!!!!!!
            pred_ins = np.zeros(v['pred_masks'].shape[0], dtype=np.int)


            for inst_id in reversed(range(v['pred_masks'].shape[1])):
                point_ids = np.argwhere(v['pred_masks'][:, inst_id] == 1.)[:, 0]
                pred_ins[point_ids] = inst_id + 1
                pred_sem[point_ids] = v['pred_classes'][inst_id] - 1

            # semantic acc
            total_true += np.sum(pred_sem == gt_sem)
            total_seen += pred_sem.shape[0]


            # TODO PARALLELIZ THIS!!!!!!!
            # pn semantic mIoU
            '''
            for j in range(gt_sem.shape[0]):
                gt_l = int(gt_sem[j])
                pred_l = int(pred_sem[j])
                gt_classes[gt_l] += 1
                positive_classes[pred_l] += 1
                true_positive_classes[gt_l] += int(gt_l == pred_l)
            '''


            uniq, counts = np.unique(pred_sem, return_counts=True)
            positive_classes[uniq] += counts

            uniq, counts = np.unique(gt_sem, return_counts=True)
            gt_classes[uniq] += counts

            uniq, counts = np.unique(gt_sem[pred_sem == gt_sem], return_counts=True)
            true_positive_classes[uniq] += counts

            # instance
            un = np.unique(pred_ins)
            pts_in_pred = [[] for _ in range(NUM_CLASSES)]
            for ig, g in enumerate(un):  # each object in prediction
                if g == -1:
                    continue
                tmp = (pred_ins == g)
                sem_seg_i = int(stats.mode(pred_sem[tmp])[0])
                pts_in_pred[sem_seg_i] += [tmp]

            un = np.unique(gt_ins)
            pts_in_gt = [[] for _ in range(NUM_CLASSES)]
            for ig, g in enumerate(un):
                tmp = (gt_ins == g)
                sem_seg_i = int(stats.mode(gt_sem[tmp])[0])
                pts_in_gt[sem_seg_i] += [tmp]

            # instance mucov & mwcov
            for i_sem in range(NUM_CLASSES):
                sum_cov = 0
                mean_cov = 0
                mean_weighted_cov = 0
                num_gt_point = 0
                for ig, ins_gt in enumerate(pts_in_gt[i_sem]):
                    ovmax = 0.
                    num_ins_gt_point = np.sum(ins_gt)
                    num_gt_point += num_ins_gt_point
                    for ip, ins_pred in enumerate(pts_in_pred[i_sem]):
                        union = (ins_pred | ins_gt)
                        intersect = (ins_pred & ins_gt)
                        iou = float(np.sum(intersect)) / np.sum(union)

                        if iou > ovmax:
                            ovmax = iou
                            ipmax = ip

                    sum_cov += ovmax
                    mean_weighted_cov += ovmax * num_ins_gt_point

                if len(pts_in_gt[i_sem]) != 0:
                    mean_cov = sum_cov / len(pts_in_gt[i_sem])
                    all_mean_cov[i_sem].append(mean_cov)

                    mean_weighted_cov /= num_gt_point
                    all_mean_weighted_cov[i_sem].append(mean_weighted_cov)

       
        matches_key = os.path.abspath(gt_file)
        # assign gt to predictions
        gt2pred, pred2gt = assign_instances_for_scan(config_general.split,config_general.OW_task, v, gt_file)
        matches[matches_key] = {}
        matches[matches_key]['gt'] = gt2pred
        matches[matches_key]['pred'] = pred2gt
        
        
        sys.stdout.write("\rscans processed: {}".format(i+1))
        sys.stdout.flush()
    
    print('')
    U_Recall, WI, A_OSE, ap_scores = evaluate_matches(matches)
    avgs = compute_averages(ap_scores)
    
    # print
    if U_Recall and WI and A_OSE:
        print_results(avgs,current_epoch,config_general.save_dir,config_general.split,config_general.OW_task,U_Recall, WI[0.5], A_OSE[0.5])
    else:
        print_results(avgs,current_epoch,config_general.save_dir,config_general.split,config_general.OW_task,0.0, 0.0, 0.0)


def main():
    pred_files = [f for f in os.listdir(opt.pred_path) if f.endswith('.txt') and f != 'semantic_instance_evaluation.txt']
    gt_files = []
    if len(pred_files) == 0:
        util.print_error('No result files found.', user_fault=True)
    for i in range(len(pred_files)):
        gt_file = os.path.join(opt.gt_path, pred_files[i])
        if not os.path.isfile(gt_file):
            util.print_error('Result file {} does not match any gt file'.format(pred_files[i]), user_fault=True)
        gt_files.append(gt_file)
        pred_files[i] = os.path.join(opt.pred_path, pred_files[i])

    # evaluate
    evaluate(pred_files, gt_files, opt.pred_path, opt.output_file)


if __name__ == '__main__':
    main()

ow_benchmark/evaluate_semantic_instance.py
aminebdj/3D-OWIS

### ScanNet200 Benchmark constants ###
VALID_CLASS_IDS_200 = (
1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 13, 14, 15, 16, 17, 18, 19, 21, 22, 23, 24, 26, 27, 28, 29, 31, 32, 33, 34, 35, 36, 38, 39, 40, 41, 42, 44, 45, 46, 47, 48, 49, 50, 51, 52, 54, 55, 56, 57, 58, 59, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71,
72, 73, 74, 75, 76, 77, 78, 79, 80, 82, 84, 86, 87, 88, 89, 90, 93, 95, 96, 97, 98, 99, 100, 101, 102, 103, 104, 105, 106, 107, 110, 112, 115, 116, 118, 120, 121, 122, 125, 128, 130, 131, 132, 134, 136, 138, 139, 140, 141, 145, 148, 154,
155, 156, 157, 159, 161, 163, 165, 166, 168, 169, 170, 177, 180, 185, 188, 191, 193, 195, 202, 208, 213, 214, 221, 229, 230, 232, 233, 242, 250, 261, 264, 276, 283, 286, 300, 304, 312, 323, 325, 331, 342, 356, 370, 392, 395, 399, 408, 417,
488, 540, 562, 570, 572, 581, 609, 748, 776, 1156, 1163, 1164, 1165, 1166, 1167, 1168, 1169, 1170, 1171, 1172, 1173, 1174, 1175, 1176, 1178, 1179, 1180, 1181, 1182, 1183, 1184, 1185, 1186, 1187, 1188, 1189, 1190, 1191,3000)

CLASS_LABELS_200 = (
'wall', 'chair', 'floor', 'table', 'door', 'couch', 'cabinet', 'shelf', 'desk', 'office chair', 'bed', 'pillow', 'sink', 'picture', 'window', 'toilet', 'bookshelf', 'monitor', 'curtain', 'book', 'armchair', 'coffee table', 'box',
'refrigerator', 'lamp', 'kitchen cabinet', 'towel', 'clothes', 'tv', 'nightstand', 'counter', 'dresser', 'stool', 'cushion', 'plant', 'ceiling', 'bathtub', 'end table', 'dining table', 'keyboard', 'bag', 'backpack', 'toilet paper',
'printer', 'tv stand', 'whiteboard', 'blanket', 'shower curtain', 'trash can', 'closet', 'stairs', 'microwave', 'stove', 'shoe', 'computer tower', 'bottle', 'bin', 'ottoman', 'bench', 'board', 'washing machine', 'mirror', 'copier',
'basket', 'sofa chair', 'file cabinet', 'fan', 'laptop', 'shower', 'paper', 'person', 'paper towel dispenser', 'oven', 'blinds', 'rack', 'plate', 'blackboard', 'piano', 'suitcase', 'rail', 'radiator', 'recycling bin', 'container',
'wardrobe', 'soap dispenser', 'telephone', 'bucket', 'clock', 'stand', 'light', 'laundry basket', 'pipe', 'clothes dryer', 'guitar', 'toilet paper holder', 'seat', 'speaker', 'column', 'bicycle', 'ladder', 'bathroom stall', 'shower wall',
'cup', 'jacket', 'storage bin', 'coffee maker', 'dishwasher', 'paper towel roll', 'machine', 'mat', 'windowsill', 'bar', 'toaster', 'bulletin board', 'ironing board', 'fireplace', 'soap dish', 'kitchen counter', 'doorframe',
'toilet paper dispenser', 'mini fridge', 'fire extinguisher', 'ball', 'hat', 'shower curtain rod', 'water cooler', 'paper cutter', 'tray', 'shower door', 'pillar', 'ledge', 'toaster oven', 'mouse', 'toilet seat cover dispenser',
'furniture', 'cart', 'storage container', 'scale', 'tissue box', 'light switch', 'crate', 'power outlet', 'decoration', 'sign', 'projector', 'closet door', 'vacuum cleaner', 'candle', 'plunger', 'stuffed animal', 'headphones', 'dish rack',
'broom', 'guitar case', 'range hood', 'dustpan', 'hair dryer', 'water bottle', 'handicap bar', 'purse', 'vent', 'shower floor', 'water pitcher', 'mailbox', 'bowl', 'paper bag', 'alarm clock', 'music stand', 'projector screen', 'divider',
'laundry detergent', 'bathroom counter', 'object', 'bathroom vanity', 'closet wall', 'laundry hamper', 'bathroom stall door', 'ceiling light', 'trash bin', 'dumbbell', 'stair rail', 'tube', 'bathroom cabinet', 'cd case', 'closet rod',
'coffee kettle', 'structure', 'shower head', 'keyboard piano', 'case of water bottles', 'coat rack', 'storage organizer', 'folded chair', 'fire alarm', 'power strip', 'calendar', 'poster', 'potted plant', 'luggage', 'mattress','unknown')


#wall and floor are excluded from the training in all splits


SPLIT_A_SUBSET_1_IDS = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 32, 33, 34, 39, 42, 43, 45, 46, 47, 48, 50, 57, 58, 59, 60, 62, 63, 66, 71, 74, 78, 79, 81, 90, 98, 99, 115, 116, 161, 170, 171, 172, 178]
SPLIT_A_SUBSET_2_IDS = [31, 35, 36, 37, 38, 40, 41, 44, 49, 51, 52, 53, 54, 55, 56, 61, 64, 65, 68, 69, 70, 72, 75, 76, 77, 80, 83, 86, 87, 88, 89, 93, 95, 96, 97, 101, 102, 103, 104, 106, 107, 108, 111, 113, 118, 123, 126, 127, 128, 132, 133, 140, 143, 144, 149, 152, 166, 167, 169, 173, 174, 175, 176, 180, 184, 189, 195, 197]
SPLIT_A_SUBSET_3_IDS = [67, 73, 82, 84, 85, 91, 92, 94, 100, 105, 109, 110, 112, 114, 117, 119, 120, 121, 122, 124, 125, 129, 130, 131, 134, 135, 136, 137, 138, 139, 141, 142, 145, 146, 147, 148, 150, 151, 153, 154, 155, 156, 157, 158, 159, 160, 162, 163, 164, 165, 168, 177, 179, 181, 182, 183, 185, 186, 187, 188, 190, 191, 192, 193, 194, 196]
SPLIT_A_SUBSET_1_LABELS = ['tv stand', 'curtain', 'blinds', 'shower curtain', 'bookshelf', 'tv', 'kitchen cabinet', 'pillow', 'lamp', 'dresser', 'monitor', 'object', 'ceiling', 'board', 'stove', 'closet wall', 'couch', 'office chair', 'kitchen counter', 'shower', 'closet', 'doorframe', 'sofa chair', 'mailbox', 'nightstand', 'washing machine', 'picture', 'book', 'sink', 'recycling bin', 'table', 'backpack', 'shower wall', 'toilet', 'copier', 'counter', 'stool', 'refrigerator', 'window', 'file cabinet', 'chair', 'plant', 'coffee table', 'stairs', 'armchair', 'cabinet', 'bathroom vanity', 'bathroom stall', 'mirror', 'blackboard', 'trash can', 'stair rail', 'box', 'towel', 'door', 'clothes', 'whiteboard', 'bed', 'bathtub', 'desk', 'wardrobe', 'clothes dryer', 'radiator', 'shelf']
SPLIT_A_SUBSET_2_LABELS = ["cushion", "end table", "dining table", "keyboard", "bag", "toilet paper", "printer", "blanket", "microwave", "shoe", "computer tower", "bottle", "bin", "ottoman", "bench", "basket", "fan", "laptop", "person", "paper towel dispenser", "oven", "rack", "piano", "suitcase", "rail", "container", "telephone", "stand", "light", "laundry basket", "pipe", "seat", "column", "bicycle", "ladder", "jacket", "storage bin", "coffee maker", "dishwasher", "machine", "mat", "windowsill", "bulletin board", "fireplace", "mini fridge", "water cooler", "shower door", "pillar", "ledge", "furniture", "cart", "decoration", "closet door", "vacuum cleaner", "dish rack", "range hood", "projector screen", "divider", "bathroom counter", "laundry hamper", "bathroom stall door", "ceiling light", "trash bin", "bathroom cabinet", "structure", "storage organizer", "potted plant", "mattress"]
SPLIT_A_SUBSET_3_LABELS = ["paper", "plate", "soap dispenser", "bucket", "clock", "guitar", "toilet paper holder", "speaker", "cup", "paper towel roll", "bar", "toaster", "ironing board", "soap dish", "toilet paper dispenser", "fire extinguisher", "ball", "hat", "shower curtain rod", "paper cutter", "tray", "toaster oven", "mouse", "toilet seat cover dispenser", "storage container", "scale", "tissue box", "light switch", "crate", "power outlet", "sign", "projector", "candle", "plunger", "stuffed animal", "headphones", "broom", "guitar case", "dustpan", "hair dryer", "water bottle", "handicap bar", "purse", "vent", "shower floor", "water pitcher", "bowl", "paper bag", "alarm clock", "music stand", "laundry detergent", "dumbbell", "tube", "cd case", "closet rod", "coffee kettle", "shower head", "keyboard piano", "case of water bottles", "coat rack", "folded chair", "fire alarm", "power strip", "calendar", "poster", "luggage"]

SPLIT_B_SUBSET_1_IDS = [0, 1, 2, 3, 5, 8, 9, 11, 16, 18, 19, 22, 25, 26, 27, 29, 30, 32, 33, 35, 36, 38, 39, 42, 44, 47, 51, 55, 56, 62, 64, 65, 72, 75, 76, 77, 79, 80, 81, 87, 88, 94, 96, 101, 102, 112, 113, 118, 121, 127, 136, 139, 140, 141, 143, 144, 145, 151, 157, 164, 165, 170, 172, 173, 177, 186, 187, 188, 189, 194, 195, 196, 197]
SPLIT_B_SUBSET_2_IDS = [4, 6, 7, 12, 14, 15, 17, 20, 37, 41, 43, 48, 52, 54, 57, 60, 63, 67, 68, 71, 74, 78, 83, 84, 85, 86, 89, 91, 93, 95, 100, 105, 106, 108, 111, 120, 124, 125, 130, 132, 133, 134, 142, 148, 155, 161, 166, 167, 175, 178, 179, 181, 184, 190, 192]
SPLIT_B_SUBSET_3_IDS = [10, 13, 21, 23, 24, 28, 31, 34, 40, 45, 46, 49, 50, 53, 58, 59, 61, 66, 69, 70, 73, 82, 90, 92, 97, 98, 99, 103, 104, 107, 109, 110, 114, 115, 116, 117, 119, 122, 123, 126, 128, 129, 131, 135, 137, 138, 146, 147, 149, 150, 152, 153, 154, 156, 158, 159, 160, 162, 163, 168, 169, 171, 174, 176, 180, 182, 183, 185, 191, 193]
SPLIT_B_SUBSET_1_LABELS = ['alarm clock', 'backpack', 'bag', 'bed', 'blanket', 'case of water bottles', 'ceiling', 'closet', 'closet door', 'closet wall', 'clothes', 'coat rack', 'container', 'curtain', 'door', 'dresser', 'dumbbell', 'fan', 'guitar case', 'hat', 'ironing board', 'lamp', 'laptop', 'laundry basket', 'laundry hamper', 'luggage', 'mattress', 'mini fridge', 'nightstand', 'object', 'pillow', 'poster', 'power outlet', 'purse', 'rack', 'recycling bin', 'shelf', 'shoe', 'sign', 'storage bin', 'storage organizer', 'suitcase', 'tissue box', 'wardrobe', 'decoration', 'armchair', 'bench', 'bicycle', 'candle', 'chair', 'coffee table', 'couch', 'dining table', 'end table', 'fireplace', 'jacket', 'keyboard piano', 'light', 'music stand', 'ottoman', 'piano', 'picture', 'pillar', 'plant', 'potted plant', 'rail', 'sofa chair', 'speaker', 'stool', 'table', 'tv', 'tv stand', 'vacuum cleaner']
SPLIT_B_SUBSET_2_LABELS = ['guitar', 'paper towel roll', 'book', 'bookshelf', 'cart', 'furniture', 'blackboard', 'projector', 'seat', 'folded chair', 'office chair', 'projector screen', 'whiteboard', 'bin', 'bucket', 'bulletin board', 'copier', 'machine', 'mailbox', 'paper cutter', 'printer', 'column', 'storage container', 'blinds', 'structure', 'water bottle', 'ball', 'board', 'box', 'cabinet', 'cd case', 'ceiling light', 'clock', 'computer tower', 'cup', 'desk', 'divider', 'file cabinet', 'headphones', 'keyboard', 'monitor', 'mouse', 'paper', 'person', 'power strip', 'radiator', 'stand', 'telephone', 'tray', 'tube', 'window', 'windowsill', 'pipe', 'stair rail', 'stairs']
SPLIT_B_SUBSET_3_LABELS = ['bar', 'basket', 'bathroom cabinet', 'bathroom counter', 'bathroom stall', 'bathroom stall door', 'bathroom vanity', 'bathtub', 'bottle', 'broom', 'clothes dryer', 'cushion', 'doorframe', 'fire alarm', 'hair dryer', 'handicap bar', 'ledge', 'light switch', 'mat', 'mirror', 'paper towel dispenser', 'plunger', 'scale', 'shower', 'shower curtain', 'shower curtain rod', 'shower door', 'shower floor', 'shower head', 'shower wall', 'sink', 'soap dish', 'soap dispenser', 'toilet', 'toilet paper', 'toilet paper dispenser', 'toilet paper holder', 'toilet seat cover dispenser', 'towel', 'trash bin', 'washing machine', 'closet rod', 'dustpan', 'laundry detergent', 'stuffed animal', 'bowl', 'calendar', 'coffee kettle', 'coffee maker', 'counter', 'dish rack', 'dishwasher', 'fire extinguisher', 'kitchen cabinet', 'kitchen counter', 'microwave', 'oven', 'paper bag', 'plate', 'range hood', 'refrigerator', 'stove', 'toaster', 'toaster oven', 'trash can', 'vent', 'water cooler', 'water pitcher', 'crate', 'ladder']

SPLIT_C_SUBSET_1_IDS = [0, 2, 4, 5, 6, 8, 20, 33, 36, 37, 38, 39, 40, 43, 46, 50, 51, 52, 57, 60, 61, 66, 69, 73, 75, 76, 77, 78, 88, 91, 94, 102, 105, 109, 111, 119, 121, 123, 124, 127, 128, 129, 132, 136, 139, 140, 144, 146, 147, 148, 154, 155, 160, 161, 164, 168, 170, 171, 172, 173, 177, 178, 180, 185, 195, 196]
SPLIT_C_SUBSET_2_IDS = [3, 10, 13, 14, 15, 16, 17, 18, 19, 22, 25, 27, 28, 34, 42, 44, 47, 48, 55, 58, 59, 62, 64, 65, 67, 70, 71, 79, 81, 83, 84, 85, 87, 89, 90, 95, 96, 97, 99, 100, 106, 112, 113, 116, 118, 120, 122, 126, 130, 131, 135, 141, 149, 151, 156, 162, 163, 165, 167, 176, 183, 184, 186, 187, 192, 194]
SPLIT_C_SUBSET_3_IDS = [1, 7, 9, 11, 12, 21, 23, 24, 26, 29, 30, 31, 32, 35, 41, 45, 49, 53, 54, 56, 63, 68, 72, 74, 80, 82, 86, 92, 93, 98, 101, 103, 104, 107, 108, 110, 114, 115, 117, 125, 133, 134, 137, 138, 142, 143, 145, 150, 152, 153, 157, 158, 159, 166, 169, 174, 175, 179, 181, 182, 188, 189, 190, 191, 193, 197]
SPLIT_C_SUBSET_1_LABELS = ['basket', 'trash can', 'stair rail', 'toaster oven', 'laundry hamper', 'bulletin board', 'dining table', 'stuffed animal', 'bathroom vanity', 'box', 'ceiling', 'potted plant', 'luggage', 'closet wall', 'paper cutter', 'desk', 'object', 'rail', 'tissue box', 'plate', 'keyboard', 'hat', 'copier', 'shower head', 'bed', 'paper towel dispenser', 'fire extinguisher', 'paper towel roll', 'backpack', 'water bottle', 'bathroom cabinet', 'stove', 'laundry basket', 'alarm clock', 'headphones', 'piano', 'guitar', 'bag', 'door', 'speaker', 'water cooler', 'shoe', 'water pitcher', 'dumbbell', 'furniture', 'decoration', 'radiator', 'plunger', 'shower', 'bar', 'hair dryer', 'suitcase', 'cabinet', 'chair', 'board', 'laundry detergent', 'whiteboard', 'vacuum cleaner', 'power outlet', 'storage bin', 'computer tower', 'mailbox', 'shelf', 'ledge', 'pillar', 'toilet paper']
SPLIT_C_SUBSET_2_LABELS = ['ironing board', 'divider', 'oven', 'dish rack', 'shower door', 'mini fridge', 'bicycle', 'laptop', 'armchair', 'couch', 'coffee kettle', 'counter', 'structure', 'pipe', 'bowl', 'shower curtain rod', 'sofa chair', 'clothes dryer', 'coffee table', 'stairs', 'toilet seat cover dispenser', 'machine', 'paper bag', 'book', 'blinds', 'monitor', 'shower wall', 'curtain', 'closet', 'telephone', 'fan', 'ball', 'bucket', 'sign', 'mirror', 'clock', 'nightstand', 'tv stand', 'handicap bar', 'poster', 'blanket', 'cup', 'recycling bin', 'lamp', 'scale', 'mouse', 'wardrobe', 'ottoman', 'paper', 'power strip', 'fireplace', 'doorframe', 'toilet', 'trash bin', 'case of water bottles', 'light', 'washing machine', 'guitar case', 'sink', 'bathtub', 'ladder', 'bookshelf', 'column', 'clothes', 'keyboard piano', 'music stand']
SPLIT_C_SUBSET_3_LABELS = ['mattress', 'toaster', 'stool', 'plant', 'folded chair', 'microwave', 'cushion', 'bench', 'soap dispenser', 'storage organizer', 'shower curtain', 'cart', 'kitchen counter', 'towel', 'blackboard', 'tv', 'printer', 'stand', 'rack', 'bathroom counter', 'closet rod', 'bottle', 'range hood', 'purse', 'candle', 'person', 'coffee maker', 'light switch', 'storage container', 'bathroom stall door', 'shower floor', 'kitchen cabinet', 'refrigerator', 'fire alarm', 'tube', 'toilet paper holder', 'ceiling light', 'picture', 'end table', 'closet door', 'file cabinet', 'crate', 'toilet paper dispenser', 'pillow', 'mat', 'bathroom stall', 'broom', 'container', 'seat', 'jacket', 'dresser', 'dustpan', 'table', 'projector', 'window', 'windowsill', 'tray', 'cd case', 'soap dish', 'office chair', 'dishwasher', 'vent', 'coat rack', 'calendar', 'bin', 'projector screen']

"""
IDS for currently known, unknown, and previously known classes for each split and task.

IMPORTANT: IDS in this script correspond to the index of the prediction of the class in the classification head, 
           and NOT the label IDS defined by scannet200 benchmark.
"""

# UNKNOWN_CLASSES_IDS: IDS of the unknown classes
UNKNOWN_CLASSES_IDS = {
    'A':{
        'task1': SPLIT_A_SUBSET_2_IDS+SPLIT_A_SUBSET_3_IDS,
        'task2': SPLIT_A_SUBSET_3_IDS,
        'task3': None
    },
    'B':{
        'task1': SPLIT_B_SUBSET_2_IDS+SPLIT_B_SUBSET_3_IDS,
        'task2': SPLIT_B_SUBSET_3_IDS,
        'task3': None
    },
    'C':{
        'task1': SPLIT_C_SUBSET_2_IDS+SPLIT_C_SUBSET_3_IDS,
        'task2': SPLIT_C_SUBSET_3_IDS,
        'task3': None
    }
}

# KNOWN_CLASSES_IDS: IDS of the currently known classes
KNOWN_CLASSES_IDS = {
    'A': {
        'task1': SPLIT_A_SUBSET_1_IDS,
        'task2': SPLIT_A_SUBSET_2_IDS,
        'task3': SPLIT_A_SUBSET_3_IDS
    },
    'B': {
        'task1': SPLIT_B_SUBSET_1_IDS,
        'task2': SPLIT_B_SUBSET_2_IDS,
        'task3': SPLIT_B_SUBSET_3_IDS
    },
    'C': {
        'task1': SPLIT_C_SUBSET_1_IDS,
        'task2': SPLIT_C_SUBSET_2_IDS,
        'task3': SPLIT_C_SUBSET_3_IDS
    }
}

# PREV_KNOWN_CLASSES_IDS: IDS of the previously known classes
PREV_KNOWN_CLASSES_IDS = {
    'A': {
        'task1': None,
        'task2': SPLIT_A_SUBSET_1_IDS,
        'task3': SPLIT_A_SUBSET_1_IDS+SPLIT_A_SUBSET_2_IDS
    },
    'B': {
        'task1': None,  
        'task2': SPLIT_B_SUBSET_1_IDS,
        'task3': SPLIT_B_SUBSET_1_IDS+SPLIT_B_SUBSET_2_IDS
    },
    'C': {
        'task1': None,  
        'task2': SPLIT_C_SUBSET_1_IDS,
        'task3': SPLIT_C_SUBSET_1_IDS+SPLIT_C_SUBSET_2_IDS
    }
}

"""
LABELS for currently known, unknown, and previously known classes for each split and task.
"""

PREV_KNOWN_CLASSES_LABELS = {
    'A': {
        'task1': None,
        'task2': SPLIT_A_SUBSET_1_LABELS,
        'task3': SPLIT_A_SUBSET_1_LABELS+SPLIT_A_SUBSET_2_LABELS
    },
    'B': {
        'task1': None,  
        'task2': SPLIT_B_SUBSET_1_LABELS,
        'task3': SPLIT_B_SUBSET_1_LABELS+SPLIT_B_SUBSET_2_LABELS
    },
    'C': {
        'task1': None,  
        'task2': SPLIT_C_SUBSET_1_LABELS,
        'task3': SPLIT_C_SUBSET_1_LABELS+SPLIT_C_SUBSET_2_LABELS
    }
}

KNOWN_CLASSES_LABELS = {
    'A' :{
        'task1': SPLIT_A_SUBSET_1_LABELS,
        'task2': SPLIT_A_SUBSET_2_LABELS,
        'task3': SPLIT_A_SUBSET_3_LABELS
    },
    'B': {
        'task1': SPLIT_B_SUBSET_1_LABELS,
        'task2': SPLIT_B_SUBSET_2_LABELS,
        'task3': SPLIT_B_SUBSET_3_LABELS
    },
    
    'C': {
        'task1': SPLIT_C_SUBSET_1_LABELS,
        'task2': SPLIT_C_SUBSET_2_LABELS,
        'task3': SPLIT_C_SUBSET_3_LABELS
    }
}

UNKNOWN_CLASSES_LABELS = {
    'A' :{
        'task1': SPLIT_A_SUBSET_2_LABELS+SPLIT_A_SUBSET_3_LABELS,
        'task2': SPLIT_A_SUBSET_3_LABELS,
        'task3': None
    },
    'B': {
        'task1': SPLIT_B_SUBSET_2_LABELS+SPLIT_B_SUBSET_3_LABELS,
        'task2': SPLIT_B_SUBSET_3_LABELS,
        'task3': None,
    },
    'C': {
        'task1': SPLIT_C_SUBSET_2_LABELS+SPLIT_C_SUBSET_3_LABELS,
        'task2': SPLIT_C_SUBSET_3_LABELS,
        'task3': None,
    }
}


ow_benchmark/owis_splits.py
aminebdj/3D-OWIS
import os, sys
import csv
try:
    import numpy as np
except:
    print("Failed to import numpy package.")
    sys.exit(-1)
try:
    import imageio
except:
    print("Please install the module 'imageio' for image processing, e.g.")
    print("pip install imageio")
    sys.exit(-1)

# print an error message and quit
def print_error(message, user_fault=False):
    sys.stderr.write('ERROR: ' + str(message) + '\n')
    if user_fault:
        sys.exit(2)
    sys.exit(-1)


# if string s represents an int
def represents_int(s):
    try:
        int(s)
        return True
    except ValueError:
        return False


def read_label_mapping(filename, label_from='raw_category', label_to='nyu40id'):
    assert os.path.isfile(filename)
    mapping = dict()
    with open(filename) as csvfile:
        reader = csv.DictReader(csvfile, delimiter='\t')
        for row in reader:
            mapping[row[label_from]] = int(row[label_to])
    # if ints convert
    if represents_int(list(mapping.keys())[0]):
        mapping = {int(k):v for k,v in mapping.items()}
    return mapping


# input: scene_types.txt or scene_types_all.txt
def read_scene_types_mapping(filename, remove_spaces=True):
    assert os.path.isfile(filename)
    mapping = dict()
    lines = open(filename).read().splitlines()
    lines = [line.split('\t') for line in lines]
    if remove_spaces:
        mapping = { x[1].strip():int(x[0]) for x in lines }
    else:
        mapping = { x[1]:int(x[0]) for x in lines }
    return mapping


# color by label
def visualize_label_image(filename, image):
    height = image.shape[0]
    width = image.shape[1]
    vis_image = np.zeros([height, width, 3], dtype=np.uint8)
    color_palette = create_color_palette()
    for idx, color in enumerate(color_palette):
        vis_image[image==idx] = color
    imageio.imwrite(filename, vis_image)


# color by different instances (mod length of color palette)
def visualize_instance_image(filename, image):
    height = image.shape[0]
    width = image.shape[1]
    vis_image = np.zeros([height, width, 3], dtype=np.uint8)
    color_palette = create_color_palette()
    instances = np.unique(image)
    for idx, inst in enumerate(instances):
        vis_image[image==inst] = color_palette[inst%len(color_palette)]
    imageio.imwrite(filename, vis_image)


# color palette for nyu40 labels
def create_color_palette():
    return [
        (0, 0, 0),
        (174, 199, 232),		# wall
        (152, 223, 138),		# floor
        (31, 119, 180), 		# cabinet
        (255, 187, 120),		# bed
        (188, 189, 34), 		# chair
        (140, 86, 75),  		# sofa
        (255, 152, 150),		# table
        (214, 39, 40),  		# door
        (197, 176, 213),		# window
        (148, 103, 189),		# bookshelf
        (196, 156, 148),		# picture
        (23, 190, 207), 		# counter
        (178, 76, 76),
        (247, 182, 210),		# desk
        (66, 188, 102),
        (219, 219, 141),		# curtain
        (140, 57, 197),
        (202, 185, 52),
        (51, 176, 203),
        (200, 54, 131),
        (92, 193, 61),
        (78, 71, 183),
        (172, 114, 82),
        (255, 127, 14), 		# refrigerator
        (91, 163, 138),
        (153, 98, 156),
        (140, 153, 101),
        (158, 218, 229),		# shower curtain
        (100, 125, 154),
        (178, 127, 135),
        (120, 185, 128),
        (146, 111, 194),
        (44, 160, 44),  		# toilet
        (112, 128, 144),		# sink
        (96, 207, 209),
        (227, 119, 194),		# bathtub
        (213, 92, 176),
        (94, 106, 211),
        (82, 84, 163),  		# otherfurn
        (100, 85, 144)
    ]

ow_benchmark/util.py
aminebdj/3D-OWIS
import os, sys
import json

try:
    import numpy as np
except:
    print("Failed to import numpy package.")
    sys.exit(-1)

try:
    from plyfile import PlyData, PlyElement
except:
    print("Please install the module 'plyfile' for PLY i/o, e.g.")
    print("pip install plyfile")
    sys.exit(-1)

import benchmark.util as util
############################################################################################################################
from datasets.scannet200.scannet200_splits import HEAD_CATS_SCANNET_200, TAIL_CATS_SCANNET_200, COMMON_CATS_SCANNET_200
from datasets.scannet200.owis_splits import UNKNOWN_CLASSES_LABELS
############################################################################################################################
# matrix: 4x4 np array
# points Nx3 np array
def transform_points(matrix, points):
    assert len(points.shape) == 2 and points.shape[1] == 3
    num_points = points.shape[0]
    p = np.concatenate([points, np.ones((num_points, 1))], axis=1)
    p = np.matmul(matrix, np.transpose(p))
    p = np.transpose(p)
    p[:,:3] /= p[:,3,None]
    return p[:,:3]


def export_ids(filename, ids):
    with open(filename, 'w') as f:
        for id in ids:
            f.write('%d\n' % id)


def load_ids(filename):
    ids = open(filename).read().splitlines()
    ids = np.array(ids, dtype=np.int64)
    return ids


def read_mesh_vertices(filename):
    assert os.path.isfile(filename)
    with open(filename, 'rb') as f:
        plydata = PlyData.read(f)
        num_verts = plydata['vertex'].count
        vertices = np.zeros(shape=[num_verts, 3], dtype=np.float32)
        vertices[:,0] = plydata['vertex'].data['x']
        vertices[:,1] = plydata['vertex'].data['y']
        vertices[:,2] = plydata['vertex'].data['z']
    return vertices


# export 3d instance labels for instance evaluation
def export_instance_ids_for_eval(filename, label_ids, instance_ids):
    assert label_ids.shape[0] == instance_ids.shape[0]
    output_mask_path_relative = 'pred_mask'
    name = os.path.splitext(os.path.basename(filename))[0]
    output_mask_path = os.path.join(os.path.dirname(filename), output_mask_path_relative)
    if not os.path.isdir(output_mask_path):
        os.mkdir(output_mask_path)
    insts = np.unique(instance_ids)
    zero_mask = np.zeros(shape=(instance_ids.shape[0]), dtype=np.int32)
    with open(filename, 'w') as f:
        for idx, inst_id in enumerate(insts):
            if inst_id == 0:  # 0 -> no instance for this vertex
                continue
            output_mask_file = os.path.join(output_mask_path_relative, name + '_' + str(idx) + '.txt')
            loc = np.where(instance_ids == inst_id)
            label_id = label_ids[loc[0][0]]
            f.write('%s %d %f\n' % (output_mask_file, label_id, 1.0))
            # write mask
            mask = np.copy(zero_mask)
            mask[loc[0]] = 1
            export_ids(output_mask_file, mask)


# ------------ Instance Utils ------------ #

class Instance(object):
    instance_id = 0
    label_id = 0
    vert_count = 0
    med_dist = -1
    dist_conf = 0.0

    def __init__(self,mesh_vert_instances, instance_id):
        if (instance_id == -1):
            return
        self.instance_id     = int(instance_id)
        self.label_id    = int(self.get_label_id(instance_id))
        self.vert_count = int(self.get_instance_verts(mesh_vert_instances, instance_id))

    def get_label_id(self, instance_id):
        return int(instance_id // 1000)

    def get_instance_verts(self, mesh_vert_instances, instance_id):
        return (mesh_vert_instances == instance_id).sum()

    def to_json(self):
        return json.dumps(self, default=lambda o: o.__dict__, sort_keys=True, indent=4)

    def to_dict(self):
        dict = {}
        dict["instance_id"] = self.instance_id
        dict["label_id"]    = self.label_id
        dict["vert_count"]  = self.vert_count
        dict["med_dist"]    = self.med_dist
        dict["dist_conf"]   = self.dist_conf
        return dict

    def from_json(self, data):
        self.instance_id     = int(data["instance_id"])
        self.label_id        = int(data["label_id"])
        self.vert_count      = int(data["vert_count"])
        if ("med_dist" in data):
            self.med_dist    = float(data["med_dist"])
            self.dist_conf   = float(data["dist_conf"])

    def __str__(self):
        return "("+str(self.instance_id)+")"

def read_instance_prediction_file(filename, pred_path):
    lines = open(filename).read().splitlines()
    instance_info = {}
    abs_pred_path = os.path.abspath(pred_path)
    for line in lines:
        parts = line.split(' ')
        if len(parts) != 3:
            util.print_error('invalid instance prediction file. Expected (per line): [rel path prediction] [label id prediction] [confidence prediction]')
        if os.path.isabs(parts[0]):
            util.print_error('invalid instance prediction file. First entry in line must be a relative path')
        mask_file = os.path.join(os.path.dirname(filename), parts[0])
        mask_file = os.path.abspath(mask_file)
        # check that mask_file lives inside prediction path
        if os.path.commonprefix([mask_file, abs_pred_path]) != abs_pred_path:
            util.print_error('predicted mask {} in prediction text file {} points outside of prediction path.'.format(mask_file,filename))

        info            = {}
        info["label_id"] = int(float(parts[1]))
        info["conf"]    = float(parts[2])
        instance_info[mask_file]  = info
    return instance_info


def get_instances(split,task, ids, class_ids, class_labels, id2label):
    instances = {}
    for label in class_labels:
        instances[label] = []
    instance_ids = np.unique(ids)
    for id in instance_ids:
        if id == 0:
            continue
        inst = Instance(ids, id)
        #common tail labels
        
        # COMMON_TAIL_LABELS = TAIL_CATS_SCANNET_200+COMMON_CATS_SCANNET_200
        # TAIL_LABELS = TAIL_CATS_SCANNET_200
        unknown_class_labels = UNKNOWN_CLASSES_LABELS[split][task]
        #get their ids
        try:
            # if task == "task1":
            #     if id2label[inst.label_id] in COMMON_TAIL_LABELS:
            #         inst.label_id = 3000
            # elif task == "task2":
            #     if id2label[inst.label_id] in TAIL_LABELS:
            #         inst.label_id = 3000
            if unknown_class_labels is not None:
                if id2label[inst.label_id] in unknown_class_labels:
                    inst.label_id = 3000
        except:
            pass
            # print(inst.label_id)
        #change them to 3000
        
        if inst.label_id in class_ids:
            #assign 3000 uknown
            instances[id2label[inst.label_id]].append(inst.to_dict())
            # if id2label[inst.label_id]=='unknown':
            #     print(inst.vert_count)
            
    return instances

ow_benchmark/util_3d.py
aminebdj/3D-OWIS
# Copyright (c) Facebook, Inc. and its affiliates.

''' Modified based on: https://github.com/erikwijmans/Pointnet2_PyTorch '''
from __future__ import (
    division,
    absolute_import,
    with_statement,
    print_function,
    unicode_literals,
)
import torch
from torch.autograd import Function
import torch.nn as nn
import third_party.pointnet2.pytorch_utils as pt_utils
import sys

try:
    import builtins
except:
    import __builtin__ as builtins

try:
    import pointnet2._ext as _ext
except ImportError:
    if not getattr(builtins, "__POINTNET2_SETUP__", False):
        raise ImportError(
            "Could not import _ext module.\n"
            "Please see the setup instructions in the README: "
            "https://github.com/erikwijmans/Pointnet2_PyTorch/blob/master/README.rst"
        )

if False:
    # Workaround for type hints without depending on the `typing` module
    from typing import *


class RandomDropout(nn.Module):
    def __init__(self, p=0.5, inplace=False):
        super(RandomDropout, self).__init__()
        self.p = p
        self.inplace = inplace

    def forward(self, X):
        theta = torch.Tensor(1).uniform_(0, self.p)[0]
        return pt_utils.feature_dropout_no_scaling(X, theta, self.train, self.inplace)


class FurthestPointSampling(Function):
    @staticmethod
    def forward(ctx, xyz, npoint):
        # type: (Any, torch.Tensor, int) -> torch.Tensor
        r"""
        Uses iterative furthest point sampling to select a set of npoint features that have the largest
        minimum distance

        Parameters
        ----------
        xyz : torch.Tensor
            (B, N, 3) tensor where N > npoint
        npoint : int32
            number of features in the sampled set

        Returns
        -------
        torch.Tensor
            (B, npoint) tensor containing the set
        """
        fps_inds = _ext.furthest_point_sampling(xyz, npoint)
        ctx.mark_non_differentiable(fps_inds)
        return fps_inds

    @staticmethod
    def backward(xyz, a=None):
        return None, None


furthest_point_sample = FurthestPointSampling.apply


class GatherOperation(Function):
    @staticmethod
    def forward(ctx, features, idx):
        # type: (Any, torch.Tensor, torch.Tensor) -> torch.Tensor
        r"""

        Parameters
        ----------
        features : torch.Tensor
            (B, C, N) tensor

        idx : torch.Tensor
            (B, npoint) tensor of the features to gather

        Returns
        -------
        torch.Tensor
            (B, C, npoint) tensor
        """

        _, C, N = features.size()

        ctx.for_backwards = (idx, C, N)

        return _ext.gather_points(features, idx)

    @staticmethod
    def backward(ctx, grad_out):
        idx, C, N = ctx.for_backwards

        grad_features = _ext.gather_points_grad(grad_out.contiguous(), idx, N)
        return grad_features, None


gather_operation = GatherOperation.apply


class ThreeNN(Function):
    @staticmethod
    def forward(ctx, unknown, known):
        # type: (Any, torch.Tensor, torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]
        r"""
            Find the three nearest neighbors of unknown in known
        Parameters
        ----------
        unknown : torch.Tensor
            (B, n, 3) tensor of known features
        known : torch.Tensor
            (B, m, 3) tensor of unknown features

        Returns
        -------
        dist : torch.Tensor
            (B, n, 3) l2 distance to the three nearest neighbors
        idx : torch.Tensor
            (B, n, 3) index of 3 nearest neighbors
        """
        dist2, idx = _ext.three_nn(unknown, known)

        return torch.sqrt(dist2), idx

    @staticmethod
    def backward(ctx, a=None, b=None):
        return None, None


three_nn = ThreeNN.apply


class ThreeInterpolate(Function):
    @staticmethod
    def forward(ctx, features, idx, weight):
        # type(Any, torch.Tensor, torch.Tensor, torch.Tensor) -> Torch.Tensor
        r"""
            Performs weight linear interpolation on 3 features
        Parameters
        ----------
        features : torch.Tensor
            (B, c, m) Features descriptors to be interpolated from
        idx : torch.Tensor
            (B, n, 3) three nearest neighbors of the target features in features
        weight : torch.Tensor
            (B, n, 3) weights

        Returns
        -------
        torch.Tensor
            (B, c, n) tensor of the interpolated features
        """
        B, c, m = features.size()
        n = idx.size(1)

        ctx.three_interpolate_for_backward = (idx, weight, m)

        return _ext.three_interpolate(features, idx, weight)

    @staticmethod
    def backward(ctx, grad_out):
        # type: (Any, torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]
        r"""
        Parameters
        ----------
        grad_out : torch.Tensor
            (B, c, n) tensor with gradients of ouputs

        Returns
        -------
        grad_features : torch.Tensor
            (B, c, m) tensor with gradients of features

        None

        None
        """
        idx, weight, m = ctx.three_interpolate_for_backward

        grad_features = _ext.three_interpolate_grad(
            grad_out.contiguous(), idx, weight, m
        )

        return grad_features, None, None


three_interpolate = ThreeInterpolate.apply


class GroupingOperation(Function):
    @staticmethod
    def forward(ctx, features, idx):
        # type: (Any, torch.Tensor, torch.Tensor) -> torch.Tensor
        r"""

        Parameters
        ----------
        features : torch.Tensor
            (B, C, N) tensor of features to group
        idx : torch.Tensor
            (B, npoint, nsample) tensor containing the indicies of features to group with

        Returns
        -------
        torch.Tensor
            (B, C, npoint, nsample) tensor
        """
        B, nfeatures, nsample = idx.size()
        _, C, N = features.size()

        ctx.for_backwards = (idx, N)

        return _ext.group_points(features, idx)

    @staticmethod
    def backward(ctx, grad_out):
        # type: (Any, torch.tensor) -> Tuple[torch.Tensor, torch.Tensor]
        r"""

        Parameters
        ----------
        grad_out : torch.Tensor
            (B, C, npoint, nsample) tensor of the gradients of the output from forward

        Returns
        -------
        torch.Tensor
            (B, C, N) gradient of the features
        None
        """
        idx, N = ctx.for_backwards

        grad_features = _ext.group_points_grad(grad_out.contiguous(), idx, N)

        return grad_features, None


grouping_operation = GroupingOperation.apply


class BallQuery(Function):
    @staticmethod
    def forward(ctx, radius, nsample, xyz, new_xyz):
        # type: (Any, float, int, torch.Tensor, torch.Tensor) -> torch.Tensor
        r"""

        Parameters
        ----------
        radius : float
            radius of the balls
        nsample : int
            maximum number of features in the balls
        xyz : torch.Tensor
            (B, N, 3) xyz coordinates of the features
        new_xyz : torch.Tensor
            (B, npoint, 3) centers of the ball query

        Returns
        -------
        torch.Tensor
            (B, npoint, nsample) tensor with the indicies of the features that form the query balls
        """
        inds = _ext.ball_query(new_xyz, xyz, radius, nsample)
        ctx.mark_non_differentiable(inds)
        return inds

    @staticmethod
    def backward(ctx, a=None):
        return None, None, None, None


ball_query = BallQuery.apply


class QueryAndGroup(nn.Module):
    r"""
    Groups with a ball query of radius

    Parameters
    ---------
    radius : float32
        Radius of ball
    nsample : int32
        Maximum number of features to gather in the ball
    """

    def __init__(self, radius, nsample, use_xyz=True, ret_grouped_xyz=False, normalize_xyz=False, sample_uniformly=False, ret_unique_cnt=False):
        # type: (QueryAndGroup, float, int, bool) -> None
        super(QueryAndGroup, self).__init__()
        self.radius, self.nsample, self.use_xyz = radius, nsample, use_xyz
        self.ret_grouped_xyz = ret_grouped_xyz
        self.normalize_xyz = normalize_xyz
        self.sample_uniformly = sample_uniformly
        self.ret_unique_cnt = ret_unique_cnt
        if self.ret_unique_cnt:
            assert(self.sample_uniformly)

    def forward(self, xyz, new_xyz, features=None):
        # type: (QueryAndGroup, torch.Tensor. torch.Tensor, torch.Tensor) -> Tuple[Torch.Tensor]
        r"""
        Parameters
        ----------
        xyz : torch.Tensor
            xyz coordinates of the features (B, N, 3)
        new_xyz : torch.Tensor
            centriods (B, npoint, 3)
        features : torch.Tensor
            Descriptors of the features (B, C, N)

        Returns
        -------
        new_features : torch.Tensor
            (B, 3 + C, npoint, nsample) tensor
        """
        idx = ball_query(self.radius, self.nsample, xyz, new_xyz)

        if self.sample_uniformly:
            unique_cnt = torch.zeros((idx.shape[0], idx.shape[1]))
            for i_batch in range(idx.shape[0]):
                for i_region in range(idx.shape[1]):
                    unique_ind = torch.unique(idx[i_batch, i_region, :])
                    num_unique = unique_ind.shape[0]
                    unique_cnt[i_batch, i_region] = num_unique
                    sample_ind = torch.randint(0, num_unique, (self.nsample - num_unique,), dtype=torch.long)
                    all_ind = torch.cat((unique_ind, unique_ind[sample_ind]))
                    idx[i_batch, i_region, :] = all_ind


        xyz_trans = xyz.transpose(1, 2).contiguous()
        grouped_xyz = grouping_operation(xyz_trans, idx)  # (B, 3, npoint, nsample)
        grouped_xyz -= new_xyz.transpose(1, 2).unsqueeze(-1)
        if self.normalize_xyz:
            grouped_xyz /= self.radius

        if features is not None:
            grouped_features = grouping_operation(features, idx)
            if self.use_xyz:
                new_features = torch.cat(
                    [grouped_xyz, grouped_features], dim=1
                )  # (B, C + 3, npoint, nsample)
            else:
                new_features = grouped_features
        else:
            assert (
                self.use_xyz
            ), "Cannot have not features and not use xyz as a feature!"
            new_features = grouped_xyz

        ret = [new_features]
        if self.ret_grouped_xyz:
            ret.append(grouped_xyz)
        if self.ret_unique_cnt:
            ret.append(unique_cnt)
        if len(ret) == 1:
            return ret[0]
        else:
            return tuple(ret)


class GroupAll(nn.Module):
    r"""
    Groups all features

    Parameters
    ---------
    """

    def __init__(self, use_xyz=True, ret_grouped_xyz=False):
        # type: (GroupAll, bool) -> None
        super(GroupAll, self).__init__()
        self.use_xyz = use_xyz

    def forward(self, xyz, new_xyz, features=None):
        # type: (GroupAll, torch.Tensor, torch.Tensor, torch.Tensor) -> Tuple[torch.Tensor]
        r"""
        Parameters
        ----------
        xyz : torch.Tensor
            xyz coordinates of the features (B, N, 3)
        new_xyz : torch.Tensor
            Ignored
        features : torch.Tensor
            Descriptors of the features (B, C, N)

        Returns
        -------
        new_features : torch.Tensor
            (B, C + 3, 1, N) tensor
        """

        grouped_xyz = xyz.transpose(1, 2).unsqueeze(2)
        if features is not None:
            grouped_features = features.unsqueeze(2)
            if self.use_xyz:
                new_features = torch.cat(
                    [grouped_xyz, grouped_features], dim=1
                )  # (B, 3 + C, 1, N)
            else:
                new_features = grouped_features
        else:
            new_features = grouped_xyz

        if self.ret_grouped_xyz:
            return new_features, grouped_xyz
        else:
            return new_features

third_party/pointnet2/pointnet2_utils.py
aminebdj/3D-OWIS
import gc
from contextlib import nullcontext
from pathlib import Path
import statistics
import shutil
import os
import math 
import pyviz3d.visualizer as vis
import matplotlib
from benchmark.evaluate_semantic_instance import evaluate
from collections import defaultdict
from sklearn.cluster import DBSCAN
from utils.votenet_utils.eval_det import eval_det
from torch_scatter import scatter_mean
from datasets.scannet200.scannet200_splits import HEAD_CATS_SCANNET_200, TAIL_CATS_SCANNET_200, COMMON_CATS_SCANNET_200, VALID_CLASS_IDS_200_VALIDATION, CLASS_LABELS_200_VALIDATION
from datasets.scannet200.owis_splits import UNKNOWN_CLASSES_IDS, PREV_KNOWN_CLASSES_IDS, KNOWN_CLASSES_LABELS,KNOWN_CLASSES_IDS
from omegaconf import OmegaConf,open_dict
import hydra
import MinkowskiEngine as ME
import numpy as np
import pytorch_lightning as pl
import torch
from models.metrics import IoU
import random
import colorsys
from typing import List, Tuple
import functools
##########################
import json
import matplotlib.pyplot as plt
from plyfile import PlyData, PlyElement
from omegaconf import OmegaConf,open_dict
from datasets.scannet200.scannet200_constants import CLASS_LABELS_200
from datasets.scannet200.scannet200_splits import HEAD_CATS_SCANNET_200, COMMON_CATS_SCANNET_200, TAIL_CATS_SCANNET_200
from datasets.scannet200.scannet200_splits import HEAD_PRED_IDS, COMMON_PRED_IDS, TAIL_PRED_IDS
from datasets.scannet200.scannet200_constants import VALID_CLASS_IDS_200
from reliability.Fitters import Fit_Weibull_3P
from torch.distributions.weibull import Weibull
from torch.distributions.transforms import AffineTransform
from torch.distributions.transformed_distribution import TransformedDistribution
import shortuuid
import shutil
from collections import deque
from sklearn.decomposition import PCA
import yaml
from yaml.loader import SafeLoader
import time

#tSNE plot
import sklearn
from sklearn.manifold import TSNE 
import seaborn as sb
import matplotlib.pyplot as plt

@functools.lru_cache(20)
def get_evenly_distributed_colors(count: int) -> List[Tuple[np.uint8, np.uint8, np.uint8]]:
    # lru cache caches color tuples
    HSV_tuples = [(x/count, 1.0, 1.0) for x in range(count)]
    random.shuffle(HSV_tuples)
    return list(map(lambda x: (np.array(colorsys.hsv_to_rgb(*x))*255).astype(np.uint8), HSV_tuples))

class RegularCheckpointing(pl.Callback):
    def on_train_epoch_end(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule"):
        general = pl_module.config.general
        trainer.save_checkpoint(f"{general.save_dir}/last-epoch.ckpt")
        print("Checkpoint created")

class InstanceSegmentation(pl.LightningModule):
    def __init__(self, config):
        super().__init__()
        
        global HEAD_PRED_IDS
        global COMMON_PRED_IDS
        global TAIL_PRED_IDS
        
        global PREV_KNOWN_CLASSES_IDS
        global UNKNOWN_CLASSES_IDS
        global KNOWN_CLASSES_LABELS
        
        
        self.decoder_id = config.general.decoder_id

        if config.model.train_on_segments:
            self.mask_type = "segment_mask"
        else:
            self.mask_type = "masks"

        self.eval_on_segments = config.general.eval_on_segments

        self.config = config
        self.save_hyperparameters()
        # model
        self.model = hydra.utils.instantiate(config.model)
        self.optional_freeze = nullcontext
        if config.general.freeze_backbone:
            self.optional_freeze = torch.no_grad
        # loss
        self.ignore_label = config.data.ignore_label
        
        #known classes for each task
        if self.config.general.OW_task == "task1":
            self.num_seen_classes = len(HEAD_CATS_SCANNET_200)
            self.KNOWN_CLASSES_IDS = HEAD_PRED_IDS
        elif self.config.general.OW_task == "task2":
            self.num_seen_classes = len(HEAD_CATS_SCANNET_200+COMMON_CATS_SCANNET_200)
            self.KNOWN_CLASSES_IDS = HEAD_PRED_IDS+COMMON_PRED_IDS
        else:
            self.num_seen_classes = 198
            self.KNOWN_CLASSES_IDS = HEAD_PRED_IDS+COMMON_PRED_IDS+TAIL_PRED_IDS
            
        #Continual learning
        self.classes_dict = {} #to store the number of exmplars per class
            
        
        matcher = hydra.utils.instantiate(config.matcher)
        if self.config.general.use_obj_loss:
            weight_dict = {"loss_ce": matcher.cost_class,
                        "loss_mask": matcher.cost_mask,
                        "loss_dice": matcher.cost_dice,
                        "c_loss": self.config.general.c_loss,
                        "_loss_obj": self.config.general.obj_loss}
        else:
            weight_dict = {"loss_ce": matcher.cost_class,
                        "loss_mask": matcher.cost_mask,
                        "loss_dice": matcher.cost_dice,
                        "c_loss": self.config.general.c_loss}

        aux_weight_dict = {}
        for i in range(self.model.num_levels * self.model.num_decoders):
            if i not in self.config.general.ignore_mask_idx:
                aux_weight_dict.update({k + f"_{i}": v for k, v in weight_dict.items()})
            else:
                aux_weight_dict.update({k + f"_{i}": 0. for k, v in weight_dict.items()})
        weight_dict.update(aux_weight_dict)

        self.preds = dict()
        self.bbox_preds = dict()
        self.bbox_gt = dict()

        self.criterion = hydra.utils.instantiate(config.loss, matcher=matcher, weight_dict=weight_dict)

        # metrics
        self.confusion = hydra.utils.instantiate(config.metrics)
        self.iou = IoU()
        # misc
        self.labels_info = dict()
        
        self.train_oracle = self.config.general.train_oracle
        
        # autolabeling
        self.topk = self.config.general.topk
        self.ukn_cls = 200
        
        #unknown distribution
        self.unknown_class_distribution_is_set = False
        
        
    def forward(self, x, point2segment=None, raw_coordinates=None, is_eval=False):
        with self.optional_freeze():
            x = self.model(x, point2segment, raw_coordinates=raw_coordinates,
                           is_eval=is_eval)
        return x

    def training_step(self, batch, batch_idx):
        
        data, target, file_names = batch 
        
        if self.config.general.use_examplars_in_train:
            target = self.task(target,eval = False ,task_ow = self.config.general.OW_task, split=self.config.general.split, File_names = file_names)
        else: 
            target = self.task(target,eval = False ,task_ow = self.config.general.OW_task, split=self.config.general.split, File_names = None)
               
        if data.features.shape[0] > self.config.general.max_batch_size:
            print("data exceeds threshold")
            raise RuntimeError("BATCH TOO BIG")

        if len(target) == 0:
            print("no targets")
            return None

        raw_coordinates = None
        if self.config.data.add_raw_coordinates:
            raw_coordinates = data.features[:, -3:]
            data.features = data.features[:, :-3]

        data = ME.SparseTensor(coordinates=data.coordinates,
                              features=data.features,
                              device=self.device)

        try:
            output = self.forward(data,
                                  point2segment=[target[i]['point2segment'] for i in range(len(target))],
                                  raw_coordinates=raw_coordinates)
        except RuntimeError as run_err:
            print(run_err)
            if 'only a single point gives nans in cross-attention' == run_err.args[0]:
                return None
            else:
                raise run_err

        try:
            if self.current_epoch >= self.config.general.start_auto_labeling and (not self.config.general.train_oracle):
                for b in range(len(target)):
                    try:
                        target = self.Auto_Labeling(b, output, target, self.topk,raw_coordinates, point2segment= [target[i]['point2segment'] for i in range(len(target))])
                    except:
                        pass
            
            losses = self.criterion(output, target, mask_type=self.mask_type, iteration = self.global_step)

            if self.config.general.learn_energy_trainig_dataset and (self.current_epoch >= self.config.general.WARM_UP_EPOCH):
                
                file_path_p = os.path.join(self.config.general.save_energy_training_dataset_in,"train_set/")
                if not os.path.exists(file_path_p):
                    os.makedirs(file_path_p) 
                
                indices = self.criterion.get_indices(output, target,
                                        mask_type=self.mask_type)
                for batch_id, (map_id, target_id) in enumerate(indices):
                    unkn_mask = target[batch_id]['labels'][target_id]==self.ukn_cls
                    pred_logits_unk = output['pred_logits'][batch_id][map_id][unkn_mask] if batch_id == 0 else torch.cat([pred_logits_unk, output['pred_logits'][batch_id][map_id][unkn_mask]])
                    pred_logits_kn = output['pred_logits'][batch_id][map_id][~unkn_mask] if batch_id == 0 else torch.cat([pred_logits_kn, output['pred_logits'][batch_id][map_id][~unkn_mask]])

                
                data_ = pred_logits_kn, pred_logits_unk           
                file_path = file_path_p+"/logits_temp/"+ shortuuid.uuid() + '.pkl'
                if not os.path.exists(file_path_p+"/logits_temp/"):
                    os.makedirs(file_path_p+"/logits_temp/")
                torch.save(data_, file_path)
                 
                
                
        except ValueError as val_err:
            print(f"ValueError: {val_err}")
            print(f"data shape: {data.shape}")
            print(f"data feat shape:  {data.features.shape}")
            print(f"data feat nans:   {data.features.isnan().sum()}")
            print(f"output: {output}")
            print(f"target: {target}")
            print(f"filenames: {file_names}")
            raise val_err

        for k in list(losses.keys()):
            if k in self.criterion.weight_dict:
                losses[k] *= self.criterion.weight_dict[k]
            else:
                # remove this loss if not specified in `weight_dict`
                losses.pop(k)

        logs = {f"train_{k}": v.detach().cpu().item() for k,v in losses.items()}

        logs['train_mean_loss_ce'+'_'+self.config.general.logg_suffix ] = statistics.mean([item for item in [v for k, v in logs.items() if "loss_ce" in k]])

        logs['train_mean_loss_mask'+'_'+self.config.general.logg_suffix ] = statistics.mean(
            [item for item in [v for k, v in logs.items() if "loss_mask" in k]])

        logs['train_mean_loss_dice'+'_'+self.config.general.logg_suffix ] = statistics.mean(
            [item for item in [v for k, v in logs.items() if "loss_dice" in k]])

        self.log_dict(logs)
        return sum(losses.values())
    
    def validation_step(self, batch, batch_idx):
        return self.eval_step(batch, batch_idx)

    def export(self, pred_masks, scores, pred_classes, file_names, decoder_id):
        root_path = f"eval_output"
        base_path = f"{root_path}/instance_evaluation_{self.config.general.experiment_name}_{self.current_epoch}/decoder_{decoder_id}"
        pred_mask_path = f"{base_path}/pred_mask"

        Path(pred_mask_path).mkdir(parents=True, exist_ok=True)

        file_name = file_names
        with open(f"{base_path}/{file_name}.txt", "w") as fout:
            real_id = -1
            for instance_id in range(len(pred_classes)):
                real_id += 1
                pred_class = pred_classes[instance_id]
                score = scores[instance_id]
                mask = pred_masks[:, instance_id].astype("uint8")

                if score > self.config.general.export_threshold:
                    # reduce the export size a bit. I guess no performance difference
                    np.savetxt(f"{pred_mask_path}/{file_name}_{real_id}.txt", mask, fmt="%d")
                    fout.write(f"pred_mask/{file_name}_{real_id}.txt {pred_class} {score}\n")

    def training_epoch_end(self, outputs):
        
        #Coninual learning
        self.classes_dict = {} #reinitilize the dictionary after every epoch
        train_loss = sum([out["loss"].cpu().item() for out in outputs]) / len(outputs)
        results = {"train_loss_mean"+"_"+self.config.general.logg_suffix : train_loss}
        self.log_dict(results)

        if self.config.general.learn_energy_trainig_dataset and (self.current_epoch >= self.config.general.WARM_UP_EPOCH):
            
            file_path_p = os.path.join(self.config.general.save_energy_training_dataset_in,"train_set/")
            temp_file_path = os.path.join(file_path_p, "logits_temp")
            new_file_path = os.path.join(file_path_p, "logits")
            for file in os.listdir(self.config.general.save_dir):
                file_lst = file.split("_")
                
                if "ap" in file_lst:
                    os.replace(temp_file_path, new_file_path)
                    self.learn_energy()
                    os.replace(os.path.join(self.config.general.save_dir, file), os.path.join(self.config.general.save_dir, "best.ckpt"))
                    
                    
            if os.path.exists(temp_file_path):    
                shutil.rmtree(temp_file_path)
                

    def validation_epoch_end(self, outputs):
        self.test_epoch_end(outputs)

    def save_visualizations(self, target_full, full_res_coords,
                            sorted_masks, sort_classes, file_name, original_colors, original_normals,
                            sort_scores_values, point_size=20, sorted_heatmaps=None,
                            query_pos=None, backbone_features=None):

        full_res_coords -= full_res_coords.mean(axis=0)

        gt_pcd_pos = []
        gt_pcd_normals = []
        gt_pcd_color = []
        gt_inst_pcd_color = []
        gt_boxes = []

        if 'labels' in target_full:
            instances_colors = torch.from_numpy(
                np.vstack(get_evenly_distributed_colors(target_full['labels'].shape[0])))
            
            unkn_gt_exists = 0
            for instance_counter, (label, mask) in enumerate(zip(target_full['labels'], target_full['masks'])):
                if label == 255:
                    continue
                if label == 3000:
                    unkn_gt_exists = 1
                    
                mask_tmp = mask.detach().cpu().numpy()
                mask_coords = full_res_coords[mask_tmp.astype(bool), :]

                if len(mask_coords) == 0:
                    continue

                gt_pcd_pos.append(mask_coords)
                mask_coords_min = full_res_coords[mask_tmp.astype(bool), :].min(axis=0)
                mask_coords_max = full_res_coords[mask_tmp.astype(bool), :].max(axis=0)
                size = mask_coords_max - mask_coords_min
                mask_coords_middle = mask_coords_min + size / 2

                gt_boxes.append({"position": mask_coords_middle, "size": size,
                                 "color": self.validation_dataset.map2color([label])[0]})

                gt_pcd_color.append(
                    self.validation_dataset.map2color([label]).repeat(gt_pcd_pos[-1].shape[0], 1)
                )
                gt_inst_pcd_color.append(instances_colors[instance_counter % len(instances_colors)].unsqueeze(0).repeat(gt_pcd_pos[-1].shape[0], 1))

                gt_pcd_normals.append(original_normals[mask_tmp.astype(bool), :])
                
                
            
            gt_pcd_pos = np.concatenate(gt_pcd_pos)
            gt_pcd_normals = np.concatenate(gt_pcd_normals)
            gt_pcd_color = np.concatenate(gt_pcd_color)
            gt_inst_pcd_color = np.concatenate(gt_inst_pcd_color)
        
        v = vis.Visualizer()
        
        v.add_points("RGB Input", full_res_coords,
                     colors=original_colors,
                     normals=original_normals,
                     visible=True,
                     point_size=point_size)
        
  
        if backbone_features is not None:
            v.add_points("PCA", full_res_coords,
                        colors=backbone_features,
                        normals=original_normals,
                        visible=False,
                        point_size=point_size)

        if 'labels' in target_full:
            v.add_points("Semantics (GT)", gt_pcd_pos,
                        colors=gt_pcd_color,
                        normals=gt_pcd_normals,
                        alpha=0.8,
                        visible=False,
                        point_size=point_size)
            v.add_points("Instances (GT)", gt_pcd_pos,
                        colors=gt_inst_pcd_color,
                        normals=gt_pcd_normals,
                        alpha=0.8,
                        visible=False,
                        point_size=point_size)
        
        pred_coords = []
        pred_normals = []
        pred_sem_color = []
        pred_inst_color = []

        for did in range(len(sorted_masks)):
            instances_colors = torch.from_numpy(
                np.vstack(get_evenly_distributed_colors(max(1, sorted_masks[did].shape[1]))))
            ukn_pred_exists = 0
            for i in reversed(range(sorted_masks[did].shape[1])):
                coords = full_res_coords[sorted_masks[did][:, i].astype(bool), :]

                mask_coords = full_res_coords[sorted_masks[did][:,i].astype(bool), :]
                mask_normals = original_normals[sorted_masks[did][:,i].astype(bool), :]

                label = sort_classes[did][i]
            
                if label==3000:
                    if len(mask_coords) == 0:
                        continue
                    ukn_pred_exists += 1
                    
                    pred_coords.append(mask_coords)
                    pred_normals.append(mask_normals)

                    pred_sem_color.append(
                        self.validation_dataset.map2color([label]).repeat(
                            mask_coords.shape[0], 1)
                    )

                    pred_inst_color.append(instances_colors[i % len(instances_colors)].unsqueeze(0).repeat(
                        mask_coords.shape[0], 1)
                    )

            if len(pred_coords) > 0:
                pred_coords = np.concatenate(pred_coords)
                pred_normals = np.concatenate(pred_normals)
                pred_sem_color = np.concatenate(pred_sem_color)
                pred_inst_color = np.concatenate(pred_inst_color)

                v.add_points("Semantics (Mask3D)", pred_coords,
                             colors=pred_sem_color,
                             normals=pred_normals,
                             visible=False,
                             alpha=0.8,
                             point_size=point_size)
                v.add_points("Instances (Mask3D)", pred_coords,
                             colors=pred_inst_color,
                             normals=pred_normals,
                             visible=False,
                             alpha=0.8,
                             point_size=point_size)
                
        if ukn_pred_exists and unkn_gt_exists:
            v.save(f"{self.config['general']['save_dir']}/visualizations/{file_name}")
        
    def eval_step(self, batch, batch_idx):
        
        
        data, target, file_names = batch
        inverse_maps = data.inverse_maps
        target_full = data.target_full
        original_colors = data.original_colors
        data_idx = data.idx
        original_normals = data.original_normals
        original_coordinates = data.original_coordinates
        
        if self.config.general.correct_unknown_cls_prob and not self.unknown_class_distribution_is_set:
            self.set_unkn_distribution(target[0]['point2segment'].device)
            
        if len(data.coordinates) == 0:
            return 0.

        raw_coordinates = None
        if self.config.data.add_raw_coordinates:
            raw_coordinates = data.features[:, -3:]
            data.features = data.features[:, :-3]

        if raw_coordinates.shape[0] == 0:
            return 0.

        data = ME.SparseTensor(coordinates=data.coordinates, features=data.features, device=self.device)


        try:
            
            output = self.forward(data,
                                  point2segment=[target[i]['point2segment'] for i in range(len(target))],
                                  raw_coordinates=raw_coordinates,
                                  is_eval=True)
            
                
        except RuntimeError as run_err:
            print(run_err)
            if 'only a single point gives nans in cross-attention' == run_err.args[0]:
                return None
            else:
                raise run_err
            

        if self.config.data.test_mode != "test":
            if self.config.trainer.deterministic:
                torch.use_deterministic_algorithms(False)

            try:
                target = self.task(target,eval = True, task_ow=self.config.general.OW_task, split=self.config.general.split)
                target_full = self.task(target_full,eval = True, task_ow=self.config.general.OW_task, split=self.config.general.split) # target_full is used instead of target when saving visualizations

                losses = self.criterion(output, target, mask_type=self.mask_type, iteration = 0)

                if self.config.general.save_KN_UKN_tSNE:
                    
                    file_path_p = self.config.general.save_features_in+"val_set/"
                    indices = self.criterion.get_indices(output, target,
                                            mask_type=self.mask_type)
                    for batch_id, (map_id, target_id) in enumerate(indices):
                        unkn_mask = target[batch_id]['labels'][target_id]==self.ukn_cls
                        queries = output['refin_queries'][batch_id][map_id] if batch_id == 0 else torch.cat([queries, output['refin_queries'][batch_id][map_id]])
                        labels = target[batch_id]['labels'][target_id] if batch_id == 0 else torch.cat([labels, target[batch_id]['labels'][target_id]])
                        
                    if not os.path.exists(file_path_p+"known_unknown_queries/"):
                        os.makedirs(file_path_p+"known_unknown_queries/")
                    file_path = file_path_p+"known_unknown_queries/"+ shortuuid.uuid() + '.pkl'
                    data_ = (queries, labels)
                    torch.save(data_, file_path)
                    
                    
            except ValueError as val_err:
                print(f"ValueError: {val_err}")
                print(f"data shape: {data.shape}")
                print(f"data feat shape:  {data.features.shape}")
                print(f"data feat nans:   {data.features.isnan().sum()}")
                print(f"output: {output}")
                print(f"target: {target}")
                print(f"filenames: {file_names}")
                raise val_err

            for k in list(losses.keys()):
                if k in self.criterion.weight_dict:
                    losses[k] *= self.criterion.weight_dict[k]
                else:
                    # remove this loss if not specified in `weight_dict`
                    losses.pop(k)
            if self.config.trainer.deterministic:
                torch.use_deterministic_algorithms(True)
#######################################################################################
        if self.config.general.save_visualizations:
            backbone_features = output['backbone_features'].F.detach().cpu().numpy()
            from sklearn import decomposition
            pca = decomposition.PCA(n_components=3)
            pca.fit(backbone_features)
            pca_features = pca.transform(backbone_features)
            rescaled_pca = 255 * (pca_features - pca_features.min()) / (pca_features.max() - pca_features.min())

        self.eval_instance_step(output, target, target_full, inverse_maps, file_names, original_coordinates,
                                original_colors, original_normals, raw_coordinates, data_idx,
                                backbone_features=rescaled_pca if self.config.general.save_visualizations else None)
#######################################################################################
        if self.config.data.test_mode != "test":
            return {f"val_{k}": v.detach().cpu().item() for k, v in losses.items()}
        else:
            return 0.
#######################################################################################
    def test_step(self, batch, batch_idx):
        return self.eval_step(batch, batch_idx)

    def get_full_res_mask(self, mask, inverse_map, point2segment_full, is_heatmap=False):
        mask = mask.detach().cpu()[inverse_map]  # full res

        if self.eval_on_segments and is_heatmap==False:
            mask = scatter_mean(mask, point2segment_full, dim=0)  # full res segments
            mask = (mask > 0.5).float()
            mask = mask.detach().cpu()[point2segment_full.cpu()]  # full res points

        return mask


    def get_mask_and_scores(self, mask_cls, mask_pred, num_queries=100, num_classes=18, device=None):
        if device is None:
            device = self.device
        labels = torch.arange(num_classes, device=device).unsqueeze(0).repeat(num_queries, 1).flatten(0, 1)

        if self.config.general.topk_per_image != -1 :
            scores_per_query, topk_indices = mask_cls.flatten(0, 1).topk(self.config.general.topk_per_image, sorted=True)
        else:
            scores_per_query, topk_indices = mask_cls.flatten(0, 1).topk(num_queries, sorted=True)

        labels_per_query = labels[topk_indices]
        topk_indices = torch.div(topk_indices, num_classes, rounding_mode='floor')
        mask_pred = mask_pred[:, topk_indices]

        result_pred_mask = (mask_pred > 0).float()
        heatmap = mask_pred.float().sigmoid()

        mask_scores_per_image = (heatmap * result_pred_mask).sum(0) / (result_pred_mask.sum(0) + 1e-6)
        score = scores_per_query * mask_scores_per_image
        classes = labels_per_query

        return score, result_pred_mask, classes, heatmap

    def eval_instance_step(self, output, target_low_res, target_full_res, inverse_maps, file_names,
                           full_res_coords, original_colors, original_normals, raw_coords, idx, first_full_res=False,
                           backbone_features=None,):
        label_offset = self.validation_dataset.label_offset
        
        prediction = output['aux_outputs']
        prediction.append({
            'pred_logits': output['pred_logits'],
            'pred_masks': output['pred_masks']
        })
        
        if self.config.general.correct_unknown_cls_prob and self.unknown_class_distribution_is_set:
            try :
                prediction[self.decoder_id]['pred_logits'] = self.correct_classes_prob(output) #correct the classification output during inference
            except:
                prediction[self.decoder_id]['pred_logits'] = torch.functional.F.softmax(
                                                            prediction[self.decoder_id]['pred_logits'],
                                                           dim=-1)[..., :-1] 
        else:
            prediction[self.decoder_id]['pred_logits'] = torch.functional.F.softmax(
                                                            prediction[self.decoder_id]['pred_logits'],
                                                           dim=-1)[..., :-1]     
      
        if self.config.general.learn_energy_trainig_dataset and os.path.exists(os.path.join(self.config.general.save_dir,'energy_dist_' + str(self.num_seen_classes) + '.pkl')):
            param_save_location = os.path.join(self.config.general.save_dir,'energy_dist_' + str(self.num_seen_classes) + '.pkl')
            if os.path.isfile(param_save_location):
                params = torch.load(param_save_location)
                unknown = params[0]
                known = params[1]
                
                self.unk_dist = self.create_distribution(unknown['scale_unk'], unknown['shape_unk'], unknown['shift_unk'])
                self.unk_dist_shift = unknown['shift_unk'] 
                self.known_dist = self.create_distribution(known['scale_known'], known['shape_known'], known['shift_known'])
                self.known_dist_shift = known['shift_known']
                
                self.energy_distribution_loaded = True
                
                try: 
                    prediction[self.decoder_id]['pred_logits'] = self.update_label_based_on_energy(prediction[self.decoder_id]['pred_logits'] )
                except:
                    pass
            else:
                print('Energy distribution is not found at ' + param_save_location)
        
           

        all_pred_classes = list()
        all_pred_masks = list()
        all_pred_scores = list()
        all_heatmaps = list()
        all_query_pos = list()

        offset_coords_idx = 0
        for bid in range(len(prediction[self.decoder_id]['pred_masks'])):
            if not first_full_res:
                if self.model.train_on_segments:
                    masks = prediction[self.decoder_id]['pred_masks'][bid].detach().cpu()[target_low_res[bid]['point2segment'].cpu()]
                else:
                    masks = prediction[self.decoder_id]['pred_masks'][bid].detach().cpu()

                if self.config.general.use_dbscan:
                    new_preds = {
                        'pred_masks': list(),
                        'pred_logits': list(),
                    }

                    curr_coords_idx = masks.shape[0]
                    curr_coords = raw_coords[offset_coords_idx:curr_coords_idx + offset_coords_idx]
                    offset_coords_idx += curr_coords_idx

                    for curr_query in range(masks.shape[1]):
                        curr_masks = masks[:, curr_query] > 0

                        if curr_coords[curr_masks].shape[0] > 0:
                            clusters = DBSCAN(eps=self.config.general.dbscan_eps,
                                              min_samples=self.config.general.dbscan_min_points,
                                              n_jobs=-1).fit(curr_coords[curr_masks]).labels_

                            new_mask = torch.zeros(curr_masks.shape, dtype=int)
                            new_mask[curr_masks] = torch.from_numpy(clusters) + 1

                            for cluster_id in np.unique(clusters):
                                original_pred_masks = masks[:, curr_query]
                                if cluster_id != -1:
                                    new_preds['pred_masks'].append(original_pred_masks * (new_mask == cluster_id + 1))
                                    new_preds['pred_logits'].append(
                                        prediction[self.decoder_id]['pred_logits'][bid, curr_query])

                    scores, masks, classes, heatmap = self.get_mask_and_scores(
                        torch.stack(new_preds['pred_logits']).cpu(),
                        torch.stack(new_preds['pred_masks']).T,
                        len(new_preds['pred_logits']),
                        self.model.num_classes - 1)
                    
                else:
                    scores, masks, classes, heatmap = self.get_mask_and_scores(
                    prediction[self.decoder_id]['pred_logits'][bid].detach().cpu(),
                    masks,
                    prediction[self.decoder_id]['pred_logits'][bid].shape[0],
                    self.model.num_classes - 1)

                masks = self.get_full_res_mask(masks,
                                               inverse_maps[bid],
                                               target_full_res[bid]['point2segment'])

                heatmap = self.get_full_res_mask(heatmap,
                                                 inverse_maps[bid],
                                                 target_full_res[bid]['point2segment'],
                                                 is_heatmap=True)

                if backbone_features is not None:
                    backbone_features = self.get_full_res_mask(torch.from_numpy(backbone_features),
                                                               inverse_maps[bid],
                                                               target_full_res[bid]['point2segment'],
                                                               is_heatmap=True)
                    backbone_features = backbone_features.numpy()
            else:
                assert False,  "not tested"
                masks = self.get_full_res_mask(prediction[self.decoder_id]['pred_masks'][bid].cpu(),
                                               inverse_maps[bid],
                                               target_full_res[bid]['point2segment'])

                scores, masks, classes, heatmap = self.get_mask_and_scores(
                    prediction[self.decoder_id]['pred_logits'][bid].cpu(),
                    masks,
                    prediction[self.decoder_id]['pred_logits'][bid].shape[0],
                    self.model.num_classes - 1,
                    device="cpu")

            masks = masks.numpy()
            heatmap = heatmap.numpy()

            sort_scores = scores.sort(descending=True)
            sort_scores_index = sort_scores.indices.cpu().numpy()
            sort_scores_values = sort_scores.values.cpu().numpy()
            sort_classes = classes[sort_scores_index]

            sorted_masks = masks[:, sort_scores_index]
            sorted_heatmap = heatmap[:, sort_scores_index]
            
            if self.config.general.filter_out_instances:
                keep_instances = set()
                pairwise_overlap = (sorted_masks.T @ sorted_masks)
                normalization = pairwise_overlap.max(axis=0)
                norm_overlaps = pairwise_overlap / normalization

                for instance_id in range(norm_overlaps.shape[0]):
                    # filter out unlikely masks and nearly empty masks
                    # if not(sort_scores_values[instance_id] < 0.3 or sorted_masks[:, instance_id].sum() < 500):
                    if not(sort_scores_values[instance_id] < self.config.general.scores_threshold):
                        # check if mask != empty
                        if not sorted_masks[:, instance_id].sum() == 0.0:
                            overlap_ids = set(np.nonzero(norm_overlaps[instance_id, :] > self.config.general.iou_threshold)[0])

                            if len(overlap_ids) == 0:
                                keep_instances.add(instance_id)
                            else:
                                if instance_id == min(overlap_ids):
                                    keep_instances.add(instance_id)

                keep_instances = sorted(list(keep_instances))
                all_pred_classes.append(sort_classes[keep_instances])
                all_pred_masks.append(sorted_masks[:, keep_instances])
                all_pred_scores.append(sort_scores_values[keep_instances])
                all_heatmaps.append(sorted_heatmap[:, keep_instances])
            else:
                all_pred_classes.append(sort_classes)
                all_pred_masks.append(sorted_masks)
                all_pred_scores.append(sort_scores_values)
                all_heatmaps.append(sorted_heatmap)

        if self.validation_dataset.dataset_name == "scannet200":
            all_pred_classes[bid][all_pred_classes[bid] == 0] = -1
            if self.config.data.test_mode != "test":
                target_full_res[bid]['labels'][target_full_res[bid]['labels'] == 0] = -1

        for bid in range(len(prediction[self.decoder_id]['pred_masks'])):
            all_pred_classes[bid] = self.validation_dataset._remap_model_output(all_pred_classes[bid].cpu() + label_offset) #from 200==>3000

            if self.config.data.test_mode != "test" and len(target_full_res) != 0:
                target_full_res[bid]['labels'] = self.validation_dataset._remap_model_output(
                    target_full_res[bid]['labels'].cpu() + label_offset)

                # PREDICTION BOX
                bbox_data = []
                for query_id in range(all_pred_masks[bid].shape[1]):  # self.model.num_queries
                    obj_coords = full_res_coords[bid][all_pred_masks[bid][:, query_id].astype(bool), :]
                    if obj_coords.shape[0] > 0:
                        obj_center = obj_coords.mean(axis=0)
                        obj_axis_length = obj_coords.max(axis=0) - obj_coords.min(axis=0)

                        bbox = np.concatenate((obj_center, obj_axis_length))

                        bbox_data.append((all_pred_classes[bid][query_id].item(), bbox,
                                          all_pred_scores[bid][query_id]
                        ))
                self.bbox_preds[file_names[bid]] = bbox_data

                # GT BOX
                bbox_data = []
                for obj_id in range(target_full_res[bid]['masks'].shape[0]):
                    if target_full_res[bid]['labels'][obj_id].item() == 255:
                        continue

                    obj_coords = full_res_coords[bid][target_full_res[bid]['masks'][obj_id, :].cpu().detach().numpy().astype(bool), :]
                    if obj_coords.shape[0] > 0:
                        obj_center = obj_coords.mean(axis=0)
                        obj_axis_length = obj_coords.max(axis=0) - obj_coords.min(axis=0)

                        bbox = np.concatenate((obj_center, obj_axis_length))
                        bbox_data.append((target_full_res[bid]['labels'][obj_id].item(), bbox))

                self.bbox_gt[file_names[bid]] = bbox_data

            if self.config.general.eval_inner_core == -1:
                self.preds[file_names[bid]] = {
                    'pred_masks': all_pred_masks[bid],
                    'pred_scores': all_pred_scores[bid],
                    'pred_classes': all_pred_classes[bid]
                }
            else:
                # prev val_dataset
                self.preds[file_names[bid]] = {
                    'pred_masks': all_pred_masks[bid][self.test_dataset.data[idx[bid]]['cond_inner']],
                    'pred_scores': all_pred_scores[bid],
                    'pred_classes': all_pred_classes[bid]
                }

            if self.config.general.save_visualizations:
                if 'cond_inner' in self.test_dataset.data[idx[bid]]:
                    target_full_res[bid]['masks'] = target_full_res[bid]['masks'][:, self.test_dataset.data[idx[bid]]['cond_inner']]
                    self.save_visualizations(target_full_res[bid],
                                             full_res_coords[bid][self.test_dataset.data[idx[bid]]['cond_inner']],
                                             [self.preds[file_names[bid]]['pred_masks']],
                                             [self.preds[file_names[bid]]['pred_classes']],
                                             file_names[bid],
                                             original_colors[bid][self.test_dataset.data[idx[bid]]['cond_inner']],
                                             original_normals[bid][self.test_dataset.data[idx[bid]]['cond_inner']],
                                             [self.preds[file_names[bid]]['pred_scores']],
                                             sorted_heatmaps=[all_heatmaps[bid][self.test_dataset.data[idx[bid]]['cond_inner']]],
                                             query_pos=all_query_pos[bid][self.test_dataset.data[idx[bid]]['cond_inner']] if len(all_query_pos) > 0 else None,
                                             backbone_features=backbone_features[self.test_dataset.data[idx[bid]]['cond_inner']],
                                             point_size=self.config.general.visualization_point_size)
                else:
                    self.save_visualizations(target_full_res[bid],
                                             full_res_coords[bid],
                                             [self.preds[file_names[bid]]['pred_masks']],
                                             [self.preds[file_names[bid]]['pred_classes']],
                                             file_names[bid],
                                             original_colors[bid],
                                             original_normals[bid],
                                             [self.preds[file_names[bid]]['pred_scores']],
                                             sorted_heatmaps=[all_heatmaps[bid]],
                                             query_pos=all_query_pos[bid] if len(all_query_pos) > 0 else None,
                                             backbone_features=backbone_features,
                                             point_size=self.config.general.visualization_point_size)

            if self.config.general.export:
                if self.validation_dataset.dataset_name == "stpls3d":
                    scan_id, _, _, crop_id = file_names[bid].split("_")
                    crop_id = int(crop_id.replace(".txt", ""))
                    file_name = f"{scan_id}_points_GTv3_0{crop_id}_inst_nostuff"

                    self.export(
                        self.preds[file_names[bid]]['pred_masks'],
                        self.preds[file_names[bid]]['pred_scores'],
                        self.preds[file_names[bid]]['pred_classes'],
                        file_name,
                        self.decoder_id
                    )
                else:
                    self.export(
                        self.preds[file_names[bid]]['pred_masks'],
                        self.preds[file_names[bid]]['pred_scores'],
                        self.preds[file_names[bid]]['pred_classes'],
                        file_names[bid],
                        self.decoder_id
                    )
#######################################################################################
    def eval_instance_epoch_end(self):
        log_prefix = f"val"
        ap_results = {}
        U_Recall_results = {}
        A_OSE_results = {}
        WI_results = {}
        head_results, tail_results, common_results = [], [], []
        if self.config.general.split == 'A':
            Split_A_T1_results, Split_A_T2_results, Split_A_T3_results = [], [], []
        elif self.config.general.split == 'B':
            Split_B_T1_results, Split_B_T2_results, Split_B_T3_results= [], [], []
        elif self.config.general.split == 'C':
            Split_C_T1_results, Split_C_T2_results, Split_C_T3_results= [], [], []  

        box_ap_50 = eval_det(self.bbox_preds, self.bbox_gt, ovthresh=0.5, use_07_metric=False)
        box_ap_25 = eval_det(self.bbox_preds, self.bbox_gt, ovthresh=0.25, use_07_metric=False)
        mean_box_ap_25 = sum([v for k, v in box_ap_25[-1].items()]) / len(box_ap_25[-1].keys())
        mean_box_ap_50 = sum([v for k, v in box_ap_50[-1].items()]) / len(box_ap_50[-1].keys())

        ap_results[f"{log_prefix}_mean_box_ap_25"] = mean_box_ap_25
        ap_results[f"{log_prefix}_mean_box_ap_50"] = mean_box_ap_50

        for class_id in box_ap_50[-1].keys():
            class_name = self.train_dataset.label_info[class_id]['name']
            ap_results[f"{log_prefix}_{class_name}_val_box_ap_50"] = box_ap_50[-1][class_id]

        for class_id in box_ap_25[-1].keys():
            class_name = self.train_dataset.label_info[class_id]['name']
            ap_results[f"{log_prefix}_{class_name}_val_box_ap_25"] = box_ap_25[-1][class_id]

        root_path = f"eval_output"
        base_path = f"{root_path}/instance_evaluation_{self.config.general.experiment_name}_{self.current_epoch}"

        if self.validation_dataset.dataset_name in ["scannet", "stpls3d", "scannet200"]:
            gt_data_path = f"{self.validation_dataset.data_dir[0]}/instance_gt/{self.validation_dataset.mode}"
        else:
            gt_data_path = f"{self.validation_dataset.data_dir[0]}/instance_gt/Area_{self.config.general.area}"

        pred_path = f"{base_path}/tmp_output.txt"

        log_prefix = f"val"

        if not os.path.exists(base_path):
            os.makedirs(base_path)

        try:
            if self.validation_dataset.dataset_name == "s3dis":
                new_preds = {}
                for key in self.preds.keys():
                    new_preds[key.replace(f"Area_{self.config.general.area}_", "")] = {
                        'pred_classes': self.preds[key]['pred_classes'] + 1,
                        'pred_masks': self.preds[key]['pred_masks'],
                        'pred_scores': self.preds[key]['pred_scores']
                    }
                mprec, mrec = evaluate(self.current_epoch,
                                       self.config.general,
                                       new_preds,
                                       gt_data_path,
                                       pred_path, dataset="s3dis")
                ap_results[f"{log_prefix}_mean_precision"+"_"+self.config.general.logg_suffix ] = mprec
                ap_results[f"{log_prefix}_mean_recall"+"_"+self.config.general.logg_suffix ] = mrec
            elif self.validation_dataset.dataset_name == "stpls3d":
                new_preds = {}
                for key in self.preds.keys():
                    new_preds[key.replace(".txt", "")] = {
                        'pred_classes': self.preds[key]['pred_classes'],
                        'pred_masks': self.preds[key]['pred_masks'],
                        'pred_scores': self.preds[key]['pred_scores']
                    }

                evaluate(self.current_epoch,
                         self.config.general,
                         new_preds, 
                         gt_data_path, 
                         pred_path, 
                         dataset="stpls3d")
            elif self.validation_dataset.dataset_name == "scannet200":
                U_Recall,WI,A_OSE = evaluate(self.current_epoch,
                                             self.config.general,
                            self.preds, 
                            gt_data_path, 
                            pred_path, 
                            dataset=self.validation_dataset.dataset_name)
                
            with open(pred_path, "r") as fin:
                for line_id, line in enumerate(fin):
                    if line_id == 0:
                        # ignore header
                        continue
                    class_name, _, ap, ap_50, ap_25 = line.strip().split(",")

                    if self.validation_dataset.dataset_name == "scannet200":
                        if class_name in VALID_CLASS_IDS_200_VALIDATION:
                            ap_results[f"{log_prefix}_{class_name}_val_ap"+"_"+self.config.general.logg_suffix ] = float(ap)
                            ap_results[f"{log_prefix}_{class_name}_val_ap_50"+"_"+self.config.general.logg_suffix ] = float(ap_50)
                            ap_results[f"{log_prefix}_{class_name}_val_ap_25"+"_"+self.config.general.logg_suffix ] = float(ap_25)

                            if class_name in HEAD_CATS_SCANNET_200:
                                head_results.append(np.array((float(ap), float(ap_50), float(ap_25))))
                            elif class_name in COMMON_CATS_SCANNET_200:
                                common_results.append(np.array((float(ap), float(ap_50), float(ap_25))))
                            elif class_name in TAIL_CATS_SCANNET_200:
                                tail_results.append(np.array((float(ap), float(ap_50), float(ap_25))))
                            
                            if self.config.general.split == 'A':
                                if class_name in KNOWN_CLASSES_LABELS[self.config.general.split]['task1']:
                                    Split_A_T1_results.append(np.array((float(ap), float(ap_50), float(ap_25))))
                                elif class_name in KNOWN_CLASSES_LABELS[self.config.general.split]['task2']:
                                    Split_A_T2_results.append(np.array((float(ap), float(ap_50), float(ap_25))))
                                elif class_name in KNOWN_CLASSES_LABELS[self.config.general.split]['task3']:
                                    Split_A_T3_results.append(np.array((float(ap), float(ap_50), float(ap_25))))
                                    
                            elif self.config.general.split == 'B':
                                if class_name in KNOWN_CLASSES_LABELS[self.config.general.split]['task1']:
                                    Split_B_T1_results.append(np.array((float(ap), float(ap_50), float(ap_25))))
                                elif class_name in KNOWN_CLASSES_LABELS[self.config.general.split]['task2']:
                                    Split_B_T2_results.append(np.array((float(ap), float(ap_50), float(ap_25))))
                                elif class_name in KNOWN_CLASSES_LABELS[self.config.general.split]['task3']:
                                    Split_B_T3_results.append(np.array((float(ap), float(ap_50), float(ap_25))))
                                # elif class_name in KNOWN_CLASSES_LABELS[self.config.general.split]['task4']:
                                #     Split_B_T4_results.append(np.array((float(ap), float(ap_50), float(ap_25))))
                            # else:
                            #     assert False, 'class not known!' #changed
                            elif self.config.general.split == 'C':
                                if class_name in KNOWN_CLASSES_LABELS[self.config.general.split]['task1']:
                                    Split_C_T1_results.append(np.array((float(ap), float(ap_50), float(ap_25))))
                                elif class_name in KNOWN_CLASSES_LABELS[self.config.general.split]['task2']:
                                    Split_C_T2_results.append(np.array((float(ap), float(ap_50), float(ap_25))))
                                elif class_name in KNOWN_CLASSES_LABELS[self.config.general.split]['task3']:
                                    Split_C_T3_results.append(np.array((float(ap), float(ap_50), float(ap_25))))
                    else:
                        ap_results[f"{log_prefix}_{class_name}_val_ap"+"_"+self.config.general.logg_suffix ] = float(ap)
                        ap_results[f"{log_prefix}_{class_name}_val_ap_50"+"_"+self.config.general.logg_suffix ] = float(ap_50)
                        ap_results[f"{log_prefix}_{class_name}_val_ap_25"+"_"+self.config.general.logg_suffix ] = float(ap_25)

            if self.validation_dataset.dataset_name == "scannet200":
                head_results = np.stack(head_results)
                common_results = np.stack(common_results)
                tail_results = np.stack(tail_results)
                
                mean_tail_results = np.nanmean(tail_results, axis=0)
                mean_common_results = np.nanmean(common_results, axis=0)
                mean_head_results = np.nanmean(head_results, axis=0)
                
                if self.config.general.split == 'A':
                    Split_A_T1_results = np.stack(Split_A_T1_results)
                    Split_A_T2_results = np.stack(Split_A_T2_results)
                    Split_A_T3_results = np.stack(Split_A_T3_results)
                    
                    mean_Split_A_T1_results = np.nanmean(Split_A_T1_results, axis = 0)
                    mean_Split_A_T2_results = np.nanmean(Split_A_T2_results, axis = 0)
                    mean_Split_A_T3_results = np.nanmean(Split_A_T3_results, axis = 0)
                    
                    if self.config.general.OW_task == 'task1':
                        ap_results[f"{log_prefix}_current_known_mAP"+"_"+self.config.general.logg_suffix] = np.nanmean(Split_A_T1_results, axis = 0)[0]
                    elif self.config.general.OW_task == 'task2':
                        ap_results[f"{log_prefix}_previously_known_mAP"+"_"+self.config.general.logg_suffix] = np.nanmean(Split_A_T1_results, axis = 0)[0]
                        ap_results[f"{log_prefix}_current_known_mAP"+"_"+self.config.general.logg_suffix] = np.nanmean(Split_A_T2_results, axis = 0)[0]
                    elif self.config.general.OW_task == 'task3':
                        ap_results[f"{log_prefix}_previously_known_mAP"+"_"+self.config.general.logg_suffix] = np.nanmean(np.concatenate([Split_A_T1_results,Split_A_T2_results], axis=0), axis = 0)[0]
                        ap_results[f"{log_prefix}_current_known_mAP"+"_"+self.config.general.logg_suffix] = np.nanmean(Split_A_T3_results, axis = 0)[0]
                    
                elif self.config.general.split == 'B':
                    try:
                        Split_B_T1_results = np.stack(Split_B_T1_results)
                        Split_B_T2_results = np.stack(Split_B_T2_results)
                        Split_B_T3_results = np.stack(Split_B_T3_results)
                        # Split_B_T4_results = np.stack(Split_B_T4_results)
                        
                        mean_Split_B_T1_results = np.nanmean(Split_B_T1_results, axis = 0)
                        mean_Split_B_T2_results = np.nanmean(Split_B_T2_results, axis = 0)
                        mean_Split_B_T3_results = np.nanmean(Split_B_T3_results, axis = 0)
                        # mean_Split_B_T4_results = np.nanmean(Split_B_T4_results, axis = 0)
                        if self.config.general.OW_task == 'task1':
                            ap_results[f"{log_prefix}_current_known_mAP"+"_"+self.config.general.logg_suffix] = np.nanmean(Split_B_T1_results, axis = 0)[0]
                        elif self.config.general.OW_task == 'task2':
                            ap_results[f"{log_prefix}_previously_known_mAP"+"_"+self.config.general.logg_suffix] = np.nanmean(Split_B_T1_results, axis = 0)[0]
                            ap_results[f"{log_prefix}_current_known_mAP"+"_"+self.config.general.logg_suffix] = np.nanmean(Split_B_T2_results, axis = 0)[0]
                        elif self.config.general.OW_task == 'task3':
                            ap_results[f"{log_prefix}_previously_known_mAP"+"_"+self.config.general.logg_suffix] = np.nanmean(np.concatenate([Split_B_T1_results,Split_B_T2_results], axis=0), axis = 0)[0]
                            ap_results[f"{log_prefix}_current_known_mAP"+"_"+self.config.general.logg_suffix] = np.nanmean(Split_B_T3_results, axis = 0)[0]
                    except:
                        pass
                elif self.config.general.split == 'C':
                    try:
                        Split_C_T1_results = np.stack(Split_C_T1_results)
                        Split_C_T2_results = np.stack(Split_C_T2_results)
                        Split_C_T3_results = np.stack(Split_C_T3_results)
                        # Split_C_T4_results = np.stack(Split_C_T4_results)
                        
                        mean_Split_C_T1_results = np.nanmean(Split_C_T1_results, axis = 0)
                        mean_Split_C_T2_results = np.nanmean(Split_C_T2_results, axis = 0)
                        mean_Split_C_T3_results = np.nanmean(Split_C_T3_results, axis = 0)
                        # mean_Split_C_T4_results = np.nanmean(Split_C_T4_results, axis = 0)
                        if self.config.general.OW_task == 'task1':
                            ap_results[f"{log_prefix}_current_known_mAP"+"_"+self.config.general.logg_suffix] = np.nanmean(Split_C_T1_results, axis = 0)[0]
                        elif self.config.general.OW_task == 'task2':
                            ap_results[f"{log_prefix}_previously_known_mAP"+"_"+self.config.general.logg_suffix] = np.nanmean(Split_C_T1_results, axis = 0)[0]
                            ap_results[f"{log_prefix}_current_known_mAP"+"_"+self.config.general.logg_suffix] = np.nanmean(Split_C_T2_results, axis = 0)[0]
                        elif self.config.general.OW_task == 'task3':
                            ap_results[f"{log_prefix}_previously_known_mAP"+"_"+self.config.general.logg_suffix] = np.nanmean(np.concatenate([Split_C_T1_results,Split_C_T2_results], axis=0), axis = 0)[0]
                            ap_results[f"{log_prefix}_current_known_mAP"+"_"+self.config.general.logg_suffix] = np.nanmean(Split_C_T3_results, axis = 0)[0]
                    except:
                        pass
                ap_results[f"{log_prefix}_mean_tail_ap"+"_"+self.config.general.logg_suffix ] = mean_tail_results[0]
                ap_results[f"{log_prefix}_mean_common_ap"+"_"+self.config.general.logg_suffix ] = mean_common_results[0]
                ap_results[f"{log_prefix}_mean_head_ap"+"_"+self.config.general.logg_suffix ] = mean_head_results[0]

                ap_results[f"{log_prefix}_mean_tail_ap_50"+"_"+self.config.general.logg_suffix ] = mean_tail_results[1]
                ap_results[f"{log_prefix}_mean_common_ap_50"+"_"+self.config.general.logg_suffix ] = mean_common_results[1]
                ap_results[f"{log_prefix}_mean_head_ap_50"+"_"+self.config.general.logg_suffix ] = mean_head_results[1]

                ap_results[f"{log_prefix}_mean_tail_ap_25"+self.config.general.logg_suffix ] = mean_tail_results[2]
                ap_results[f"{log_prefix}_mean_common_ap_25"+self.config.general.logg_suffix ] = mean_common_results[2]
                ap_results[f"{log_prefix}_mean_head_ap_25"+self.config.general.logg_suffix ] = mean_head_results[2]
                
                if self.config.general.split == 'A':
                    ap_results[f"{log_prefix}_mean_Split_A_T1_ap"+"_"+self.config.general.logg_suffix ] = mean_Split_A_T1_results[0]
                    ap_results[f"{log_prefix}_mean_Split_A_T2_ap"+"_"+self.config.general.logg_suffix ] = mean_Split_A_T2_results[0]
                    ap_results[f"{log_prefix}_mean_Split_A_T3_ap"+"_"+self.config.general.logg_suffix ] = mean_Split_A_T3_results[0]

                    ap_results[f"{log_prefix}_mean_Split_A_T1_ap_50"+"_"+self.config.general.logg_suffix ] = mean_Split_A_T1_results[1]
                    ap_results[f"{log_prefix}_mean_Split_A_T2_ap_50"+"_"+self.config.general.logg_suffix ] = mean_Split_A_T2_results[1]
                    ap_results[f"{log_prefix}_mean_Split_A_T3_ap_50"+"_"+self.config.general.logg_suffix ] = mean_Split_A_T3_results[1]
                    
                    ap_results[f"{log_prefix}_mean_Split_A_T1_ap_25"+"_"+self.config.general.logg_suffix ] = mean_Split_A_T1_results[2]
                    ap_results[f"{log_prefix}_mean_Split_A_T2_ap_25"+"_"+self.config.general.logg_suffix ] = mean_Split_A_T2_results[2]
                    ap_results[f"{log_prefix}_mean_Split_A_T3_ap_25"+"_"+self.config.general.logg_suffix ] = mean_Split_A_T3_results[2]
                    
                    
                    
                elif self.config.general.split == 'B':
                    try:
                        ap_results[f"{log_prefix}_mean_Split_C_T1_ap"+"_"+self.config.general.logg_suffix ] = mean_Split_B_T1_results[0]
                        ap_results[f"{log_prefix}_mean_Split_B_T2_ap"+"_"+self.config.general.logg_suffix ] = mean_Split_B_T2_results[0]
                        ap_results[f"{log_prefix}_mean_Split_B_T3_ap"+"_"+self.config.general.logg_suffix ] = mean_Split_B_T3_results[0]
                        # ap_results[f"{log_prefix}_mean_Split_B_T4_ap"+"_"+self.config.general.logg_suffix ] = mean_Split_B_T4_results[0]

                        ap_results[f"{log_prefix}_mean_Split_B_T1_ap_50"+"_"+self.config.general.logg_suffix ] = mean_Split_B_T1_results[1]
                        ap_results[f"{log_prefix}_mean_Split_B_T2_ap_50"+"_"+self.config.general.logg_suffix ] = mean_Split_B_T2_results[1]
                        ap_results[f"{log_prefix}_mean_Split_B_T3_ap_50"+"_"+self.config.general.logg_suffix ] = mean_Split_B_T3_results[1]
                        # ap_results[f"{log_prefix}_mean_Split_B_T4_ap_50"+"_"+self.config.general.logg_suffix ] = mean_Split_B_T4_results[1]
                        
                        ap_results[f"{log_prefix}_mean_Split_B_T1_ap_25"+"_"+self.config.general.logg_suffix ] = mean_Split_B_T1_results[2]
                        ap_results[f"{log_prefix}_mean_Split_B_T2_ap_25"+"_"+self.config.general.logg_suffix ] = mean_Split_B_T2_results[2]
                        ap_results[f"{log_prefix}_mean_Split_B_T3_ap_25"+"_"+self.config.general.logg_suffix ] = mean_Split_B_T3_results[2]
                        # ap_results[f"{log_prefix}_mean_Split_B_T4_ap_25"+"_"+self.config.general.logg_suffix ] = mean_Split_B_T4_results[2]
                    except:
                        pass
                elif self.config.general.split == 'C':
                    try:
                        ap_results[f"{log_prefix}_mean_Split_C_T1_ap"+"_"+self.config.general.logg_suffix ] = mean_Split_C_T1_results[0]
                        ap_results[f"{log_prefix}_mean_Split_C_T2_ap"+"_"+self.config.general.logg_suffix ] = mean_Split_C_T2_results[0]
                        ap_results[f"{log_prefix}_mean_Split_C_T3_ap"+"_"+self.config.general.logg_suffix ] = mean_Split_C_T3_results[0]
                        # ap_results[f"{log_prefix}_mean_Split_B_T4_ap"+"_"+self.config.general.logg_suffix ] = mean_Split_B_T4_results[0]

                        ap_results[f"{log_prefix}_mean_Split_C_T1_ap_50"+"_"+self.config.general.logg_suffix ] = mean_Split_C_T1_results[1]
                        ap_results[f"{log_prefix}_mean_Split_C_T2_ap_50"+"_"+self.config.general.logg_suffix ] = mean_Split_C_T2_results[1]
                        ap_results[f"{log_prefix}_mean_Split_C_T3_ap_50"+"_"+self.config.general.logg_suffix ] = mean_Split_C_T3_results[1]
                        # ap_results[f"{log_prefix}_mean_Split_B_T4_ap_50"+"_"+self.config.general.logg_suffix ] = mean_Split_B_T4_results[1]
                        
                        ap_results[f"{log_prefix}_mean_Split_C_T1_ap_25"+"_"+self.config.general.logg_suffix ] = mean_Split_C_T1_results[2]
                        ap_results[f"{log_prefix}_mean_Split_C_T2_ap_25"+"_"+self.config.general.logg_suffix ] = mean_Split_C_T2_results[2]
                        ap_results[f"{log_prefix}_mean_Split_C_T3_ap_25"+"_"+self.config.general.logg_suffix ] = mean_Split_C_T3_results[2]
                        # ap_results[f"{log_prefix}_mean_Split_B_T4_ap_25"+"_"+self.config.general.logg_suffix ] = mean_Split_B_T4_results[2]
                    except:
                        pass
                overall_ap_results = np.nanmean(np.vstack((head_results, common_results, tail_results)), axis=0)

                ap_results[f"{log_prefix}_mean_ap"+"_"+self.config.general.logg_suffix ] = overall_ap_results[0]
                ap_results[f"{log_prefix}_mean_ap_50"] = overall_ap_results[1]
                ap_results[f"{log_prefix}_mean_ap_25"+"_"+self.config.general.logg_suffix ] = overall_ap_results[2]

                ap_results = {key: 0. if math.isnan(score) else score for key, score in ap_results.items()}
                if U_Recall:
                    U_Recall_results["Unknown Recall at 0.5_"+self.config.general.logg_suffix ] = U_Recall[50]
                    U_Recall_results["Unknown Recall at 0.9_"+self.config.general.logg_suffix ] = U_Recall[90]
                    U_Recall_results["Unknown Recall at 0.8_"+self.config.general.logg_suffix ] = U_Recall[80]
                    
                    WI_results["Wilderness Impact "+self.config.general.logg_suffix] = WI[0.5]
                    A_OSE_results["Absolute open set error "+self.config.general.logg_suffix] = A_OSE[0.5]
                    
                    self.log_dict(U_Recall_results)
                    self.log_dict(A_OSE_results)
                    self.log_dict(WI_results)
                    
            else:
                mean_ap = statistics.mean([item for key, item in ap_results.items() if key.endswith("val_ap")])
                mean_ap_50 = statistics.mean([item for key, item in ap_results.items() if key.endswith("val_ap_50")])
                mean_ap_25 = statistics.mean([item for key, item in ap_results.items() if key.endswith("val_ap_25")])

                ap_results[f"{log_prefix}_mean_ap"+"_"+self.config.general.logg_suffix ] = mean_ap
                ap_results[f"{log_prefix}_mean_ap_50"] = mean_ap_50
                ap_results[f"{log_prefix}_mean_ap_25"+"_"+self.config.general.logg_suffix ] = mean_ap_25

                ap_results = {key: 0. if math.isnan(score) else score for key, score in ap_results.items()}
        except (IndexError, OSError) as e:
            print("NO SCORES!!!")
            ap_results[f"{log_prefix}_mean_ap"+"_"+self.config.general.logg_suffix ] = 0.
            ap_results[f"{log_prefix}_mean_ap_50"] = 0.
            ap_results[f"{log_prefix}_mean_ap_25"+"_"+self.config.general.logg_suffix ] = 0.

        self.log_dict(ap_results)
            
        if not self.config.general.export:
            shutil.rmtree(base_path)
            
        if self.config.general.save_KN_UKN_tSNE:
            file_path_p = self.config.general.save_features_in+"val_set/"+"known_unknown_queries/"

            if os.path.exists(file_path_p):
                # start = time.time()
                quer_features = torch.tensor([])
                pred_id = torch.tensor([])
                for file in os.listdir(file_path_p):
                    queries, labels = torch.load(os.path.join(file_path_p,file))
                    quer_features = torch.cat([quer_features,queries.detach().cpu()])
                    pred_id = torch.cat([pred_id,labels.detach().cpu()])
                    
                label_ids = pred_id.numpy()
                features = quer_features.numpy()
                mask = label_ids==198

                file_path_p = self.config.general.save_dir
                if not os.path.exists(file_path_p+"/plots"):
                    os.makedirs(file_path_p+"/plots")
                features_final = TSNE(perplexity=30).fit_transform(features) 
                palette = np.array(sb.color_palette("hls", 255))
                palette[198] = np.array([0,0,0]) #unknown is assigned a black color
                f = plt.figure(figsize=(8, 8))
                ax = plt.subplot(aspect='equal')
                sc = ax.scatter(features_final[:,0], features_final[:,1], lw=0, s=10, c=palette[label_ids.astype(int)])
                plt.savefig(file_path_p+"/plots"+'/tSNE_plot.png')
                f1 = plt.figure(figsize=(8, 8))
                ax1 = plt.subplot(aspect='equal')
                sc1 = ax1.scatter(features_final[:,0][mask], features_final[:,1][mask], lw=0, s=10, c=palette[label_ids[mask].astype(int)])
                plt.savefig(file_path_p+"/plots"+'/tSNE_plot_unkn.png')
                f2 = plt.figure(figsize=(8, 8))
                ax2 = plt.subplot(aspect='equal')
                sc2 = ax2.scatter(features_final[:,0][~mask], features_final[:,1][~mask], lw=0, s=10, c=palette[label_ids[~mask].astype(int)])
                plt.savefig(file_path_p+"/plots"+'/tSNE_plot_kn.png')
                
        del self.preds
        del self.bbox_preds
        del self.bbox_gt

        gc.collect()

        self.preds = dict()
        self.bbox_preds = dict()
        self.bbox_gt = dict()
        
        
                
    def test_epoch_end(self, outputs):
        if self.config.general.export:
            return

        self.eval_instance_epoch_end()

        dd = defaultdict(list)
        for output in outputs:
            for key, val in output.items():  # .items() in Python 3.
                dd[key].append(val)

        dd = {k: statistics.mean(v) for k, v in dd.items()}

        dd['val_mean_loss_ce'+"_"+self.config.general.logg_suffix ] = statistics.mean([item for item in [v for k,v in dd.items() if "loss_ce" in k]])
        dd['val_mean_loss_mask'+"_"+self.config.general.logg_suffix ] = statistics.mean([item for item in [v for k,v in dd.items() if "loss_mask" in k]])
        dd['val_mean_loss_dice'+"_"+self.config.general.logg_suffix ] = statistics.mean([item for item in [v for k,v in dd.items() if "loss_dice" in k]])
        
        self.log_dict(dd)

    def configure_optimizers(self):
        optimizer = hydra.utils.instantiate(
            self.config.optimizer, params=self.parameters()
        )
        if "steps_per_epoch" in self.config.scheduler.scheduler.keys():
            self.config.scheduler.scheduler.steps_per_epoch = len(
                self.train_dataloader()
            )
        lr_scheduler = hydra.utils.instantiate(
            self.config.scheduler.scheduler, optimizer=optimizer
        )
        scheduler_config = {"scheduler": lr_scheduler}
        scheduler_config.update(self.config.scheduler.pytorch_lightning_params)
        return [optimizer], [scheduler_config]

    def prepare_data(self):
        self.train_dataset = hydra.utils.instantiate(self.config.data.train_dataset)
        self.validation_dataset = hydra.utils.instantiate( 
            self.config.data.validation_dataset
        )
        self.test_dataset = hydra.utils.instantiate(self.config.data.test_dataset)
        self.labels_info = self.train_dataset.label_info

    def train_dataloader(self):
        c_fn = hydra.utils.instantiate(self.config.data.train_collation)
        return hydra.utils.instantiate(
            self.config.data.train_dataloader,
            self.train_dataset,
            collate_fn=c_fn,
        )

    def val_dataloader(self):
        c_fn = hydra.utils.instantiate(self.config.data.validation_collation)
        return hydra.utils.instantiate(
            self.config.data.validation_dataloader,
            self.validation_dataset,
            collate_fn=c_fn,
        )

    def test_dataloader(self):
        c_fn = hydra.utils.instantiate(self.config.data.test_collation)
        return hydra.utils.instantiate(
            self.config.data.test_dataloader,
            self.test_dataset,
            collate_fn=c_fn,
        )
    
        
    """
    AUTOLABELER
    """
    def Auto_Labeling(self, batch_idx, output, target, topk, raw_coord,point2segment,return_ukn_idxs = False):
    
        IoU_th = 0.01
        Conf_th = self.config.general.Conf_th
        #IoU per point
        IoU_matrix = self.get_IoU(target[batch_idx]['segment_mask'][(target[batch_idx]['labels']!=253)*(target[batch_idx]['labels']!=198)],output['pred_masks'][batch_idx])
        
        min_id = sum([point2segment[i].shape[0] for i in range(batch_idx)]) if batch_idx!=0 else 0
        max_id = min_id+point2segment[batch_idx].shape[0]
        if self.config.general.multi_scale_autolabeler_flag:
            score = list()
            for L in range(len(output['aux_outputs'])):
                
                pred_logits = output['aux_outputs'][L]['pred_logits']
                pred_logits = torch.functional.F.softmax(
                    pred_logits ,
                    dim=-1)[..., :-1]
                if self.model.train_on_segments:
                    masks = output['aux_outputs'][L]['pred_masks'][batch_idx][target[batch_idx]['point2segment']]
                else:
                    masks = output['aux_outputs'][L]['pred_masks'][batch_idx]
                
                # if self.config.general.use_obj_loss:
                #     score_L = output['obj']
                # else:
                score_L, _, _ = self.get_scores(
                        pred_logits[batch_idx],
                        masks,
                        raw_coord[min_id:max_id],
                        point2segment[batch_idx])
                score.append(score_L)
                
            scores = torch.max(torch.stack(score), dim=0).values
        else:
            if self.model.train_on_segments:
                masks = output['pred_masks'][batch_idx][target[batch_idx]['point2segment']]
            else:
                masks = output['pred_masks'][batch_idx]
            pred_logits = output['pred_logits']
            pred_logits = torch.functional.F.softmax(
                pred_logits ,
                dim=-1)[..., :-1]
            
            # if self.config.general.use_obj_loss:
            #     scores = output['obj'][batch_idx]
            # else:
            scores, _, _ = self.get_scores(
                    pred_logits[batch_idx],
                    masks,
                    raw_coord[min_id:max_id],
                    point2segment[batch_idx])
        # potential_ukn = torch.Tensor([i for i in range(self.model.num_queries)]).long()
        
        # for i in range(indices[batch_idx][0].shape[0]): 
        #     potential_ukn = potential_ukn[potential_ukn!=indices[batch_idx][0][i]]
        scores = scores.detach().cpu()  
        
        
        # topk_scores_indices = scores_indices[potential_ukn][:topk]
        
        
        if self.config.general.use_conf_th:
            topk_scores_indices = torch.where(scores>Conf_th)[0]
        else:
            _, scores_indices = scores.sort(descending=True)
            topk_scores_indices = scores_indices[:topk]
            
        max_IoU_per_GT = torch.max(IoU_matrix, dim=0).values
        
        unk_mask = max_IoU_per_GT<IoU_th
        IoU_indices = torch.where(unk_mask)[0].detach().cpu()
            #indices of masks with topk score and no intersection with the GT
        indx_ukns = torch.from_numpy(np.intersect1d(IoU_indices,topk_scores_indices))
        # print(indx_ukns)
        if return_ukn_idxs:
            pred_labels = torch.argmax(pred_logits, dim = -1)
            indx_ukns_c = indx_ukns.to(pred_labels.device).clone()
            return pred_labels.permute(1,0)[indx_ukns_c].clone()
        else:
            # output['pred_logits'] = output['pred_logits'].permute(0,2,1)
            # output['aux_outputs'][-1]['pred_logits'] = output['aux_outputs'][-1]['pred_logits'].permute(0,2,1)
            # output['pred_logits'][batch_idx][-2][indx_ukns] = 100000
            # output['aux_outputs'][-1]['pred_logits'][batch_idx][-2][indx_ukns] = 100000
            # output['pred_logits'] = output['pred_logits'].permute(0,2,1)
            # output['aux_outputs'][-1]['pred_logits'] = output['aux_outputs'][-1]['pred_logits'].permute(0,2,1)
            target[batch_idx]['segment_mask'] = torch.cat((target[batch_idx]['segment_mask'],output['pred_masks'][batch_idx].permute(1,0)[indx_ukns]), dim=0)
            target[batch_idx]['labels'] = torch.cat((target[batch_idx]['labels'],(torch.ones_like(indx_ukns)*self.ukn_cls).to(target[batch_idx]['labels'].device)))
            return target
    
    def get_IoU(self, GT_segment_mask, P_segment_mask):
        threshold = 0.5
        if GT_segment_mask.nelement() !=0 :
            intersection = GT_segment_mask.float()@(torch.nn.Sigmoid()(P_segment_mask)>threshold).float()
            union = torch.stack(tuple(torch.sum((GT_segment_mask[inst_id,:].float()+(torch.nn.Sigmoid()(P_segment_mask)>threshold).float().T),dim=1)-intersection[inst_id,:] for inst_id in range(GT_segment_mask.shape[0])), dim = 0)
            IoU = intersection/union
        else:
            IoU = torch.zeros(P_segment_mask.shape[-1])
        return IoU
    def get_scores(self, mask_cls, mask_pred, raw_coord, point2segment):
        
        scores_per_query = torch.max(mask_cls, dim = 1).values

        result_pred_mask = (mask_pred > 0).float()
        heatmap = mask_pred.float().sigmoid()

        mask_scores_per_image = (heatmap * result_pred_mask).sum(0) / (result_pred_mask.sum(0) + 1e-6)
        score = scores_per_query * mask_scores_per_image
        
        return score, result_pred_mask, heatmap
    
        
    """
    DEFINING TASKS AND SPLITS
        task(): is method that assigns a label ignore (253) during training, 
                or unknown (200) during evaluation for a specific set of classes
                that depend on the task and the split
                
        NB: please refer to the paper for more details on the three splits 
    """

    def task(self,target, eval = False, task_ow = 'task1', split = 'A',File_names = None):
        unknown_classes_ids = UNKNOWN_CLASSES_IDS[split][task_ow]
        prev_classes_ids = PREV_KNOWN_CLASSES_IDS[split][task_ow]
        examplar_list_path = "data/processed/scannet200/exemplars/"+split+"/"+task_ow+"/"+"examplar_list.yaml"
        if os.path.exists(examplar_list_path):
            with open(examplar_list_path) as user_file:
                examplar_list = yaml.load(user_file, Loader=SafeLoader)
                
        if (unknown_classes_ids != None) and (prev_classes_ids != None):
            if eval:
                for batch_id in range(len(target)):
                    for k in unknown_classes_ids:
                        try:
                            target[batch_id]['labels'][target[batch_id]['labels']==k]=198
                        except:
                            pass
                
                    
            elif not eval and self.train_oracle:
                for batch_id in range(len(target)):
                    for k, h in zip(unknown_classes_ids,prev_classes_ids):
                        try:
                            target[batch_id]['labels'][target[batch_id]['labels']==k]=self.ukn_cls #set the tail classes as known
                        except:
                            pass
            else:        
                for batch_id in range(len(target)):
                    for k, h in zip(unknown_classes_ids,prev_classes_ids):
                        try:
                            if not self.config.general.finetune:
                                if File_names == None:
                                    target[batch_id]['segment_mask'][target[batch_id]['labels']==h] = False #remove the ground truth masks from the scene
                                    target[batch_id]['labels'][target[batch_id]['labels']==k]=253 #train wihtout the unknown labels
                                    target[batch_id]['labels'][target[batch_id]['labels']==h]=253 #ignore the known labels in task1
                                else:
                                    for file_n in File_names:
                                        if file_n in examplar_list:
                                            target[batch_id]['labels'][target[batch_id]['labels']==k]=253
                                        else:
                                            target[batch_id]['segment_mask'][target[batch_id]['labels']==h] = False #remove the ground truth masks from the scene
                                            target[batch_id]['labels'][target[batch_id]['labels']==k]=253 #train wihtout the unknown labels
                                            target[batch_id]['labels'][target[batch_id]['labels']==h]=253 #ignore the known labels in task1
                            else:
                                target[batch_id]['labels'][target[batch_id]['labels']==k]=253
                        except:
                            pass
        elif prev_classes_ids != None:
            if not self.train_oracle: 
                for batch_id in range(len(target)):
                    for h in prev_classes_ids:
                        try:
                            if not self.config.general.finetune:
                                if File_names == None:
                                    target[batch_id]['segment_mask'][target[batch_id]['labels']==h] = False #remove the ground truth masks from the scene
                                    target[batch_id]['labels'][target[batch_id]['labels']==h]=253 #ignore the known labels in task1
                                else:
                                    for file_n in File_names:
                                        if file_n not in examplar_list:
                                            target[batch_id]['segment_mask'][target[batch_id]['labels']==h] = False #remove the ground truth masks from the scene
                                            target[batch_id]['labels'][target[batch_id]['labels']==h]=253 #ignore the known labels in task1
                        except:
                            pass
        elif unknown_classes_ids != None:
            if eval:
                for batch_id in range(len(target)):
                    for k in unknown_classes_ids:
                        try:
                            target[batch_id]['labels'][target[batch_id]['labels']==k]=self.ukn_cls-2
                        except:
                            pass
            elif not eval and self.train_oracle:
                for batch_id in range(len(target)):
                    for k in unknown_classes_ids:
                        try:
                            target[batch_id]['labels'][target[batch_id]['labels']==k]=self.ukn_cls
                        except:
                            pass
            elif not eval and not self.train_oracle:        
                for batch_id in range(len(target)):
                    for k in unknown_classes_ids:
                        try:
                            target[batch_id]['labels'][target[batch_id]['labels']==k]=253
                        except:
                            pass
        return target
    """
    ENERGY BASED UNKNOWN IDENTIFIER
    
    DETAILS: This method is not used in our pipeline, please refer to the original paper
             "Toward open world object detection(https://arxiv.org/abs/2103.02603)" for 
             details. 
    """
    def learn_energy(self, temp = 0.05):
        #the energy is computed over the known classes
        known_classes_ids = list(set(self.KNOWN_CLASSES_IDS))
        if self.ukn_cls in known_classes_ids:
            known_classes_ids.remove(self.ukn_cls)#remove the label of the unknown class
        
        file_path_p = os.path.join(os.path.join(self.config.general.save_energy_training_dataset_in,"train_set"), "logits")
        if os.path.exists(file_path_p):
            temp = self.config.general.ENERGY_TEMP
            file_path = file_path_p
            files = os.listdir(file_path_p)   
            for id, file in enumerate(files):
                path = os.path.join(file_path_p, file)
                pred_logits_kn_id, pred_logits_unk_id = torch.load(path)
                if id == 0:
                    logits_ukn = pred_logits_unk_id
                    logits_kn = pred_logits_kn_id
                else:
                    logits_ukn = torch.cat([logits_ukn, pred_logits_unk_id], dim = 0)
                    logits_kn = torch.cat([logits_kn, pred_logits_kn_id], dim = 0)
                

            lse_unkn = (temp * torch.logsumexp(logits_ukn[:, torch.tensor(known_classes_ids)] / temp, dim=1)).detach().cpu().tolist()
            lse_kn = (temp * torch.logsumexp(logits_kn[:, torch.tensor(known_classes_ids)] / temp, dim=1)).detach().cpu().tolist()
            
            wb_dist_param = []
            save_WB_in = self.config.general.save_dir+"/energy_dist_"+str(self.num_seen_classes)+".pkl"
            wb_unk = Fit_Weibull_3P(failures=lse_unkn, show_probability_plot=False, print_results=False)
            wb_kn = Fit_Weibull_3P(failures=lse_kn, show_probability_plot=False, print_results=False)
            
            wb_dist_param.append({"scale_unk": wb_unk.alpha, "shape_unk": wb_unk.beta, "shift_unk": wb_unk.gamma})
            wb_dist_param.append({"scale_known": wb_kn.alpha, "shape_known": wb_kn.beta, "shift_known": wb_kn.gamma})
            
            torch.save(wb_dist_param, save_WB_in)
            plt.hist(lse_kn, density = True,alpha=0.5, label='known')
            plt.hist(lse_unkn, density = True, alpha=0.5, label='unk')
            plt.legend(loc='upper right')
            plt.savefig(os.path.join(self.config.general.save_dir, 'energy.png'))
            plt.clf()
            try:
                shutil.rmtree(file_path)
            except:
                pass
        else: 
            print(f"generate {file_path_p} first")
            assert 1==0

    def compute_prob(self, x, distribution, shift):
        eps_radius = 0.5 # in the baseline of object detection is 0.5
        num_eval_points = 100
        start_x = x - eps_radius
        end_x = x + eps_radius
        step = (end_x - start_x) / num_eval_points
        dx = torch.linspace(x - eps_radius, x + eps_radius, num_eval_points)
        dx = dx[dx>shift]
        pdf = distribution.log_prob(dx).exp()
        prob = torch.sum(pdf * step)
        return prob
    
    def create_distribution(self, scale, shape, shift):
        wd = Weibull(scale=scale, concentration=shape)
        transforms = AffineTransform(loc=shift, scale=1.)
        weibull = TransformedDistribution(wd, transforms)
        return weibull
    
    def update_label_based_on_energy(self, logits, temp = 0.05):
        known_classes_ids = list(set(self.KNOWN_CLASSES_IDS))
        if self.ukn_cls in known_classes_ids:
            known_classes_ids.remove(self.ukn_cls)
        lse = (temp * torch.logsumexp(logits[:, :,torch.tensor(known_classes_ids)] / temp, dim=-1)).detach().cpu().tolist()
        pred_clss = torch.functional.F.softmax(
            logits ,
            dim=-1)[..., :-1]
        pred_clss = torch.argmax(pred_clss, dim=-1)
        
        for bid in range(len(lse)):
            for i, energy in enumerate(lse[bid]):
                p_unk = self.compute_prob(energy, self.unk_dist, self.unk_dist_shift)
                p_known = self.compute_prob(energy, self.known_dist, self.known_dist_shift)
                if p_unk <= p_known:
                    if pred_clss[bid, i] == self.ukn_cls:
                        logits[bid,i,-2] = -100 
                else:
                    if pred_clss[bid, i] != self.ukn_cls:
                        logits[bid, i,-2] = 100000
                    
        return logits
    
    """
    PROBABILITY CORRECTION
    
    DETAILS: The implementation of the probability correction
    """
    #unknown distribution 
    def set_unkn_distribution(self, device):
        self.unkn_dist = unknown_cls_distribution(save_dir = self.config.general.save_dir,
                                                  split = self.config.general.split,
                                                  ow_task = self.config.general.OW_task,
                                                  device = device,
                                                  margin = self.config.general.margin)  
        self.unknown_class_distribution_is_set = True  
        return 
    
    def correct_classes_prob(self, output, train = False):
        
        prediction = torch.functional.F.softmax(    output['pred_logits'],
                                                        dim=-1)
        known_cls_indc = list(range(self.ukn_cls))
        known_cls_indc.append(201)
        if PREV_KNOWN_CLASSES_IDS[self.config.general.split][self.config.general.OW_task] != None:
            sum = PREV_KNOWN_CLASSES_IDS[self.config.general.split][self.config.general.OW_task]+KNOWN_CLASSES_IDS[self.config.general.split][self.config.general.OW_task]
            
        else:
            sum = KNOWN_CLASSES_IDS[self.config.general.split][self.config.general.OW_task]
            
        if self.ukn_cls in sum:
            sum.remove(self.ukn_cls) #remove the ID of the unknown class
            
        a = 1-torch.sum(prediction[:,:,sum], dim=-1)

        unkn_prob = self.unkn_dist.get_prob(output["refin_queries"]).to(self.device)*(a)
        unkn_prob = prediction[:,:,-2]+unkn_prob-unkn_prob*prediction[:,:,-2]
        prediction[:,:,known_cls_indc] = (1-unkn_prob)[:,:,None].repeat(1,1,201)*prediction[:,:,known_cls_indc]/prediction[:,:,known_cls_indc].sum(dim=-1)[:,:,None].repeat(1,1,201)
        
        return prediction[..., :-1]

class unknown_cls_distribution():
    def __init__(self, save_dir = ".saved/",split = 'A',ow_task = 'task1', device = 'cpu',hidden_dim = 128, margin=1.0):
        
        self.device = device
        self.save_dir = save_dir
        self.hidden_dim=hidden_dim
        self.ukn_cls = 200
        self.enable_inference = True
        self.enable_correction = True
        self.ow_task = ow_task
        self.split = split
        self.margin = margin
        self.params_are_set = False

    
    def set_means(self):
        if os.path.exists(self.save_dir +'/'+self.ow_task+'_classes_means.pkl'):
            self.reach_unkn_means = torch.load(self.save_dir +'/'+self.ow_task+'_classes_means.pkl').detach().cpu()
            self.kn_mean = self.reach_unkn_means[:-1,:]
        else:
            self.enable_correction = False
    
    def set_params(self):
        alpha = 0.95
        beta = 0.05
        self.scale = self.margin/(2*np.log((alpha/beta)*((1-beta)/(1-alpha))))
        self.shift = self.margin+self.scale*np.log((1-alpha)/alpha)
        self.sensitivity = self.margin
        self.params_are_set = True
            
    
    def get_prob(self, queries):
        self.set_means()
        if not self.params_are_set:
            self.set_params()
        if self.enable_correction:
            queries = queries.detach().cpu()
            return torch.sigmoid(((torch.min(torch.norm(queries[None,:,:,:]-self.kn_mean[:,None,None,:],dim = -1), dim=0).values-self.shift)/self.scale-self.margin/2))
        else:
            self.enable_correction = True
            return torch.zeros((queries.shape[0], queries.shape[1]))
        
    
        
        
        
        
           
    
trainer/trainer.py
aminebdj/3D-OWIS
from typing import Tuple

import torch
from torch.autograd import Function
import torch.nn as nn

import pointops2_cuda as pointops_cuda
import time

class FurthestSampling(Function):
    @staticmethod
    def forward(ctx, xyz, offset, new_offset):
        """
        input: xyz: (n, 3), offset: (b), new_offset: (b)
        output: idx: (m)
        """
        assert xyz.is_contiguous()
        n, b, n_max = xyz.shape[0], offset.shape[0], offset[0]
        for i in range(1, b):
            n_max = max(offset[i] - offset[i-1], n_max)
        idx = torch.cuda.IntTensor(new_offset[b-1].item()).zero_()
        tmp = torch.cuda.FloatTensor(n).fill_(1e10)
        pointops_cuda.furthestsampling_cuda(b, n_max, xyz, offset, new_offset, tmp, idx)
        del tmp
        return idx

furthestsampling = FurthestSampling.apply


class KNNQuery(Function):
    @staticmethod
    def forward(ctx, nsample, xyz, new_xyz, offset, new_offset):
        """
        input: xyz: (n, 3), new_xyz: (m, 3), offset: (b), new_offset: (b)
        output: idx: (m, nsample), dist2: (m, nsample)
        """
        if new_xyz is None: new_xyz = xyz
        assert xyz.is_contiguous() and new_xyz.is_contiguous()
        m = new_xyz.shape[0]
        idx = torch.cuda.IntTensor(m, nsample).zero_()
        dist2 = torch.cuda.FloatTensor(m, nsample).zero_()
        pointops_cuda.knnquery_cuda(m, nsample, xyz, new_xyz, offset, new_offset, idx, dist2)
        return idx, torch.sqrt(dist2)

knnquery = KNNQuery.apply


class Grouping(Function):
    @staticmethod
    def forward(ctx, input, idx):
        """
        input: input: (n, c), idx : (m, nsample)
        output: (m, nsample, c)
        """
        assert input.is_contiguous() and idx.is_contiguous()
        m, nsample, n, c = idx.shape[0], idx.shape[1], input.shape[0], input.shape[1]
        output = torch.cuda.FloatTensor(m, nsample, c)
        pointops_cuda.grouping_forward_cuda(m, nsample, c, input, idx, output)
        ctx.n = n
        ctx.save_for_backward(idx)
        return output

    @staticmethod
    def backward(ctx, grad_output):
        """
        input: grad_out: (m, c, nsample)
        output: (n, c), None
        """
        n = ctx.n
        idx, = ctx.saved_tensors
        m, nsample, c = grad_output.shape
        grad_input = torch.cuda.FloatTensor(n, c).zero_()
        pointops_cuda.grouping_backward_cuda(m, nsample, c, grad_output, idx, grad_input)
        return grad_input, None

grouping = Grouping.apply

class AttentionStep1(Function):
    @staticmethod
    def forward(ctx, q, k, index0, index1):
        """
        input: q: (N, h, C//h), k: (N, h, C//h), index0: (M), index1: (M)
        output: output: [N, h, C//h]
        """
        assert q.is_contiguous() and k.is_contiguous() and index0.is_contiguous() and index1.is_contiguous()

        N_q, h, C_div_h = q.shape
        N_k = k.shape[0]
        M = index0.shape[0]
        C = int(C_div_h * h)

        output = torch.cuda.FloatTensor(M, h).zero_()
        pointops_cuda.attention_step1_forward_cuda(N_k, M, h, C, q, k, index0, index1, output)
        ctx.N_q = N_q
        ctx.N_k = N_k
        ctx.C = C
        ctx.save_for_backward(q, k, index0, index1)
        return output

    @staticmethod
    def backward(ctx, grad_output):
        """
        input: grad_output: (N, h, C//h)
        output: (M, h), (N, h, C//h), None, None
        """
        
        N_q = ctx.N_q
        N_k = ctx.N_k
        C = ctx.C
        q, k, index0, index1 = ctx.saved_tensors
        M, h = grad_output.shape
        
        grad_output = grad_output.contiguous()
        # print("grad_output.is_contiguous(): ", grad_output.is_contiguous())
        assert q.is_contiguous() and k.is_contiguous() and index0.is_contiguous() and index1.is_contiguous() and grad_output.is_contiguous()

        # print("back: attn[:5,:5]: ", attn[:5, :5])

        # print("attn.shape: {} v.shape: {}, index0.shape: {}, index1.shape: {}".format(attn.shape, v.shape, index0.shape, index1.shape))

        grad_q = torch.cuda.FloatTensor(N_q, h, C//h).zero_()
        grad_k = torch.cuda.FloatTensor(N_k, h, C//h).zero_()

        # torch.cuda.synchronize()
        # start = time.time()
        
        pointops_cuda.attention_step1_backward_cuda(N_q, M, h, C, grad_output, index0, index1, q, k, grad_q, grad_k)
        
        # torch.cuda.synchronize()
        # end = time.time()
        # print("time v7: {}".format(end - start))
        # # input()
        
        return grad_q, grad_k, None, None

attention_step1 = AttentionStep1.apply

class AttentionStep1_v2(Function):
    @staticmethod
    def forward(ctx, q, k, index1, index0_offsets, n_max):
        """
        input: q: (N, h, C//h), k: (N, h, C//h), index0: (M), index1: (M)
        output: output: [N, h, C//h]
        """
        assert q.is_contiguous() and k.is_contiguous() and index0_offsets.is_contiguous() and index1.is_contiguous()
        assert n_max <= 1024

        N_q, h, C_div_h = q.shape
        N_k = k.shape[0]
        M = index1.shape[0]
        C = int(C_div_h * h)

        output = torch.cuda.FloatTensor(M, h).zero_()
        pointops_cuda.attention_step1_forward_cuda_v2(N_k, M, h, C, n_max, q, k, index0_offsets, index1, output)
        ctx.N_q = N_q
        ctx.N_k = N_k
        ctx.C = C
        ctx.n_max = n_max
        ctx.save_for_backward(q, k, index0_offsets, index1)
        return output

    @staticmethod
    def backward(ctx, grad_output):
        """
        input: grad_output: (N, h, C//h)
        output: (M, h), (N, h, C//h), None, None
        """
        
        N_q = ctx.N_q
        N_k = ctx.N_k
        C = ctx.C
        n_max = ctx.n_max
        q, k, index0_offsets, index1 = ctx.saved_tensors
        M, h = grad_output.shape
        
        grad_output = grad_output.contiguous()
        # print("grad_output.is_contiguous(): ", grad_output.is_contiguous())
        assert q.is_contiguous() and k.is_contiguous() and index0_offsets.is_contiguous() and index1.is_contiguous() and grad_output.is_contiguous()

        # print("back: attn[:5,:5]: ", attn[:5, :5])

        # print("attn.shape: {} v.shape: {}, index0.shape: {}, index1.shape: {}".format(attn.shape, v.shape, index0.shape, index1.shape))

        grad_q = torch.cuda.FloatTensor(N_q, h, C//h).zero_()
        grad_k = torch.cuda.FloatTensor(N_k, h, C//h).zero_()

        # torch.cuda.synchronize()
        # start = time.time()
        
        pointops_cuda.attention_step1_backward_cuda_v2(N_q, M, h, C, n_max, grad_output, index0_offsets, index1, q, k, grad_q, grad_k)
        
        # torch.cuda.synchronize()
        # end = time.time()
        # print("time v7: {}".format(end - start))
        # # input()
        
        return grad_q, grad_k, None, None, None

attention_step1_v2 = AttentionStep1_v2.apply



class AttentionStep2(Function):
    @staticmethod
    def forward(ctx, attn, v, index0, index1):
        """
        input: attn: (M, h), v: (N, h, C//h), index0: (M), index1: (M)
        output: output: [N, h, C//h]
        """
        assert attn.is_contiguous() and v.is_contiguous() and index0.is_contiguous() and index1.is_contiguous()

        M, h = attn.shape
        N_q = index0.max().item() + 1
        N_v, h, C_div_h = v.shape
        C = int(C_div_h * h)

        output = torch.cuda.FloatTensor(N_q, h, C//h).zero_()
        pointops_cuda.attention_step2_forward_cuda(N_q, M, h, C, attn, v, index0, index1, output)
        ctx.M = M

        # print("attn[:5,:5]: ", attn[:5, :5])

        ctx.save_for_backward(attn, v, index0, index1)
        return output

    @staticmethod
    def backward(ctx, grad_output):
        """
        input: grad_output: (N, h, C//h)
        output: (M, h), (N, h, C//h), None, None
        """
        M = ctx.M
        attn, v, index0, index1 = ctx.saved_tensors
        N_v = v.shape[0]
        N_q, h, C_div_h = grad_output.shape
        C = h * C_div_h
        
        grad_output = grad_output.contiguous()
        # print("grad_output.is_contiguous(): ", grad_output.is_contiguous())
        assert attn.is_contiguous() and v.is_contiguous() and index0.is_contiguous() and index1.is_contiguous() and grad_output.is_contiguous()

        # print("back: attn[:5,:5]: ", attn[:5, :5])

        # print("attn.shape: {} v.shape: {}, index0.shape: {}, index1.shape: {}".format(attn.shape, v.shape, index0.shape, index1.shape))

        grad_attn = torch.cuda.FloatTensor(M, h).zero_()
        grad_v = torch.cuda.FloatTensor(N_v, h, C//h).zero_()

        # torch.cuda.synchronize()
        # start = time.time()
        
        pointops_cuda.attention_step2_backward_cuda(N_q, M, h, C, grad_output, index0, index1, attn, v, grad_attn, grad_v)
        
        # torch.cuda.synchronize()
        # end = time.time()
        # print("time v8: {}".format(end - start))
        # # input()
        
        return grad_attn, grad_v, None, None

attention_step2 = AttentionStep2.apply


class AttentionStep2_v2(Function):
    @staticmethod
    def forward(ctx, attn, v, index0, index1):
        """
        input: attn: (M, h), v: (N, h, C//h), index0: (M), index1: (M)
        output: output: [L, h, C//h]
        """
        assert attn.is_contiguous() and v.is_contiguous() and index0.is_contiguous() and index1.is_contiguous()
        
        L = int(index0.max().item()) + 1

        M, h = attn.shape
        N, h, C_div_h = v.shape
        C = int(C_div_h * h)

        output = torch.cuda.FloatTensor(L, h, C//h).zero_()
        pointops_cuda.attention_step2_forward_cuda(N, M, h, C, attn, v, index0, index1, output)
        ctx.M = M

        # print("attn[:5,:5]: ", attn[:5, :5])

        ctx.save_for_backward(attn, v, index0, index1)
        return output

    @staticmethod
    def backward(ctx, grad_output):
        """
        input: grad_output: (L, h, C//h)
        output: (M, h), (N, h, C//h), None, None
        """
        M = ctx.M
        attn, v, index0, index1 = ctx.saved_tensors
        L, h, C_div_h = grad_output.shape
        N = v.shape[0]
        C = h * C_div_h
        
        grad_output = grad_output.contiguous()
        # print("grad_output.is_contiguous(): ", grad_output.is_contiguous())
        assert attn.is_contiguous() and v.is_contiguous() and index0.is_contiguous() and index1.is_contiguous() and grad_output.is_contiguous()

        # print("back: attn[:5,:5]: ", attn[:5, :5])

        # print("attn.shape: {} v.shape: {}, index0.shape: {}, index1.shape: {}".format(attn.shape, v.shape, index0.shape, index1.shape))

        grad_attn = torch.cuda.FloatTensor(M, h).zero_()
        grad_v = torch.cuda.FloatTensor(N, h, C//h).zero_()

        pointops_cuda.attention_step2_backward_cuda(N, M, h, C, grad_output, index0, index1, attn, v, grad_attn, grad_v)
        return grad_attn, grad_v, None, None

attention_step2_v2 = AttentionStep2_v2.apply

class DotProdWithIdx(Function):
    @staticmethod
    def forward(ctx, q, index, table, rel_idx):
        """
        input: q: (N, h, hdim), index: (M), table: (L, h, hdim, 3), rel_idx: (M, 3)
        output: output: [M, h]
        """
        assert q.is_contiguous() and index.is_contiguous() and table.is_contiguous() and rel_idx.is_contiguous()

        N, h, hdim = q.shape
        M = index.shape[0]

        output = torch.cuda.FloatTensor(M, h).zero_()
        pointops_cuda.dot_prod_with_idx_forward_cuda(N, M, h, hdim, q, index, table, rel_idx, output)
        ctx.save_for_backward(q, index, table, rel_idx)
        return output

    @staticmethod
    def backward(ctx, grad_output):
        """
        input: grad_output: [M, h]
        output: (N, h, hdim), None, (L, h, hdim, 3), None
        """
        q, index, table, rel_idx = ctx.saved_tensors
        M, h = grad_output.shape
        N, _, hdim = q.shape
        L = table.shape[0]
        
        grad_output = grad_output.contiguous()
        assert q.is_contiguous() and index.is_contiguous() and table.is_contiguous() and rel_idx.is_contiguous() and grad_output.is_contiguous()

        # print("back: attn[:5,:5]: ", attn[:5, :5])

        # print("attn.shape: {} v.shape: {}, index0.shape: {}, index1.shape: {}".format(attn.shape, v.shape, index0.shape, index1.shape))

        grad_q = torch.cuda.FloatTensor(N, h, hdim).zero_()
        grad_table = torch.cuda.FloatTensor(L, h, hdim, 3).zero_()

        # torch.cuda.synchronize()
        # start = time.time()
        
        pointops_cuda.dot_prod_with_idx_backward_cuda(N, M, h, hdim, grad_output, q, index, table, rel_idx, grad_q, grad_table)
        
        # torch.cuda.synchronize()
        # end = time.time()
        # print("time v9: {}".format(end - start))
        # # input()
        
        return grad_q, None, grad_table, None

dot_prod_with_idx = DotProdWithIdx.apply

class DotProdWithIdx_v2(Function):
    @staticmethod
    def forward(ctx, q, index_q, k, index_k, table_q, table_k, rel_idx):
        """
        input: q: (N, h, hdim), index_q: (M), k: (N, h, hdim), index_k: (M), table_q: (L, h, hdim, 3), table_k: (L, h, hdim, 3), rel_idx: (M, 3)
        output: output: [M, h]
        """
        assert q.is_contiguous() and index_q.is_contiguous() and k.is_contiguous() and index_k.is_contiguous() and table_q.is_contiguous() and table_k.is_contiguous() and rel_idx.is_contiguous()

        N, h, hdim = q.shape
        M = index_q.shape[0]
        L = table_q.shape[0]
        assert table_k.shape[0] == L and index_k.shape[0] == M

        # obtain the mapping from block_idx to m_idx
        rel_idx_merge = rel_idx[:, 0] + rel_idx[:, 1] * L + rel_idx[:, 2] * (L ** 2) #[M, ]
        sorted_values, sort_indices = torch.sort(rel_idx_merge)
        _, counts = torch.unique_consecutive(sorted_values, return_counts=True)
        rel_idx_offsets = torch.cumsum(counts, dim=-1) #[T,]
        rel_idx_offsets = torch.cat([torch.zeros(1, dtype=torch.long).cuda(), rel_idx_offsets], 0) #[T+1,]
        n_max = counts.max()
        T = counts.shape[0]

        # print("M: {}, L: {}, n_max: {}, T: {}".format(M, L, n_max, T))
        # print("rel_idx_merge.shape: {}, sorted_values.shape: {}".format(rel_idx_merge.shape, sorted_values.shape))
        # print("counts.shape: {}".format(counts.shape))

        output = torch.cuda.FloatTensor(M, h).zero_()
        # pointops_cuda.dot_prod_with_idx_forward_cuda(N, M, h, hdim, q, index, table, rel_idx, output)
        pointops_cuda.dot_prod_with_idx_forward_cuda_v2(N, M, h, hdim, n_max, T, q, index_q, k, index_k, table_q, table_k, rel_idx, rel_idx_offsets.int(), sort_indices.int(), output)
        
        ctx.n_max = n_max
        ctx.T = T
        ctx.save_for_backward(q, index_q, k, index_k, table_q, table_k, rel_idx, rel_idx_offsets, sort_indices)
        return output

    @staticmethod
    def backward(ctx, grad_output):
        """
        input: grad_output: [M, h]
        output: (N, h, hdim), None, (L, h, hdim, 3), None
        """
        q, index_q, k, index_k, table_q, table_k, rel_idx, rel_idx_offsets, sort_indices = ctx.saved_tensors
        M, h = grad_output.shape
        N, _, hdim = q.shape
        L = table_q.shape[0]
        T, n_max = ctx.T, ctx.n_max
        
        grad_output = grad_output.contiguous()
        assert q.is_contiguous() and index_q.is_contiguous() and k.is_contiguous() and index_k.is_contiguous() and table_q.is_contiguous() and table_k.is_contiguous() and rel_idx.is_contiguous() and rel_idx_offsets.is_contiguous() and sort_indices.is_contiguous() and grad_output.is_contiguous()

        # print("back: attn[:5,:5]: ", attn[:5, :5])

        # print("attn.shape: {} v.shape: {}, index0.shape: {}, index1.shape: {}".format(attn.shape, v.shape, index0.shape, index1.shape))

        grad_q = torch.cuda.FloatTensor(N, h, hdim).zero_()
        grad_table_q = torch.cuda.FloatTensor(L, h, hdim, 3).zero_()
        grad_k = torch.cuda.FloatTensor(N, h, hdim).zero_()
        grad_table_k = torch.cuda.FloatTensor(L, h, hdim, 3).zero_()

        # torch.cuda.synchronize()
        # start = time.time()
        
        pointops_cuda.dot_prod_with_idx_backward_cuda_v2(N, M, h, hdim, n_max, T, grad_output, q, index_q, k, index_k, table_q, table_k, rel_idx, rel_idx_offsets.int(), sort_indices.int(), grad_q, grad_k, grad_table_q, grad_table_k)
        
        # torch.cuda.synchronize()
        # end = time.time()
        # print("time v9: {}".format(end - start))
        # # input()
        return grad_q, None, grad_k, None, grad_table_q, grad_table_k, None

dot_prod_with_idx_v2 = DotProdWithIdx_v2.apply


class DotProdWithIdx_v3(Function):
    @staticmethod
    def forward(ctx, q, index_q_offsets, n_max, k, index_k, table_q, table_k, rel_idx):
        """
        input: q: (N, h, hdim), index_q: (M), k: (N, h, hdim), index_k: (M), table_q: (L, h, hdim, 3), table_k: (L, h, hdim, 3), rel_idx: (M, 3)
        output: output: [M, h]
        """
        assert q.is_contiguous() and index_q_offsets.is_contiguous() and k.is_contiguous() and index_k.is_contiguous() and table_q.is_contiguous() and table_k.is_contiguous() and rel_idx.is_contiguous()

        N, h, hdim = q.shape
        M = index_k.shape[0]
        L = table_q.shape[0]
        assert table_k.shape[0] == L

        # # obtain the mapping from block_idx to m_idx
        # rel_idx_merge = rel_idx[:, 0] + rel_idx[:, 1] * L + rel_idx[:, 2] * (L ** 2) #[M, ]
        # sorted_values, sort_indices = torch.sort(rel_idx_merge)
        # _, counts = torch.unique_consecutive(sorted_values, return_counts=True)
        # rel_idx_offsets = torch.cumsum(counts, dim=-1) #[T,]
        # rel_idx_offsets = torch.cat([torch.zeros(1, dtype=torch.long).cuda(), rel_idx_offsets], 0) #[T+1,]
        # n_max = counts.max()
        # T = counts.shape[0]

        # print("M: {}, L: {}, n_max: {}, T: {}".format(M, L, n_max, T))
        # print("rel_idx_merge.shape: {}, sorted_values.shape: {}".format(rel_idx_merge.shape, sorted_values.shape))
        # print("counts.shape: {}".format(counts.shape))

        # print("M: {}, L: {}, n_max: {}".format(M, L, n_max))

        output = torch.cuda.FloatTensor(M, h).zero_()
        # pointops_cuda.dot_prod_with_idx_forward_cuda(N, M, h, hdim, q, index, table, rel_idx, output)
        pointops_cuda.dot_prod_with_idx_forward_cuda_v3(N, M, h, hdim, n_max, q, index_q_offsets, k, index_k, table_q, table_k, rel_idx, output)
        
        ctx.n_max = n_max
        # ctx.T = T
        ctx.save_for_backward(q, index_q_offsets, k, index_k, table_q, table_k, rel_idx)
        return output

    @staticmethod
    def backward(ctx, grad_output):
        """
        input: grad_output: [M, h]
        output: (N, h, hdim), None, (L, h, hdim, 3), None
        """
        q, index_q_offsets, k, index_k, table_q, table_k, rel_idx = ctx.saved_tensors
        M, h = grad_output.shape
        N, _, hdim = q.shape
        L = table_q.shape[0]
        n_max = ctx.n_max
        
        grad_output = grad_output.contiguous()
        assert q.is_contiguous() and index_q_offsets.is_contiguous() and k.is_contiguous() and index_k.is_contiguous() and table_q.is_contiguous() and table_k.is_contiguous() and rel_idx.is_contiguous() and grad_output.is_contiguous()

        # print("back: attn[:5,:5]: ", attn[:5, :5])

        # print("attn.shape: {} v.shape: {}, index0.shape: {}, index1.shape: {}".format(attn.shape, v.shape, index0.shape, index1.shape))

        grad_q = torch.cuda.FloatTensor(N, h, hdim).zero_()
        grad_table_q = torch.cuda.FloatTensor(L, h, hdim, 3).zero_()
        grad_k = torch.cuda.FloatTensor(N, h, hdim).zero_()
        grad_table_k = torch.cuda.FloatTensor(L, h, hdim, 3).zero_()

        # torch.cuda.synchronize()
        # start = time.time()
        
        pointops_cuda.dot_prod_with_idx_backward_cuda_v3(N, M, h, hdim, n_max, grad_output, q, index_q_offsets, k, index_k, table_q, table_k, rel_idx, grad_q, grad_k, grad_table_q, grad_table_k)
        
        # torch.cuda.synchronize()
        # end = time.time()
        # print("time v9: {}".format(end - start))
        # # input()
        return grad_q, None, None, grad_k, None, grad_table_q, grad_table_k, None

dot_prod_with_idx_v3 = DotProdWithIdx_v3.apply

class AttentionStep2WithRelPosValue(Function):
    @staticmethod
    def forward(ctx, attn, v, index0, index1, table, rel_idx):
        """
        input: attn: (M, h), v: (N, h, hdim), index0: (M), index1: (M), table: (L, h, hdim, 3), rel_idx: (M, 3)
        output: output: [N, h, hdim]
        """
        assert attn.is_contiguous() and v.is_contiguous() and index0.is_contiguous() and index1.is_contiguous() and table.is_contiguous() and rel_idx.is_contiguous()

        M, h = attn.shape
        N_v, h, hdim = v.shape
        N_q = index0.max().item() + 1

        output = torch.cuda.FloatTensor(N_q, h, hdim).zero_()
        pointops_cuda.attention_step2_with_rel_pos_value_forward_cuda(N_q, M, h, hdim, attn, v, index0, index1, table, rel_idx, output)

        # print("attn[:5,:5]: ", attn[:5, :5])

        ctx.save_for_backward(attn, v, index0, index1, table, rel_idx)
        return output

    @staticmethod
    def backward(ctx, grad_output):
        """
        input: grad_output: (N, h, C//h)
        output: (M, h), (N, h, C//h), None, None, (L, h, hdim, 3), None
        """
        attn, v, index0, index1, table, rel_idx = ctx.saved_tensors
        N_q, h, hdim = grad_output.shape
        N_v = v.shape[0]
        M = attn.shape[0]
        L = table.shape[0]

        grad_output = grad_output.contiguous()
        # print("grad_output.is_contiguous(): ", grad_output.is_contiguous())
        assert attn.is_contiguous() and v.is_contiguous() and index0.is_contiguous() and index1.is_contiguous() and grad_output.is_contiguous() and table.is_contiguous() and rel_idx.is_contiguous()

        # print("back: attn[:5,:5]: ", attn[:5, :5])

        # print("attn.shape: {} v.shape: {}, index0.shape: {}, index1.shape: {}".format(attn.shape, v.shape, index0.shape, index1.shape))

        grad_attn = torch.cuda.FloatTensor(M, h).zero_()
        grad_v = torch.cuda.FloatTensor(N_v, h, hdim).zero_()
        grad_table = torch.cuda.FloatTensor(L, h, hdim, 3).zero_()

        # print("attn.shape: {}, grad_attn.shape: {}".format(attn.shape, grad_attn.shape))
        # print("v.shape: {}, grad_v.shape: {}".format(v.shape, grad_v.shape))
        # print("table.shape: {}, grad_table.shape: {}".format(table.shape, grad_table.shape))

        # torch.cuda.synchronize()
        # start = time.time()
        
        pointops_cuda.attention_step2_with_rel_pos_value_backward_cuda(N_q, M, h, hdim, grad_output, index0, index1, attn, v, table, rel_idx, grad_attn, grad_v, grad_table)
        
        # torch.cuda.synchronize()
        # end = time.time()
        # print("time v10: {}".format(end - start))
        # # input()
        return grad_attn, grad_v, None, None, grad_table, None

attention_step2_with_rel_pos_value = AttentionStep2WithRelPosValue.apply


class AttentionStep2WithRelPosValue_v2(Function):
    @staticmethod
    def forward(ctx, attn, v, index0_offsets, n_max, index1, table, rel_idx):
        """
        input: attn: (M, h), v: (N, h, hdim), index0_offsets: (M), index1: (M), table: (L, h, hdim, 3), rel_idx: (M, 3)
        output: output: [N, h, hdim]
        """
        assert attn.is_contiguous() and v.is_contiguous() and index0_offsets.is_contiguous() and index1.is_contiguous() and table.is_contiguous() and rel_idx.is_contiguous()

        M, h = attn.shape
        N, h, hdim = v.shape
        # N_q = int(index0_offsets.max().item()) + 1

        output = torch.cuda.FloatTensor(N, h, hdim).zero_()
        pointops_cuda.attention_step2_with_rel_pos_value_forward_cuda_v2(N, M, h, hdim, n_max, attn, v, index0_offsets, index1, table, rel_idx, output)

        # print("attn[:5,:5]: ", attn[:5, :5])

        ctx.n_max = n_max
        ctx.save_for_backward(attn, v, index0_offsets, index1, table, rel_idx)
        return output

    @staticmethod
    def backward(ctx, grad_output):
        """
        input: grad_output: (N, h, C//h)
        output: (M, h), (N, h, C//h), None, None, (L, h, hdim, 3), None
        """
        n_max = ctx.n_max
        attn, v, index0_offsets, index1, table, rel_idx = ctx.saved_tensors
        N, h, hdim = grad_output.shape
        N = v.shape[0]
        M = attn.shape[0]
        L = table.shape[0]

        # grad_output = grad_output.contiguous()
        # print("grad_output.is_contiguous(): ", grad_output.is_contiguous())
        assert attn.is_contiguous() and v.is_contiguous() and index0_offsets.is_contiguous() and index1.is_contiguous() and grad_output.is_contiguous() and table.is_contiguous() and rel_idx.is_contiguous()

        # print("back: attn[:5,:5]: ", attn[:5, :5])

        # print("attn.shape: {} v.shape: {}, index0_offsets.shape: {}, index1.shape: {}".format(attn.shape, v.shape, index0_offsets.shape, index1.shape))

        grad_attn = torch.cuda.FloatTensor(M, h).zero_()
        grad_v = torch.cuda.FloatTensor(N, h, hdim).zero_()
        grad_table = torch.cuda.FloatTensor(L, h, hdim, 3).zero_()

        # print("attn.shape: {}, grad_attn.shape: {}".format(attn.shape, grad_attn.shape))
        # print("v.shape: {}, grad_v.shape: {}".format(v.shape, grad_v.shape))
        # print("table.shape: {}, grad_table.shape: {}".format(table.shape, grad_table.shape))

        # torch.cuda.synchronize()
        # start = time.time()
        
        pointops_cuda.attention_step2_with_rel_pos_value_backward_cuda_v2(N, M, h, hdim, n_max, grad_output, index0_offsets, index1, attn, v, table, rel_idx, grad_attn, grad_v, grad_table)
        
        # torch.cuda.synchronize()
        # end = time.time()
        # print("time v10: {}".format(end - start))
        
        return grad_attn, grad_v, None, None, None, grad_table, None

attention_step2_with_rel_pos_value_v2 = AttentionStep2WithRelPosValue_v2.apply

def queryandgroup(nsample, xyz, new_xyz, feat, idx, offset, new_offset, use_xyz=True, return_indx=False):
    """
    input: xyz: (n, 3), new_xyz: (m, 3), feat: (n, c), idx: (m, nsample), offset: (b), new_offset: (b)
    output: new_feat: (m, c+3, nsample), grouped_idx: (m, nsample)
    """
    assert xyz.is_contiguous() and new_xyz.is_contiguous() and feat.is_contiguous()
    if new_xyz is None:
        new_xyz = xyz
    if idx is None:
        idx, _ = knnquery(nsample, xyz, new_xyz, offset, new_offset) # (m, nsample)

    n, m, c = xyz.shape[0], new_xyz.shape[0], feat.shape[1]
    grouped_xyz = xyz[idx.view(-1).long(), :].view(m, nsample, 3) # (m, nsample, 3)
    #grouped_xyz = grouping(xyz, idx) # (m, nsample, 3)
    # 相对位置
    grouped_xyz -= new_xyz.unsqueeze(1) # (m, nsample, 3)
    grouped_feat = feat[idx.view(-1).long(), :].view(m, nsample, c) # (m, nsample, c)
    #grouped_feat = grouping(feat, idx) # (m, nsample, c)
    if use_xyz:
        if return_indx:
            return torch.cat((grouped_xyz, grouped_feat), -1), idx # (m, nsample, 3+c)
        else:
            return torch.cat((grouped_xyz, grouped_feat), -1)
    else:
        if return_indx:
            return grouped_feat, idx
        else:
            return grouped_feat


def Divide2Patch(nsample, xyz, offset, return_offset=False, anchor_scale=None):
    # nsample: 16  xyz: (n, 3)  offset: (b)
    downsample_scale = anchor_scale or nsample
    new_offset, count = [offset[0].item() // downsample_scale], offset[0].item() // downsample_scale
    for i in range(1, offset.shape[0]):
        count += (offset[i].item() - offset[i-1].item()) // downsample_scale
        new_offset.append(count)
    # print("donw sample scale:", downsample_scale,"offset:", offset, "newoffset:", new_offset)
    new_offset = torch.cuda.IntTensor(new_offset)
    idx = furthestsampling(xyz, offset, new_offset) # (m)
    new_xyz = xyz[idx.long()]
    p_idx, _ = knnquery(nsample, xyz, new_xyz, offset, new_offset) # (m, nsample)
    if return_offset:
        return p_idx, new_offset
    else:
        return p_idx
    
class Subtraction(Function):
    @staticmethod
    def forward(ctx, input1, input2, idx):
        """
        input: input1: (n, c), input2: (n, c), idx: (n, nsample)
        output:  (n, nsample, c)
        """
        assert input1.is_contiguous() and input2.is_contiguous()
        n, c = input1.shape; nsample = idx.shape[-1]
        output = torch.cuda.FloatTensor(n, nsample, c).zero_()
        pointops_cuda.subtraction_forward_cuda(n, nsample, c, input1, input2, idx, output)
        ctx.save_for_backward(idx)
        return output

    @staticmethod
    def backward(ctx, grad_output):
        """
        input: grad_out: (n, nsample, c)
        output: grad_input1: (n, c), grad_input2: (n, c)
        """
        idx, = ctx.saved_tensors
        n, nsample, c = grad_output.shape
        grad_input1 = torch.cuda.FloatTensor(n, c).zero_()
        grad_input2 = torch.cuda.FloatTensor(n, c).zero_()
        pointops_cuda.subtraction_backward_cuda(n, nsample, c, idx, grad_output, grad_input1, grad_input2)
        return grad_input1, grad_input2, None

subtraction = Subtraction.apply


class Aggregation(Function):
    @staticmethod
    def forward(ctx, input, position, weight, idx):
        """
        input: input: (n, c), position: (n, nsample, c), weight : (n, nsample, c'), idx: (n, nsample)
        output: (n, c)
        """
        assert input.is_contiguous() and position.is_contiguous() and weight.is_contiguous()
        n, nsample, c = position.shape; w_c = weight.shape[-1]
        output = torch.cuda.FloatTensor(n, c).zero_()
        pointops_cuda.aggregation_forward_cuda(n, nsample, c, w_c, input, position, weight, idx, output)
        ctx.save_for_backward(input, position, weight, idx)
        return output

    @staticmethod
    def backward(ctx, grad_output):
        """
        input: grad_out: (n, c)
        output: grad_input: (n, c), grad_position: (n, nsample, c), grad_weight : (n, nsample, c')
        """
        input, position, weight, idx = ctx.saved_tensors
        n, nsample, c = position.shape; w_c = weight.shape[-1]
        grad_input = torch.cuda.FloatTensor(n, c).zero_()
        grad_position = torch.cuda.FloatTensor(n, nsample, c).zero_()
        grad_weight = torch.cuda.FloatTensor(n, nsample, w_c).zero_()
        pointops_cuda.aggregation_backward_cuda(n, nsample, c, w_c, input, position, weight, idx, grad_output, grad_input, grad_position, grad_weight)
        return grad_input, grad_position, grad_weight, None

aggregation = Aggregation.apply


def interpolation(xyz, new_xyz, feat, offset, new_offset, k=3):
    """
    input: xyz: (m, 3), new_xyz: (n, 3), feat: (m, c), offset: (b), new_offset: (b)
    output: (n, c)
    """
    assert xyz.is_contiguous() and new_xyz.is_contiguous() and feat.is_contiguous()
    idx, dist = knnquery(k, xyz, new_xyz, offset, new_offset) # (n, 3), (n, 3)
    dist_recip = 1.0 / (dist + 1e-8) # (n, 3)
    norm = torch.sum(dist_recip, dim=1, keepdim=True)
    weight = dist_recip / norm # (n, 3)

    new_feat = torch.cuda.FloatTensor(new_xyz.shape[0], feat.shape[1]).zero_()
    for i in range(k):
        new_feat += feat[idx[:, i].long(), :] * weight[:, i].unsqueeze(-1)
    return new_feat


def interpolation_v2(xyz, new_xyz, feat, offset, new_offset, k=3):
    """
    input: xyz: (m, 3), new_xyz: (n, 3), feat: (m, c), offset: (b), new_offset: (b)
    output: (n, c)
    """
    assert xyz.is_contiguous() and new_xyz.is_contiguous() and feat.is_contiguous()

    idx, _ = knnquery(k, xyz, new_xyz, offset, new_offset) # (n, 3), (n, 3)

    # print("e3: idx.shape: {}, idx[:5]: {}".format(idx.shape, idx[:5]))

    dist = torch.sqrt(((new_xyz.unsqueeze(1) - xyz[idx.long()]) ** 2).sum(-1) + 1e-8)

    # print("e4: dist.shape: {}, dist[:5]: {}".format(dist.shape, dist[:5]))
    # print("((_-dist)**2).max(): ", ((_-dist)**2).max())
    # input()

    dist_recip = 1.0 / (dist + 1e-8) # (n, 3)
    norm = torch.sum(dist_recip, dim=1, keepdim=True)
    weight = dist_recip / norm # (n, 3)

    new_feat = torch.cuda.FloatTensor(new_xyz.shape[0], feat.shape[1]).zero_()
    for i in range(k):
        new_feat += feat[idx[:, i].long(), :] * weight[:, i].unsqueeze(-1)
    return new_feat


class Interpolation(Function):
    @staticmethod
    def forward(ctx, xyz, new_xyz, input, offset, new_offset, k=3):
        """
        input: xyz: (m, 3), new_xyz: (n, 3), input: (m, c), offset: (b), new_offset: (b)
        output: (n, c)
        """
        assert xyz.is_contiguous() and new_xyz.is_contiguous() and input.is_contiguous()
        idx, dist = knnquery(k, xyz, new_xyz, offset, new_offset) # (n, k), (n, k)
        dist_recip = 1.0 / (dist + 1e-8) # (n, k)
        norm = torch.sum(dist_recip, dim=1, keepdim=True)
        weight = dist_recip / norm # (n, k)

        n, c, m = new_xyz.shape[0], input.shape[1], input.shape[0]
        output = torch.cuda.FloatTensor(n, c).zero_()
        pointops_cuda.interpolation_forward_cuda(n, c, k, input, idx, weight, output)
        ctx.m, ctx.k = m, k
        ctx.save_for_backward(idx, weight)
        return output

    @staticmethod
    def backward(ctx, grad_output):
        """
        input: xyz: (m, 3), new_xyz: (n, 3), input: (m, c), offset: (b), new_offset: (b)
        output: (n, c)
        """
        m, k = ctx.m, ctx.k
        idx, weight = ctx.saved_tensors
        n, c = grad_output.shape
        grad_input = torch.cuda.FloatTensor(m, c).zero_()
        pointops_cuda.interpolation_backward_cuda(n, c, k, grad_output, idx, weight, grad_input)
        return None, None, grad_input, None, None, None

interpolation2 = Interpolation.apply

utils/pointops2/functions/pointops.py
aminebdj/3D-OWIS
from typing import Tuple

import torch
from torch.autograd import Function
import torch.nn as nn

import pointops2_cuda as pointops_cuda


class FurthestSampling(Function):
    @staticmethod
    def forward(ctx, xyz, offset, new_offset):
        """
        input: xyz: (n, 3), offset: (b), new_offset: (b)
        output: idx: (m)
        """
        assert xyz.is_contiguous()
        n, b, n_max = xyz.shape[0], offset.shape[0], offset[0]
        for i in range(1, b):
            n_max = max(offset[i] - offset[i-1], n_max)
        idx = torch.cuda.IntTensor(new_offset[b-1].item()).zero_()
        tmp = torch.cuda.FloatTensor(n).fill_(1e10)
        pointops_cuda.furthestsampling_cuda(b, n_max, xyz, offset, new_offset, tmp, idx)
        del tmp
        return idx

furthestsampling = FurthestSampling.apply


class KNNQuery(Function):
    @staticmethod
    def forward(ctx, nsample, xyz, new_xyz, offset, new_offset):
        """
        input: xyz: (n, 3), new_xyz: (m, 3), offset: (b), new_offset: (b)
        output: idx: (m, nsample), dist2: (m, nsample)
        """
        if new_xyz is None: new_xyz = xyz
        assert xyz.is_contiguous() and new_xyz.is_contiguous()
        m = new_xyz.shape[0]
        idx = torch.cuda.IntTensor(m, nsample).zero_()
        dist2 = torch.cuda.FloatTensor(m, nsample).zero_()
        pointops_cuda.knnquery_cuda(m, nsample, xyz, new_xyz, offset, new_offset, idx, dist2)
        return idx, torch.sqrt(dist2)

knnquery = KNNQuery.apply


class Grouping(Function):
    @staticmethod
    def forward(ctx, input, idx):
        """
        input: input: (n, c), idx : (m, nsample)
        output: (m, nsample, c)
        """
        assert input.is_contiguous() and idx.is_contiguous()
        m, nsample, n, c = idx.shape[0], idx.shape[1], input.shape[0], input.shape[1]
        output = torch.cuda.FloatTensor(m, nsample, c)
        pointops_cuda.grouping_forward_cuda(m, nsample, c, input, idx, output)
        ctx.n = n
        ctx.save_for_backward(idx)
        return output

    @staticmethod
    def backward(ctx, grad_output):
        """
        input: grad_out: (m, c, nsample)
        output: (n, c), None
        """
        n = ctx.n
        idx, = ctx.saved_tensors
        m, nsample, c = grad_output.shape
        grad_input = torch.cuda.FloatTensor(n, c).zero_()
        pointops_cuda.grouping_backward_cuda(m, nsample, c, grad_output, idx, grad_input)
        return grad_input, None

grouping = Grouping.apply


def queryandgroup(nsample, xyz, new_xyz, feat, idx, offset, new_offset, use_xyz=True):
    """
    input: xyz: (n, 3), new_xyz: (m, 3), feat: (n, c), idx: (m, nsample), offset: (b), new_offset: (b)
    output: new_feat: (m, c+3, nsample), grouped_idx: (m, nsample)
    """
    assert xyz.is_contiguous() and new_xyz.is_contiguous() and feat.is_contiguous()
    if new_xyz is None:
        new_xyz = xyz
    if idx is None:
        idx, _ = knnquery(nsample, xyz, new_xyz, offset, new_offset) # (m, nsample)

    n, m, c = xyz.shape[0], new_xyz.shape[0], feat.shape[1]
    grouped_xyz = xyz[idx.view(-1).long(), :].view(m, nsample, 3) # (m, nsample, 3)
    #grouped_xyz = grouping(xyz, idx) # (m, nsample, 3)
    grouped_xyz -= new_xyz.unsqueeze(1) # (m, nsample, 3)
    grouped_feat = feat[idx.view(-1).long(), :].view(m, nsample, c) # (m, nsample, c)
    #grouped_feat = grouping(feat, idx) # (m, nsample, c)

    if use_xyz:
        return torch.cat((grouped_xyz, grouped_feat), -1) # (m, nsample, 3+c)
    else:
        return grouped_feat


class Subtraction(Function):
    @staticmethod
    def forward(ctx, input1, input2, idx):
        """
        input: input1: (n, c), input2: (n, c), idx: (n, nsample)
        output:  (n, nsample, c)
        """
        assert input1.is_contiguous() and input2.is_contiguous()
        n, c = input1.shape; nsample = idx.shape[-1]
        output = torch.cuda.FloatTensor(n, nsample, c).zero_()
        pointops_cuda.subtraction_forward_cuda(n, nsample, c, input1, input2, idx, output)
        ctx.save_for_backward(idx)
        return output

    @staticmethod
    def backward(ctx, grad_output):
        """
        input: grad_out: (n, nsample, c)
        output: grad_input1: (n, c), grad_input2: (n, c)
        """
        idx, = ctx.saved_tensors
        n, nsample, c = grad_output.shape
        grad_input1 = torch.cuda.FloatTensor(n, c).zero_()
        grad_input2 = torch.cuda.FloatTensor(n, c).zero_()
        pointops_cuda.subtraction_backward_cuda(n, nsample, c, idx, grad_output, grad_input1, grad_input2)
        return grad_input1, grad_input2, None

subtraction = Subtraction.apply


class Aggregation(Function):
    @staticmethod
    def forward(ctx, input, position, weight, idx):
        """
        input: input: (n, c), position: (n, nsample, c), weight : (n, nsample, c'), idx: (n, nsample)
        output: (n, c)
        """
        assert input.is_contiguous() and position.is_contiguous() and weight.is_contiguous()
        n, nsample, c = position.shape; w_c = weight.shape[-1]
        output = torch.cuda.FloatTensor(n, c).zero_()
        pointops_cuda.aggregation_forward_cuda(n, nsample, c, w_c, input, position, weight, idx, output)
        ctx.save_for_backward(input, position, weight, idx)
        return output

    @staticmethod
    def backward(ctx, grad_output):
        """
        input: grad_out: (n, c)
        output: grad_input: (n, c), grad_position: (n, nsample, c), grad_weight : (n, nsample, c')
        """
        input, position, weight, idx = ctx.saved_tensors
        n, nsample, c = position.shape; w_c = weight.shape[-1]
        grad_input = torch.cuda.FloatTensor(n, c).zero_()
        grad_position = torch.cuda.FloatTensor(n, nsample, c).zero_()
        grad_weight = torch.cuda.FloatTensor(n, nsample, w_c).zero_()
        pointops_cuda.aggregation_backward_cuda(n, nsample, c, w_c, input, position, weight, idx, grad_output, grad_input, grad_position, grad_weight)
        return grad_input, grad_position, grad_weight, None

aggregation = Aggregation.apply


def interpolation(xyz, new_xyz, feat, offset, new_offset, k=3):
    """
    input: xyz: (m, 3), new_xyz: (n, 3), feat: (m, c), offset: (b), new_offset: (b)
    output: (n, c)
    """
    assert xyz.is_contiguous() and new_xyz.is_contiguous() and feat.is_contiguous()
    idx, dist = knnquery(k, xyz, new_xyz, offset, new_offset) # (n, 3), (n, 3)
    dist_recip = 1.0 / (dist + 1e-8) # (n, 3)
    norm = torch.sum(dist_recip, dim=1, keepdim=True)
    weight = dist_recip / norm # (n, 3)

    new_feat = torch.cuda.FloatTensor(new_xyz.shape[0], feat.shape[1]).zero_()
    for i in range(k):
        new_feat += feat[idx[:, i].long(), :] * weight[:, i].unsqueeze(-1)
    return new_feat


class Interpolation(Function):
    @staticmethod
    def forward(ctx, xyz, new_xyz, input, offset, new_offset, k=3):
        """
        input: xyz: (m, 3), new_xyz: (n, 3), input: (m, c), offset: (b), new_offset: (b)
        output: (n, c)
        """
        assert xyz.is_contiguous() and new_xyz.is_contiguous() and input.is_contiguous()
        idx, dist = knnquery(k, xyz, new_xyz, offset, new_offset) # (n, k), (n, k)
        dist_recip = 1.0 / (dist + 1e-8) # (n, k)
        norm = torch.sum(dist_recip, dim=1, keepdim=True)
        weight = dist_recip / norm # (n, k)

        n, c, m = new_xyz.shape[0], input.shape[1], input.shape[0]
        output = torch.cuda.FloatTensor(n, c).zero_()
        pointops_cuda.interpolation_forward_cuda(n, c, k, input, idx, weight, output)
        ctx.m, ctx.k = m, k
        ctx.save_for_backward(idx, weight)
        return output

    @staticmethod
    def backward(ctx, grad_output):
        """
        input: xyz: (m, 3), new_xyz: (n, 3), input: (m, c), offset: (b), new_offset: (b)
        output: (n, c)
        """
        m, k = ctx.m, ctx.k
        idx, weight = ctx.saved_tensors
        n, c = grad_output.shape
        grad_input = torch.cuda.FloatTensor(m, c).zero_()
        pointops_cuda.interpolation_backward_cuda(n, c, k, grad_output, idx, weight, grad_input)
        return None, None, grad_input, None, None, None

interpolation2 = Interpolation.apply

utils/pointops2/functions/pointops2.py
aminebdj/3D-OWIS
from typing import Tuple

import torch
from torch.autograd import Function
import torch.nn as nn

import pointops2_cuda as pointops_cuda


class FurthestSampling(Function):
    @staticmethod
    def forward(ctx, xyz, offset, new_offset):
        """
        input: xyz: (n, 3), offset: (b), new_offset: (b)
        output: idx: (m)
        """
        assert xyz.is_contiguous()
        n, b, n_max = xyz.shape[0], offset.shape[0], offset[0]
        for i in range(1, b):
            n_max = max(offset[i] - offset[i-1], n_max)
        idx = torch.cuda.IntTensor(new_offset[b-1].item()).zero_()
        tmp = torch.cuda.FloatTensor(n).fill_(1e10)
        pointops_cuda.furthestsampling_cuda(b, n_max, xyz, offset, new_offset, tmp, idx)
        del tmp
        return idx

furthestsampling = FurthestSampling.apply


class KNNQuery(Function):
    @staticmethod
    def forward(ctx, nsample, xyz, new_xyz, offset, new_offset):
        """
        input: xyz: (n, 3), new_xyz: (m, 3), offset: (b), new_offset: (b)
        output: idx: (m, nsample), dist2: (m, nsample)
        """
        if new_xyz is None: new_xyz = xyz
        assert xyz.is_contiguous() and new_xyz.is_contiguous()
        m = new_xyz.shape[0]
        idx = torch.cuda.IntTensor(m, nsample).zero_()
        dist2 = torch.cuda.FloatTensor(m, nsample).zero_()
        pointops_cuda.knnquery_cuda(m, nsample, xyz, new_xyz, offset, new_offset, idx, dist2)
        return idx, torch.sqrt(dist2)

knnquery = KNNQuery.apply


class Grouping(Function):
    @staticmethod
    def forward(ctx, input, idx):
        """
        input: input: (n, c), idx : (m, nsample)
        output: (m, nsample, c)
        """
        assert input.is_contiguous() and idx.is_contiguous()
        m, nsample, n, c = idx.shape[0], idx.shape[1], input.shape[0], input.shape[1]
        output = torch.cuda.FloatTensor(m, nsample, c)
        pointops_cuda.grouping_forward_cuda(m, nsample, c, input, idx, output)
        ctx.n = n
        ctx.save_for_backward(idx)
        return output

    @staticmethod
    def backward(ctx, grad_output):
        """
        input: grad_out: (m, c, nsample)
        output: (n, c), None
        """
        n = ctx.n
        idx, = ctx.saved_tensors
        m, nsample, c = grad_output.shape
        grad_input = torch.cuda.FloatTensor(n, c).zero_()
        pointops_cuda.grouping_backward_cuda(m, nsample, c, grad_output, idx, grad_input)
        return grad_input, None

grouping = Grouping.apply


def queryandgroup(nsample, xyz, new_xyz, feat, idx, offset, new_offset, use_xyz=True, relative=True):
    """
    input: xyz: (n, 3), new_xyz: (m, 3), feat: (n, c), idx: (m, nsample), offset: (b), new_offset: (b)
    output: new_feat: (m, c+3, nsample), grouped_idx: (m, nsample)
    """
    assert xyz.is_contiguous() and new_xyz.is_contiguous() and feat.is_contiguous()
    if new_xyz is None:
        new_xyz = xyz
    if idx is None:
        idx, _ = knnquery(nsample, xyz, new_xyz, offset, new_offset) # (m, nsample)

    n, m, c = xyz.shape[0], new_xyz.shape[0], feat.shape[1]
    grouped_xyz = xyz[idx.view(-1).long(), :].view(m, nsample, 3) # (m, nsample, 3)
    #grouped_xyz = grouping(xyz, idx) # (m, nsample, 3)
    if relative:
        grouped_xyz -= new_xyz.unsqueeze(1) # (m, nsample, 3)
    grouped_feat = feat[idx.view(-1).long(), :].view(m, nsample, c) # (m, nsample, c)
    #grouped_feat = grouping(feat, idx) # (m, nsample, c)

    if use_xyz:
        return torch.cat((grouped_xyz, grouped_feat), -1) # (m, nsample, 3+c)
    else:
        return grouped_feat


class Subtraction(Function):
    @staticmethod
    def forward(ctx, input1, input2, idx):
        """
        input: input1: (n, c), input2: (n, c), idx: (n, nsample)
        output:  (n, nsample, c)
        """
        assert input1.is_contiguous() and input2.is_contiguous()
        n, c = input1.shape; nsample = idx.shape[-1]
        output = torch.cuda.FloatTensor(n, nsample, c).zero_()
        pointops_cuda.subtraction_forward_cuda(n, nsample, c, input1, input2, idx, output)
        ctx.save_for_backward(idx)
        return output

    @staticmethod
    def backward(ctx, grad_output):
        """
        input: grad_out: (n, nsample, c)
        output: grad_input1: (n, c), grad_input2: (n, c)
        """
        idx, = ctx.saved_tensors
        n, nsample, c = grad_output.shape
        grad_input1 = torch.cuda.FloatTensor(n, c).zero_()
        grad_input2 = torch.cuda.FloatTensor(n, c).zero_()
        pointops_cuda.subtraction_backward_cuda(n, nsample, c, idx, grad_output, grad_input1, grad_input2)
        return grad_input1, grad_input2, None

subtraction = Subtraction.apply


class Aggregation(Function):
    @staticmethod
    def forward(ctx, input, position, weight, idx):
        """
        input: input: (n, c), position: (n, nsample, c), weight : (n, nsample, c'), idx: (n, nsample)
        output: (n, c)
        """
        assert input.is_contiguous() and position.is_contiguous() and weight.is_contiguous()
        n, nsample, c = position.shape; w_c = weight.shape[-1]
        output = torch.cuda.FloatTensor(n, c).zero_()
        pointops_cuda.aggregation_forward_cuda(n, nsample, c, w_c, input, position, weight, idx, output)
        ctx.save_for_backward(input, position, weight, idx)
        return output

    @staticmethod
    def backward(ctx, grad_output):
        """
        input: grad_out: (n, c)
        output: grad_input: (n, c), grad_position: (n, nsample, c), grad_weight : (n, nsample, c')
        """
        input, position, weight, idx = ctx.saved_tensors
        n, nsample, c = position.shape; w_c = weight.shape[-1]
        grad_input = torch.cuda.FloatTensor(n, c).zero_()
        grad_position = torch.cuda.FloatTensor(n, nsample, c).zero_()
        grad_weight = torch.cuda.FloatTensor(n, nsample, w_c).zero_()
        pointops_cuda.aggregation_backward_cuda(n, nsample, c, w_c, input, position, weight, idx, grad_output, grad_input, grad_position, grad_weight)
        return grad_input, grad_position, grad_weight, None

aggregation = Aggregation.apply


def interpolation(xyz, new_xyz, feat, offset, new_offset, k=3):
    """
    input: xyz: (m, 3), new_xyz: (n, 3), feat: (m, c), offset: (b), new_offset: (b)
    output: (n, c)
    """
    assert xyz.is_contiguous() and new_xyz.is_contiguous() and feat.is_contiguous()
    idx, dist = knnquery(k, xyz, new_xyz, offset, new_offset) # (n, 3), (n, 3)
    dist_recip = 1.0 / (dist + 1e-8) # (n, 3)
    norm = torch.sum(dist_recip, dim=1, keepdim=True)
    weight = dist_recip / norm # (n, 3)

    new_feat = torch.cuda.FloatTensor(new_xyz.shape[0], feat.shape[1]).zero_()
    for i in range(k):
        new_feat += feat[idx[:, i].long(), :] * weight[:, i].unsqueeze(-1)
    return new_feat


class Interpolation(Function):
    @staticmethod
    def forward(ctx, xyz, new_xyz, input, offset, new_offset, k=3):
        """
        input: xyz: (m, 3), new_xyz: (n, 3), input: (m, c), offset: (b), new_offset: (b)
        output: (n, c)
        """
        assert xyz.is_contiguous() and new_xyz.is_contiguous() and input.is_contiguous()
        idx, dist = knnquery(k, xyz, new_xyz, offset, new_offset) # (n, k), (n, k)
        dist_recip = 1.0 / (dist + 1e-8) # (n, k)
        norm = torch.sum(dist_recip, dim=1, keepdim=True)
        weight = dist_recip / norm # (n, k)

        n, c, m = new_xyz.shape[0], input.shape[1], input.shape[0]
        output = torch.cuda.FloatTensor(n, c).zero_()
        pointops_cuda.interpolation_forward_cuda(n, c, k, input, idx, weight, output)
        ctx.m, ctx.k = m, k
        ctx.save_for_backward(idx, weight)
        return output

    @staticmethod
    def backward(ctx, grad_output):
        """
        input: xyz: (m, 3), new_xyz: (n, 3), input: (m, c), offset: (b), new_offset: (b)
        output: (n, c)
        """
        m, k = ctx.m, ctx.k
        idx, weight = ctx.saved_tensors
        n, c = grad_output.shape
        grad_input = torch.cuda.FloatTensor(m, c).zero_()
        pointops_cuda.interpolation_backward_cuda(n, c, k, grad_output, idx, weight, grad_input)
        return None, None, grad_input, None, None, None

interpolation2 = Interpolation.apply

utils/pointops2/functions/pointops_ablation.py
aminebdj/3D-OWIS
import torch
import pointops
from torch_scatter import scatter_max, scatter_mean, scatter_add, scatter_min, scatter_sum

torch.manual_seed(1)

M = 800000
N = 35000
C = 96
h = 6
query = torch.rand(N, h, C//h).cuda()
key = torch.rand(N, h, C//h).cuda()

index_0 = torch.rand(M)
index_0[index_0 < 0] = 0
index_0 = (index_0*N).long().cuda()

index_1 = torch.rand(M)
index_1[index_1 < 0] = 0
index_1 = (index_1*N).long().cuda()

query.requires_grad = True
key.requires_grad = True

# rearrange index for acceleration
index_0, indices = torch.sort(index_0) #[M,]
index_1 = index_1[indices] #[M,]
index_0_counts = index_0.bincount()

print("index_0_counts.shape: ", index_0_counts.shape)

n_max = index_0_counts.max()
index_0_offsets = index_0_counts.cumsum(dim=-1) #[N]

print("v1 index_0_offsets.shape: ", index_0_offsets.shape)

index_0_offsets = torch.cat([torch.zeros(1, dtype=torch.long).cuda(), index_0_offsets], 0) #[N+1]
        
# print("index_0[:100]: ", index_0[:100])
print("n_max: ", n_max)
print("index_0_offsets.shape: ", index_0_offsets.shape)
# input()

print("index_0_offsets[:100]: ", index_0_offsets[:100])
print("index_1[300:320]: ", index_1[300:320])

            
attn_flat = pointops.attention_step1(query.float(), key.float(), index_0.int(), index_1.int())
# loss = attn_flat.sum()
# loss.backward()
print("attn_flat.shape: {}, attn_flat[300:320,:10]: {}".format(attn_flat.shape, attn_flat[300:320,:10]))
# print("query.grad[:5, :3, :5]: ", query.grad[:5, :3, :5])
# print("key.grad[:5, :3, :5]: ", key.grad[:5, :3, :5])
# input()

print("query.is_contiguous(): ", query.is_contiguous())
print("key.is_contiguous(): ", key.is_contiguous())
print("index_0.is_contiguous(): ", index_0.is_contiguous())
print("index_1.is_contiguous(): ", index_1.is_contiguous())

attn_flat_v2 = pointops.attention_step1_v2(query.float(), key.float(), index_1.int(), index_0_offsets.int(), n_max)
# loss = attn_flat_v2.sum()
# loss.backward()
print("attn_flat_v2.shape: {}, attn_flat_v2[300:320,:10]: {}".format(attn_flat_v2.shape, attn_flat_v2[300:320,:10]))
# print("query.grad[:5, :3, :5]: ", query.grad[:5, :3, :5])
# print("key.grad[:5, :3, :5]: ", key.grad[:5, :3, :5])
# input()

mask = attn_flat_v2.sum(-1) != 0
print("mask.sum(): ", mask.sum())
print("attn_flat_v2[mask] - attn_flat[mask]: ", ((attn_flat_v2[mask] - attn_flat[mask])**2).max())


print("((attn_flat-attn_flat_v2)**2 < 1e-8).all(): ", ((attn_flat-attn_flat_v2)**2 < 1e-8).all())

selected = 10000
print("torch.max((attn_flat[:selected]-attn_flat_v2[:selected])**2, 0): ", torch.max((attn_flat[:selected]-attn_flat_v2[:selected])**2, 0))


utils/pointops2/functions/test_attention_op_step1.py
aminebdj/3D-OWIS
import torch
import pointops
from torch_scatter import scatter_max, scatter_mean, scatter_add, scatter_min, scatter_sum

torch.manual_seed(1)

M = 800000
N = 35000
C = 96
h = 6
query = torch.rand(N, h, C//h).cuda()
key = torch.rand(N, h, C//h).cuda()

index_0 = torch.rand(M)
index_0[index_0 < 0] = 0
index_0 = (index_0*N).long().cuda()

index_1 = torch.rand(M)
index_1[index_1 < 0] = 0
index_1 = (index_1*N).long().cuda()

query.requires_grad = True
key.requires_grad = True


attn_flat = pointops.attention_step1(query.float(), key.float(), index_0.int(), index_1.int())
loss = attn_flat.sum()
loss.backward()
print("attn_flat.shape: {}, attn_flat[:20,:10]: {}".format(attn_flat.shape, attn_flat[:20,:10]))
print("query.grad[:5, :3, :5]: ", query.grad[:5, :3, :5])
print("key.grad[:5, :3, :5]: ", key.grad[:5, :3, :5])
input()



# rearrange index for acceleration
index_0, indices = torch.sort(index_0) #[M,]
index_1 = index_1[indices] #[M,]
index_0_counts = index_0.bincount()

print("index_0_counts.shape: ", index_0_counts.shape)

n_max = index_0_counts.max()
index_0_offsets = index_0_counts.cumsum(dim=-1) #[N]

print("v1 index_0_offsets.shape: ", index_0_offsets.shape)

index_0_offsets = torch.cat([torch.zeros(1, dtype=torch.long).cuda(), index_0_offsets], 0) #[N+1]
        
# print("index_0[:100]: ", index_0[:100])
print("n_max: ", n_max)
print("index_0_offsets.shape: ", index_0_offsets.shape)
# input()

print("index_0_offsets[:100]: ", index_0_offsets[:100])
print("index_1[:20]: ", index_1[:20])

            
attn_flat = pointops.attention_step1(query.float(), key.float(), index_0.int(), index_1.int())
# loss = attn_flat.sum()
# loss.backward()
# # attn_flat = pointops.attention_step1(query.float(), key.float(), index_0.int(), index_1.int())
# # loss = attn_flat.sum()
# # loss.backward()
# print("attn_flat.shape: {}, attn_flat[:20,:10]: {}".format(attn_flat.shape, attn_flat[:20,:10]))
# print("query.grad[:5, :3, :5]: ", query.grad[:5, :3, :5])
# print("key.grad[:5, :3, :5]: ", key.grad[:5, :3, :5])
# input()

print("query.is_contiguous(): ", query.is_contiguous())
print("key.is_contiguous(): ", key.is_contiguous())
print("index_0.is_contiguous(): ", index_0.is_contiguous())
print("index_1.is_contiguous(): ", index_1.is_contiguous())

attn_flat_v2 = pointops.attention_step1_v2(query.float(), key.float(), index_1.int(), index_0_offsets.int(), n_max)
loss = attn_flat_v2.sum()
loss.backward()

# attn_flat_v2 = pointops.attention_step1_v2(query.float(), key.float(), index_1.int(), index_0_offsets.int(), n_max)
# loss = attn_flat_v2.sum()
# loss.backward()

print("attn_flat_v2.shape: {}, attn_flat_v2[:20,:10]: {}".format(attn_flat_v2.shape, attn_flat_v2[:20,:10]))
print("query.grad[:5, :3, :5]: ", query.grad[:5, :3, :5])
print("key.grad[:5, :3, :5]: ", key.grad[:5, :3, :5])
# input()

# mask = attn_flat_v2.sum(-1) != 0
# print("mask.sum(): ", mask.sum())
# print("attn_flat_v2[mask] - attn_flat[mask]: ", ((attn_flat_v2[mask] - attn_flat[mask])**2).max())


print("((attn_flat-attn_flat_v2)**2 < 1e-8).all(): ", ((attn_flat-attn_flat_v2)**2 < 1e-8).all())

selected = 10000
print("torch.max((attn_flat[:selected]-attn_flat_v2[:selected])**2, 0): ", torch.max((attn_flat[:selected]-attn_flat_v2[:selected])**2, 0))


utils/pointops2/functions/test_attention_op_step1_v2.py
aminebdj/3D-OWIS
import torch
import pointops
from torch_scatter import scatter_max, scatter_mean, scatter_add, scatter_min, scatter_sum

torch.manual_seed(1)

M = 800000
N = 35000
C = 96
h = 6
softmax_attn_flat = torch.rand(M, h).cuda()
value = torch.rand(N, h, C//h).cuda()

index_0 = torch.rand(M)
index_0[index_0 < 0] = 0
index_0 = (index_0*N).long().cuda()

index_1 = torch.rand(M)
index_1[index_1 < 0] = 0
index_1 = (index_1*N).long().cuda()

softmax_attn_flat.requires_grad = True
value.requires_grad = True

# value_flat = value[index_1] #[M, num_heads, C // num_heads]
# x = (softmax_attn_flat.unsqueeze(-1) * value_flat).reshape(M, C)
# x = scatter_sum(src=x, index=index_0, dim=0, dim_size=N) #[N, C]
# loss = x.sum()
# loss.backward()

# print("x.shape: {}, x[:5,:10]: {}".format(x.shape, x[:5,:10]))
# print("softmax_attn_flat.grad[:5, :10]: ", softmax_attn_flat.grad[:5, :10])
# print("value.grad[:5, :3, :5]: ", value.grad[:5, :3, :5])
# input()

print("softmax_attn_flat.is_contiguous(): ", softmax_attn_flat.is_contiguous())
print("value.is_contiguous(): ", value.is_contiguous())
print("index_0.is_contiguous(): ", index_0.is_contiguous())
print("index_1.is_contiguous(): ", index_1.is_contiguous())

x_v2 = pointops.attention_step2(softmax_attn_flat.float(), value.float(), index_0.int(), index_1.int())
x_v2 = x_v2.view(N, C)
loss = x_v2.sum()
loss.backward()

print("x_v2.shape: {}, x_v2[:5,:10]: {}".format(x_v2.shape, x_v2[:5,:10]))

print("softmax_attn_flat.grad[:5, :10]: ", softmax_attn_flat.grad[:5, :10])
print("value.grad[:5, :3, :5]: ", value.grad[:5, :3, :5])
input()

print("((x-x_v2)**2 < 1e-8).all(): ", ((x-x_v2)**2 < 1e-8).all())

print("torch.max((x-x_v2)**2): ", torch.max((x-x_v2)**2))


utils/pointops2/functions/test_attention_op_step2.py
aminebdj/3D-OWIS
import torch
import pointops
from torch_scatter import scatter_max, scatter_mean, scatter_add, scatter_min, scatter_sum

torch.manual_seed(1)

M = 80000
N = 3500
hdim = 16
h = 6
L = 31
query = torch.rand(N, h, hdim).cuda()
table = torch.rand(L, h, hdim, 3).cuda()

index = torch.rand(M)
index[index < 0] = 0
index = (index*N).long().cuda()

rel_index = torch.rand(M, 3)
rel_index[rel_index < 0] = 0
rel_index = (rel_index*L).long().cuda()

query.requires_grad = True
table.requires_grad = True

# query_flat = query[index] #[M, h, hdim]
# table_x, table_y, table_z = table[:,:,:,0], table[:,:,:,1], table[:,:,:,2] #[L, h, hdim]
# rel_index_x, rel_index_y, rel_index_z = rel_index[:,0], rel_index[:,1], rel_index[:,2] #[M]
# rel_pos_encoding = table_x[rel_index_x] + table_y[rel_index_y] + table_z[rel_index_z] #[M, h, hdim]
# output = (query_flat * rel_pos_encoding).sum(-1) #[M, h]
# loss = output.mean()
# loss.backward()

# print("output.shape: {}, output[:5,:10]: {}".format(output.shape, output[:5,:10]))
# print("query.grad[:5, :3, :5]: ", query.grad[:5, :3, :5])
# print("table.grad[:5, :3, :5, :2]: ", table.grad[:5, :3, :5, :2])
# input()

# print("query.is_contiguous(): ", query.is_contiguous())
# print("key.is_contiguous(): ", key.is_contiguous())
# print("index_0.is_contiguous(): ", index_0.is_contiguous())
# print("index_1.is_contiguous(): ", index_1.is_contiguous())

output_v2 = pointops.dot_prod_with_idx(query, index.int(), table, rel_index.int())
loss = output_v2.mean()
loss.backward()

print("output_v2.shape: {}, output_v2[:5,:10]: {}".format(output_v2.shape, output_v2[:5,:10]))
print("v2: query.grad[:5, :3, :5]: ", query.grad[:5, :3, :5])
print("v2: table.grad[:5, :3, :5, :2]: ", table.grad[:5, :3, :5, :2])
input()

# print("((output-output_v2)**2).max(): ", ((output-output_v2)**2).max())

# print("torch.max((attn_flat-attn_flat_v2)**2): ", torch.max((attn_flat-attn_flat_v2)**2))


utils/pointops2/functions/test_relative_pos_encoding_op_step1.py
aminebdj/3D-OWIS
import torch
import pointops
from torch_scatter import scatter_max, scatter_mean, scatter_add, scatter_min, scatter_sum

torch.manual_seed(1)

M = 80000
N = 3500
hdim = 16
h = 6
L = 31
query = torch.rand(N, h, hdim).cuda()
table_q = torch.rand(L, h, hdim, 3).cuda()
key = torch.rand(N, h, hdim).cuda()
table_k = torch.rand(L, h, hdim, 3).cuda()

index_q = torch.rand(M)
index_q[index_q < 0] = 0
index_q = (index_q*N).long().cuda()

index_k = torch.rand(M)
index_k[index_k < 0] = 0
index_k = (index_k*N).long().cuda()

rel_index = torch.rand(M, 3)
rel_index[rel_index < 0] = 0
rel_index = (rel_index*L).long().cuda()

query.requires_grad = True
table_q.requires_grad = True
key.requires_grad = True
table_k.requires_grad = True

output1 = pointops.dot_prod_with_idx(query, index_q.int(), table_q, rel_index.int())
output2 = pointops.dot_prod_with_idx(key, index_k.int(), table_k, rel_index.int())
output = output1 + output2
# loss = output.mean()
# loss.backward()

# print("output.shape: {}, output[:5,:10]: {}".format(output.shape, output[:5,:10]))
# print("query.grad[:5, :3, :5]: ", query.grad[:5, :3, :5])
# print("table_q.grad[:5, :3, :5, :2]: ", table_q.grad[:5, :3, :5, :2])
# print("key.grad[:5, :3, :5]: ", key.grad[:5, :3, :5])
# print("table_k.grad[:5, :3, :5, :2]: ", table_k.grad[:5, :3, :5, :2])
# input()

# print("query.is_contiguous(): ", query.is_contiguous())
# print("key.is_contiguous(): ", key.is_contiguous())
# print("index_0.is_contiguous(): ", index_0.is_contiguous())
# print("index_1.is_contiguous(): ", index_1.is_contiguous())

output_v2 = pointops.dot_prod_with_idx_v2(query, index_q.int(), key, index_k.int(), table_q, table_k, rel_index.int())
loss = output_v2.mean()
loss.backward()

print("output_v2.shape: {}, output_v2[:5,:10]: {}".format(output_v2.shape, output_v2[:5,:10]))
print("v2 query.grad[:5, :3, :5]: ", query.grad[:5, :3, :5])
print("v2 table_q.grad[:5, :3, :5, :2]: ", table_q.grad[:5, :3, :5, :2])
print("v2 key.grad[:5, :3, :5]: ", key.grad[:5, :3, :5])
print("v2 table_k.grad[:5, :3, :5, :2]: ", table_k.grad[:5, :3, :5, :2])
# input()

print("((output-output_v2)**2).max(): ", ((output-output_v2)**2).max())


utils/pointops2/functions/test_relative_pos_encoding_op_step1_v2.py
aminebdj/3D-OWIS
import torch
import pointops
from torch_scatter import scatter_max, scatter_mean, scatter_add, scatter_min, scatter_sum

torch.manual_seed(1)

M = 80000
N = 3500
# M = 80
# N = 5
hdim = 16
h = 6
L = 31
query = torch.rand(N, h, hdim).cuda()
table_q = torch.rand(L, h, hdim, 3).cuda()
key = torch.rand(N, h, hdim).cuda()
table_k = torch.rand(L, h, hdim, 3).cuda()

index_q = torch.rand(M)
index_q[index_q < 0] = 0
index_q = (index_q*N).long().cuda()

index_k = torch.rand(M)
index_k[index_k < 0] = 0
index_k = (index_k*N).long().cuda()

rel_index = torch.rand(M, 3)
rel_index[rel_index < 0] = 0
rel_index = (rel_index*L).long().cuda()


# rearrange index for acceleration
index_q, indices = torch.sort(index_q) #[M,]
index_k = index_k[indices] #[M,]
rel_index = rel_index[indices]
index_q_counts = index_q.bincount()

print("index_q_counts.shape: ", index_q_counts.shape)

n_max = index_q_counts.max()
index_q_offsets = index_q_counts.cumsum(dim=-1) #[N]

print("v1 index_q_offsets.shape: ", index_q_offsets.shape)

index_q_offsets = torch.cat([torch.zeros(1, dtype=torch.long).cuda(), index_q_offsets], 0) #[N+1]
        
# print("index_q[:100]: ", index_q[:100])
print("n_max: ", n_max)
print("index_q_offsets.shape: ", index_q_offsets.shape)
# input()

print("index_q_offsets[:100]: ", index_q_offsets[:100])
print("index_k[:20]: ", index_k[:20])

query.requires_grad = True
table_q.requires_grad = True
key.requires_grad = True
table_k.requires_grad = True

output1 = pointops.dot_prod_with_idx(query, index_q.int(), table_q, rel_index.int())
output2 = pointops.dot_prod_with_idx(key, index_k.int(), table_k, rel_index.int())
output = output1 + output2
loss = output.mean()
loss.backward()

# print("output.shape: {}, output[:5,:10]: {}".format(output.shape, output[:5,:10]))
# print("query.grad[:5, :3, :5]: ", query.grad[:5, :3, :5])
# print("table_q.grad[:5, :3, :5, :2]: ", table_q.grad[:5, :3, :5, :2])
# print("key.grad[:5, :3, :5]: ", key.grad[:5, :3, :5])
# print("table_k.grad[:5, :3, :5, :2]: ", table_k.grad[:5, :3, :5, :2])
# input()

# print("query.is_contiguous(): ", query.is_contiguous())
# print("key.is_contiguous(): ", key.is_contiguous())
# print("index_q.is_contiguous(): ", index_q.is_contiguous())
# print("index_k.is_contiguous(): ", index_k.is_contiguous())

output_v2 = pointops.dot_prod_with_idx_v3(query, index_q_offsets.int(), n_max, key, index_k.int(), table_q, table_k, rel_index.int())
# loss = output_v2.mean()
# loss.backward()

# print("output_v2.shape: {}, output_v2[:5,:10]: {}".format(output_v2.shape, output_v2[:5,:10]))
# print("v2 query.grad[:5, :3, :5]: ", query.grad[:5, :3, :5])
# print("v2 table_q.grad[:5, :3, :5, :2]: ", table_q.grad[:5, :3, :5, :2])
# print("v2 key.grad[:5, :3, :5]: ", key.grad[:5, :3, :5])
# print("v2 table_k.grad[:5, :3, :5, :2]: ", table_k.grad[:5, :3, :5, :2])
# input()

print("((output-output_v2)**2).max(): ", ((output-output_v2)**2).max())


utils/pointops2/functions/test_relative_pos_encoding_op_step1_v3.py
aminebdj/3D-OWIS
import torch
import pointops
from torch_scatter import scatter_max, scatter_mean, scatter_add, scatter_min, scatter_sum

torch.manual_seed(1)

M = 80000
N = 3500
hdim = 16
h = 6
L = 31
attn = torch.rand(M, h).cuda()
v = torch.rand(N, h, hdim).cuda()
table = torch.rand(L, h, hdim, 3).cuda()

index_0 = torch.rand(M)
index_0[index_0 < 0] = 0
index_0 = (index_0*N).long().cuda()

index_1 = torch.rand(M)
index_1[index_1 < 0] = 0
index_1 = (index_1*N).long().cuda()

rel_index = torch.rand(M, 3)
rel_index[rel_index < 0] = 0
rel_index = (rel_index*L).long().cuda()

attn.requires_grad = True
v.requires_grad = True
table.requires_grad = True

v_flat = v[index_1] #[M, h, hdim]
table_x, table_y, table_z = table[:,:,:,0], table[:,:,:,1], table[:,:,:,2] #[L, h, hdim]
rel_index_x, rel_index_y, rel_index_z = rel_index[:,0], rel_index[:,1], rel_index[:,2] #[M]
rel_pos_encoding = table_x[rel_index_x] + table_y[rel_index_y] + table_z[rel_index_z] #[M, h, hdim]
v_flat_new = v_flat + rel_pos_encoding #[M, h, hdim]
output = attn.unsqueeze(-1) * v_flat_new #[M, h, hdim] 
output = scatter_sum(src=output, index=index_0, dim=0, dim_size=N) #[N, h, hdim]
loss = output.mean()
loss.backward()

print("output.shape: {}, output[:5,:10,:5]: {}".format(output.shape, output[:5,:10, :5]))
print("attn.grad[:5, :3]: ", attn.grad[:5, :3])
print("v.grad[:5, :3, :5]: ", v.grad[:5, :3, :5])
print("table.grad[:5, :3, :5, :2]: ", table.grad[:5, :3, :5, :2])
input()

# print("query.is_contiguous(): ", query.is_contiguous())
# print("key.is_contiguous(): ", key.is_contiguous())
# print("index_0.is_contiguous(): ", index_0.is_contiguous())
# print("index_1.is_contiguous(): ", index_1.is_contiguous())

# output_v2 = pointops.attention_step2_with_rel_pos_value(attn, v, index_0.int(), index_1.int(), table, rel_index.int())
# loss = output_v2.mean()
# loss.backward()

# print("output_v2.shape: {}, output_v2[:5,:10,:5]: {}".format(output_v2.shape, output_v2[:5,:10,:5]))
# print("v2 attn.grad[:5, :3]: ", attn.grad[:5, :3])
# print("v2 v.grad[:5, :3, :5]: ", v.grad[:5, :3, :5])
# print("v2 table.grad[:5, :3, :5, :2]: ", table.grad[:5, :3, :5, :2])
# input()

# print("((output-output_v2)**2).max(): ", ((output-output_v2)**2).max())

# print("torch.max((attn_flat-attn_flat_v2)**2): ", torch.max((attn_flat-attn_flat_v2)**2))


utils/pointops2/functions/test_relative_pos_encoding_op_step2.py
aminebdj/3D-OWIS
import torch
import pointops
from torch_scatter import scatter_max, scatter_mean, scatter_add, scatter_min, scatter_sum

torch.manual_seed(1)

M = 80000
N = 3500
hdim = 16
h = 6
L = 31
attn = torch.rand(M, h).cuda()
v = torch.rand(N, h, hdim).cuda()
table = torch.rand(L, h, hdim, 3).cuda()

index_0 = torch.rand(M)
index_0[index_0 < 0] = 0
index_0 = (index_0*N).long().cuda()

index_1 = torch.rand(M)
index_1[index_1 < 0] = 0
index_1 = (index_1*N).long().cuda()

rel_index = torch.rand(M, 3)
rel_index[rel_index < 0] = 0
rel_index = (rel_index*L).long().cuda()


# rearrange index for acceleration
index_0, indices = torch.sort(index_0) #[M,]
index_1 = index_1[indices] #[M,]
rel_index = rel_index[indices]
index_0_counts = index_0.bincount()

print("index_0_counts.shape: ", index_0_counts.shape)

n_max = index_0_counts.max()
index_0_offsets = index_0_counts.cumsum(dim=-1) #[N]

print("v1 index_0_offsets.shape: ", index_0_offsets.shape)

index_0_offsets = torch.cat([torch.zeros(1, dtype=torch.long).cuda(), index_0_offsets], 0) #[N+1]


attn.requires_grad = True
v.requires_grad = True
table.requires_grad = True


output = pointops.attention_step2_with_rel_pos_value(attn, v, index_0.int(), index_1.int(), table, rel_index.int())
loss = output.mean()
loss.backward()

print("output.shape: {}, output[:5,:10,:5]: {}".format(output.shape, output[:5,:10, :5]))
print("attn.grad[:5, :3]: ", attn.grad[:5, :3])
print("v.grad[:5, :3, :5]: ", v.grad[:5, :3, :5])
print("table.grad[:5, :3, :5, :2]: ", table.grad[:5, :3, :5, :2])
# input()

attn_grad = attn.grad.clone()
v_grad = v.grad.clone()
table_grad = table.grad.clone()

attn.grad.zero_()
v.grad.zero_()
table.grad.zero_()

# print("query.is_contiguous(): ", query.is_contiguous())
# print("key.is_contiguous(): ", key.is_contiguous())
# print("index_0.is_contiguous(): ", index_0.is_contiguous())
# print("index_1.is_contiguous(): ", index_1.is_contiguous())

output_v2 = pointops.attention_step2_with_rel_pos_value_v2(attn, v, index_0_offsets.int(), n_max, index_1.int(), table, rel_index.int())
loss = output_v2.mean()
loss.backward()

print("output_v2.shape: {}, output_v2[:5,:10,:5]: {}".format(output_v2.shape, output_v2[:5,:10,:5]))
print("v2 attn.grad[:5, :3]: ", attn.grad[:5, :3])
print("v2 v.grad[:5, :3, :5]: ", v.grad[:5, :3, :5])
print("v2 table.grad[:5, :3, :5, :2]: ", table.grad[:5, :3, :5, :2])
# input()

print("((output-output_v2)**2).max(): ", ((output-output_v2)**2).max())

print("((attn_grad-attn.grad)**2).max(): ", ((attn_grad-attn.grad)**2).max())

print("((v_grad-v.grad)**2).max(): ", ((v_grad-v.grad)**2).max())

print("((table_grad-table.grad)**2).max(): ", ((table_grad-table.grad)**2).max())

# print("torch.max((attn_flat-attn_flat_v2)**2): ", torch.max((attn_flat-attn_flat_v2)**2))


utils/pointops2/functions/test_relative_pos_encoding_op_step2_v2.py
aminebdj/3D-OWIS
#python3 setup.py install
from setuptools import setup
from torch.utils.cpp_extension import BuildExtension, CUDAExtension
import os
from distutils.sysconfig import get_config_vars

(opt,) = get_config_vars('OPT')
os.environ['OPT'] = " ".join(
    flag for flag in opt.split() if flag != '-Wstrict-prototypes'
)

setup(
    name='pointops2',
    ext_modules=[
        CUDAExtension('pointops2_cuda', [
            'src/pointops_api.cpp',
            'src/knnquery/knnquery_cuda.cpp',
            'src/knnquery/knnquery_cuda_kernel.cu',
            'src/sampling/sampling_cuda.cpp',
            'src/sampling/sampling_cuda_kernel.cu',
            'src/grouping/grouping_cuda.cpp',
            'src/grouping/grouping_cuda_kernel.cu',
            'src/interpolation/interpolation_cuda.cpp',
            'src/interpolation/interpolation_cuda_kernel.cu',
            'src/subtraction/subtraction_cuda.cpp',
            'src/subtraction/subtraction_cuda_kernel.cu',
            'src/aggregation/aggregation_cuda.cpp',
            'src/aggregation/aggregation_cuda_kernel.cu',
            'src/attention/attention_cuda.cpp',
            'src/attention/attention_cuda_kernel.cu',
            'src/rpe/relative_pos_encoding_cuda.cpp',
            'src/rpe/relative_pos_encoding_cuda_kernel.cu',
            'src/attention_v2/attention_cuda_v2.cpp',
            'src/attention_v2/attention_cuda_kernel_v2.cu',
            'src/rpe_v2/relative_pos_encoding_cuda_v2.cpp',
            'src/rpe_v2/relative_pos_encoding_cuda_kernel_v2.cu',
            ],
        extra_compile_args={'cxx': ['-g'], 'nvcc': ['-O2']}
        )
    ],
    cmdclass={'build_ext': BuildExtension}
)

utils/pointops2/setup.py
aminebdj/3D-OWIS
import sys

if sys.version_info[:2] >= (3, 8):
    from collections.abc import MutableMapping
else:
    from collections import MutableMapping

import torch
from loguru import logger


def flatten_dict(d, parent_key="", sep="_"):
    """
    https://stackoverflow.com/questions/6027558/flatten-nested-dictionaries-compressing-keys
    """
    items = []
    for k, v in d.items():
        new_key = parent_key + sep + k if parent_key else k
        if isinstance(v, MutableMapping):
            items.extend(flatten_dict(v, new_key, sep=sep).items())
        else:
            items.append((new_key, v))
    return dict(items)


def load_baseline_model(cfg, model):
    # if it is Minkoski weights
    cfg.model.in_channels = 3
    cfg.model.config.conv1_kernel_size = 5
    cfg.data.add_normals = False
    cfg.data.train_dataset.color_mean_std = [(0.5, 0.5, 0.5), (1, 1, 1)]
    cfg.data.validation_dataset.color_mean_std = [(0.5, 0.5, 0.5), (1, 1, 1)]
    cfg.data.test_dataset.color_mean_std = [(0.5, 0.5, 0.5), (1, 1, 1)]
    cfg.data.voxel_size = 0.02
    model = model(cfg)
    state_dict = torch.load(cfg.general.checkpoint)["state_dict"]
    model.model.load_state_dict(state_dict)
    return cfg, model

def load_backbone_checkpoint_with_missing_or_exsessive_keys(cfg, model):
    state_dict = torch.load(cfg.general.backbone_checkpoint)["state_dict"]
    correct_dict = dict(model.state_dict())

    # if parametrs not found in checkpoint they will be randomly initialized
    for key in state_dict.keys():
        if correct_dict.pop(f"model.backbone.{key}", None) is None:
            logger.warning(f"Key not found, it will be initialized randomly: {key}")

    # if parametrs have different shape, it will randomly initialize
    state_dict = torch.load(cfg.general.backbone_checkpoint)["state_dict"]
    correct_dict = dict(model.state_dict())
    for key in correct_dict.keys():
        if key.replace("model.backbone.", "") not in state_dict:
            logger.warning(
                f"{key} not in loaded checkpoint"
            )
            state_dict.update({key.replace("model.backbone.", ""): correct_dict[key]})
        elif state_dict[key.replace("model.backbone.", "")].shape != correct_dict[key].shape:
            logger.warning(
                f"incorrect shape {key}:{state_dict[key.replace('model.backbone.', '')].shape} vs {correct_dict[key].shape}"
            )
            state_dict.update({key: correct_dict[key]})

    # if we have more keys just discard them
    correct_dict = dict(model.state_dict())
    new_state_dict = dict()
    for key in state_dict.keys():
        if f"model.backbone.{key}" in correct_dict.keys():
            new_state_dict.update({f"model.backbone.{key}": state_dict[key]})
        elif key in correct_dict.keys():
            new_state_dict.update({key: correct_dict[key]})
        else:
            logger.warning(f"excessive key: {key}")
    model.load_state_dict(new_state_dict)
    return cfg, model

def load_checkpoint_with_missing_or_exsessive_keys(cfg, model):
    state_dict = torch.load(cfg.general.checkpoint)["state_dict"]
    correct_dict = dict(model.state_dict())

    # if parametrs not found in checkpoint they will be randomly initialized
    for key in state_dict.keys():
        if correct_dict.pop(key, None) is None:
            logger.warning(f"Key not found, it will be initialized randomly: {key}")

    # if parametrs have different shape, it will randomly initialize
    state_dict = torch.load(cfg.general.checkpoint)["state_dict"]
    correct_dict = dict(model.state_dict())
    for key in correct_dict.keys():
        if key not in state_dict:
            logger.warning(
                f"{key} not in loaded checkpoint"
            )
            state_dict.update({key: correct_dict[key]})
        elif state_dict[key].shape != correct_dict[key].shape:
            logger.warning(
                f"incorrect shape {key}:{state_dict[key].shape} vs {correct_dict[key].shape}"
            )
            state_dict.update({key: correct_dict[key]})

    # if we have more keys just discard them
    correct_dict = dict(model.state_dict())
    new_state_dict = dict()
    for key in state_dict.keys():
        if key in correct_dict.keys():
            new_state_dict.update({key: state_dict[key]})
        else:
            logger.warning(f"excessive key: {key}")
    model.load_state_dict(new_state_dict)
    return cfg, model


def freeze_until(net, param_name: str = None):
    """
    Freeze net until param_name
    https://opendatascience.slack.com/archives/CGK4KQBHD/p1588373239292300?thread_ts=1588105223.275700&cid=CGK4KQBHD
    Args:
        net:
        param_name:
    Returns:
    """
    found_name = False
    for name, params in net.named_parameters():
        if name == param_name:
            found_name = True
        params.requires_grad = found_name

utils/utils.py
aminebdj/3D-OWIS
# Copyright (c) Facebook, Inc. and its affiliates.
# 
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

""" Generic Code for Object Detection Evaluation

    Input:
    For each class:
        For each image:
            Predictions: box, score
            Groundtruths: box
    
    Output:
    For each class:
        precision-recal and average precision
    
    Author: Charles R. Qi
    
    Ref: https://raw.githubusercontent.com/rbgirshick/py-faster-rcnn/master/lib/datasets/voc_eval.py
"""
import numpy as np

def voc_ap(rec, prec, use_07_metric=False):
    """ ap = voc_ap(rec, prec, [use_07_metric])
    Compute VOC AP given precision and recall.
    If use_07_metric is true, uses the
    VOC 07 11 point method (default:False).
    """
    if use_07_metric:
        # 11 point metric
        ap = 0.
        for t in np.arange(0., 1.1, 0.1):
            if np.sum(rec >= t) == 0:
                p = 0
            else:
                p = np.max(prec[rec >= t])
            ap = ap + p / 11.
    else:
        # correct AP calculation
        # first append sentinel values at the end
        mrec = np.concatenate(([0.], rec, [1.]))
        mpre = np.concatenate(([0.], prec, [0.]))

        # compute the precision envelope
        for i in range(mpre.size - 1, 0, -1):
            mpre[i - 1] = np.maximum(mpre[i - 1], mpre[i])

        # to calculate area under PR curve, look for points
        # where X axis (recall) changes value
        i = np.where(mrec[1:] != mrec[:-1])[0]

        # and sum (\Delta recall) * prec
        ap = np.sum((mrec[i + 1] - mrec[i]) * mpre[i + 1])
    return ap

import os
import sys
#BASE_DIR = os.path.dirname(os.path.abspath(__file__))
from utils.votenet_utils.metric_util import calc_iou # axis-aligned 3D box IoU
def get_iou(bb1, bb2):
    """ Compute IoU of two bounding boxes.
        ** Define your bod IoU function HERE **
    """
    #pass
    iou3d = calc_iou(bb1, bb2)
    return iou3d

from box_util import box3d_iou
def get_iou_obb(bb1,bb2):
    iou3d, iou2d = box3d_iou(bb1,bb2)
    return iou3d

def get_iou_main(get_iou_func, args):
    return get_iou_func(*args)

def eval_det_cls(pred, gt, ovthresh=0.25, use_07_metric=False, get_iou_func=get_iou):
    """ Generic functions to compute precision/recall for object detection
        for a single class.
        Input:
            pred: map of {img_id: [(bbox, score)]} where bbox is numpy array
            gt: map of {img_id: [bbox]}
            ovthresh: scalar, iou threshold
            use_07_metric: bool, if True use VOC07 11 point method
        Output:
            rec: numpy array of length nd
            prec: numpy array of length nd
            ap: scalar, average precision
    """

    # construct gt objects
    class_recs = {} # {img_id: {'bbox': bbox list, 'det': matched list}}
    npos = 0
    for img_id in gt.keys():
        bbox = np.array(gt[img_id])
        det = [False] * len(bbox)
        npos += len(bbox)
        class_recs[img_id] = {'bbox': bbox, 'det': det}
    # pad empty list to all other imgids
    for img_id in pred.keys():
        if img_id not in gt:
            class_recs[img_id] = {'bbox': np.array([]), 'det': []}

    # construct dets
    image_ids = []
    confidence = []
    BB = []
    for img_id in pred.keys():
        for box,score in pred[img_id]:
            image_ids.append(img_id)
            confidence.append(score)
            BB.append(box)
    confidence = np.array(confidence)
    BB = np.array(BB) # (nd,4 or 8,3 or 6)

    # sort by confidence
    sorted_ind = np.argsort(-confidence)
    sorted_scores = np.sort(-confidence)
    BB = BB[sorted_ind, ...]
    image_ids = [image_ids[x] for x in sorted_ind]

    # go down dets and mark TPs and FPs
    nd = len(image_ids)
    tp = np.zeros(nd)
    fp = np.zeros(nd)
    for d in range(nd):
        #if d%100==0: print(d)
        R = class_recs[image_ids[d]]
        bb = BB[d,...].astype(float)
        ovmax = -np.inf
        BBGT = R['bbox'].astype(float)

        if BBGT.size > 0:
            # compute overlaps
            for j in range(BBGT.shape[0]):
                iou = get_iou_main(get_iou_func, (bb, BBGT[j,...]))
                if iou > ovmax:
                    ovmax = iou
                    jmax = j

        #print d, ovmax
        if ovmax > ovthresh:
            if not R['det'][jmax]:
                tp[d] = 1.
                R['det'][jmax] = 1
            else:
                fp[d] = 1.
        else:
            fp[d] = 1.

    # compute precision recall
    fp = np.cumsum(fp)
    tp = np.cumsum(tp)
    rec = tp / float(npos)
    #print('NPOS: ', npos)
    # avoid divide by zero in case the first detection matches a difficult
    # ground truth
    prec = tp / np.maximum(tp + fp, np.finfo(np.float64).eps)
    ap = voc_ap(rec, prec, use_07_metric)

    return rec, prec, ap

def eval_det_cls_wrapper(arguments):
    pred, gt, ovthresh, use_07_metric, get_iou_func = arguments
    rec, prec, ap = eval_det_cls(pred, gt, ovthresh, use_07_metric, get_iou_func)
    return (rec, prec, ap)

def eval_det(pred_all, gt_all, ovthresh=0.25, use_07_metric=False, get_iou_func=get_iou):
    """ Generic functions to compute precision/recall for object detection
        for multiple classes.
        Input:
            pred_all: map of {img_id: [(classname, bbox, score)]}
            gt_all: map of {img_id: [(classname, bbox)]}
            ovthresh: scalar, iou threshold
            use_07_metric: bool, if true use VOC07 11 point method
        Output:
            rec: {classname: rec}
            prec: {classname: prec_all}
            ap: {classname: scalar}
    """
    pred = {} # map {classname: pred}
    gt = {} # map {classname: gt}
    for img_id in pred_all.keys():
        for classname, bbox, score in pred_all[img_id]:
            if classname not in pred: pred[classname] = {}
            if img_id not in pred[classname]:
                pred[classname][img_id] = []
            if classname not in gt: gt[classname] = {}
            if img_id not in gt[classname]:
                gt[classname][img_id] = []
            pred[classname][img_id].append((bbox,score))
    for img_id in gt_all.keys():
        for classname, bbox in gt_all[img_id]:
            if classname not in gt: gt[classname] = {}
            if img_id not in gt[classname]:
                gt[classname][img_id] = []
            # JONAS ADAPTATION TODO
            if classname not in pred: pred[classname] = {}
            if img_id not in pred[classname]:
                pred[classname][img_id] = []
            # =====================
            gt[classname][img_id].append(bbox)

    rec = {}
    prec = {}
    ap = {}
    for classname in gt.keys():
        print('Computing AP for class: ', classname)
        rec[classname], prec[classname], ap[classname] = eval_det_cls(pred[classname], gt[classname], ovthresh, use_07_metric, get_iou_func)
        print(classname, ap[classname])
    
    return rec, prec, ap 

from multiprocessing import Pool
def eval_det_multiprocessing(pred_all, gt_all, ovthresh=0.25, use_07_metric=False, get_iou_func=get_iou):
    """ Generic functions to compute precision/recall for object detection
        for multiple classes.
        Input:
            pred_all: map of {img_id: [(classname, bbox, score)]}
            gt_all: map of {img_id: [(classname, bbox)]}
            ovthresh: scalar, iou threshold
            use_07_metric: bool, if true use VOC07 11 point method
        Output:
            rec: {classname: rec}
            prec: {classname: prec_all}
            ap: {classname: scalar}
    """
    pred = {} # map {classname: pred}
    gt = {} # map {classname: gt}
    for img_id in pred_all.keys():
        for classname, bbox, score in pred_all[img_id]:
            if classname not in pred: pred[classname] = {}
            if img_id not in pred[classname]:
                pred[classname][img_id] = []
            if classname not in gt: gt[classname] = {}
            if img_id not in gt[classname]:
                gt[classname][img_id] = []
            pred[classname][img_id].append((bbox,score))
    for img_id in gt_all.keys():
        for classname, bbox in gt_all[img_id]:
            if classname not in gt: gt[classname] = {}
            if img_id not in gt[classname]:
                gt[classname][img_id] = []
            gt[classname][img_id].append(bbox)

    rec = {}
    prec = {}
    ap = {}
    p = Pool(processes=10)
    ret_values = p.map(eval_det_cls_wrapper, [(pred[classname], gt[classname], ovthresh, use_07_metric, get_iou_func) for classname in gt.keys() if classname in pred])
    p.close()
    for i, classname in enumerate(gt.keys()):
        if classname in pred:
            rec[classname], prec[classname], ap[classname] = ret_values[i]
        else:
            rec[classname] = 0
            prec[classname] = 0
            ap[classname] = 0
        print(classname, ap[classname])
    
    return rec, prec, ap 

utils/votenet_utils/eval_det.py
aminebdj/3D-OWIS
import os
import yaml
import numpy as np
import open3d as o3d
import random
from datasets.scannet200.owis_splits import KNOWN_CLASSES_LABELS,UNKNOWN_CLASSES_LABELS
from datasets.scannet200.scannet200_constants import CLASS_LABELS_200,VALID_CLASS_IDS_200
from copy import copy
# OUTPUT_2 = "output-scans/predictions-ours-2classes/"
label2id = {l:id for id,l in zip(VALID_CLASS_IDS_200, CLASS_LABELS_200)}
def generate_color_palette(x):
    palette = []

    while len(palette) < x:
        r, g, b = random.sample(range(256), 3)
        color = (r/255.0, g/255.0, b/255.0)
        if color[2] > 0:
            continue
        palette.append(color)

    return palette

def read_ply_file(file_path):
    # Read the ply file
    point_cloud = o3d.io.read_triangle_mesh(file_path)
    return point_cloud

model_name = "experiment_name"

for split in ['A','B','C']:
    for task in ['task1','task2','task3']:
        PATH = f"eval_output/{split}/{task}/instance_evaluation_{model_name}_0/decoder_-1/"
        label_database = "data/processed/scannet200/label_database.yaml"
        OUTPUT = f"output-scans/{model_name}/{split}/{task}/"
        OUTPUT_2 = f"output-scans/gt/{split}/{task}/"
        OUTPUT_3 = f"output-scans/input/{split}/{task}/"
        with open(label_database) as f:
            x = yaml.safe_load(f) 
            
        # read label database
        cls_ids = x.keys()
        # from dict keys to list
        cls_ids = [int(i) for i in cls_ids]
        for file in os.listdir(PATH):
            pc = []
            pco = []
            SCENE = file.split(".")[0]
            # if SCENE in good_scenes:
            PATH_PLY = f"data/raw/scannet_test_segments/scans/{SCENE}/{SCENE}_vh_clean_2.labels.ply"
            PATH_PLY_1 = f"data/raw/scannet_test_segments/scans/{SCENE}/{SCENE}_vh_clean_2.ply"
            
            # read label database
            try:
                label = f"data/processed/scannet200/instance_gt/validation/{SCENE}.txt"
                # read label file
                with open(label) as f:
                    lines = f.readlines()
                unique_ids = np.unique(np.asarray(lines))
                unique_ids = [int(i[:-1]) for i in unique_ids]
                classes = [i//1000 for i in unique_ids if i !=0]
                # if task != 'task3':
                if UNKNOWN_CLASSES_LABELS[split][task] != None:
                    UNKNOWNS =  [label2id[i] for i in UNKNOWN_CLASSES_LABELS[split][task]]
                else:
                    UNKNOWNS = []
                    
                instances = []
                if not os.path.exists(OUTPUT):
                    os.makedirs(OUTPUT)
                if not os.path.exists(OUTPUT_2):
                    os.makedirs(OUTPUT_2)
                if not os.path.exists(OUTPUT_3):
                    os.makedirs(OUTPUT_3)  
                    
                with open(f"{PATH}{SCENE}.txt") as f:
                    instances = f.readlines()
                # sort instances by probablity instances.split(" ")[2]
                instances = sorted(instances, key=lambda x: float(x.split(" ")[2]), reverse=False)

                

                pc = read_ply_file(PATH_PLY)
                pco = read_ply_file(PATH_PLY_1)
                points = np.asarray(pc.vertices)
                indices = []
                for instance in instances:
                    path = instance.split(" ")[0].strip()
                    cls = instance.split(" ")[1].strip()
                    # if(int(cls) not in cls_ids):
                    #     continue
                    
                    with open(os.path.join(PATH, path), "r") as f:
                        mask = f.readlines()
                    for idx, point in enumerate(mask):
                        if(point == "1\n"):
                            indices.append(idx)
                            # change color of point based on its instance id
                            if(int(cls) == 0 or int(cls) == 1 or int(cls) == 3):
                                pc.vertex_colors[idx] = [0.5, 0.5, 0.5]
                            elif(int(cls) == 3000):
                                pc.vertex_colors[idx] = [31/255,119/255,180/255]
                            else:
                                pc.vertex_colors[idx] = [42/255,157/255,37/255]
                    indices = np.unique(np.asarray(indices))
                        # indices that are not in the list
                    left_indices = np.setdiff1d(np.arange(len(mask)), indices)
                    for index in left_indices:
                        pc.vertex_colors[index] = [0.5, 0.5, 0.5]
                    # save new point cloud
                    o3d.io.write_triangle_mesh(os.path.join(OUTPUT, SCENE + "-3d_owis.ply"), pc)
                            
                        
                        
                    pc = read_ply_file(PATH_PLY)
                    points = np.asarray(pc.vertices)
                    for i in range(points.shape[0]):
                    # change color of point based on its instance id
                        if(int(lines[i][:-1])//1000 == 0 or int(lines[i][:-1])//1000 == 3 or int(lines[i][:-1])//1000 == 1): # make gray if 0 (nothing) or 3 (floor) or 1 (wall)
                            pc.vertex_colors[i] = [0.5, 0.5, 0.5]
                        elif(int(lines[i][:-1])//1000 in UNKNOWNS):
                            pc.vertex_colors[i] = [31/255,119/255,180/255]#[23/255,124/255,172/255]
                        else:
                            pc.vertex_colors[i] = [42/255,157/255,37/255] # green
                    o3d.io.write_triangle_mesh(os.path.join(OUTPUT_2, SCENE + "-gt.ply"), pc)
                    o3d.io.write_triangle_mesh(os.path.join(OUTPUT_3, SCENE + "-input.ply"), pco)
            except:
                pass
visualization/convert_2_ply.py
aminebdj/3D-OWIS