├── README.md ├── data ├── dataset_creator.py ├── dataset_creator_azure.py ├── dataset_creator_parallel.py ├── dataset_creator_scannet_depth_completion.py └── merge_kinect_pickle_files.py ├── dataset.py ├── demo.sh ├── demo_dataset ├── color │ ├── color_000000.png │ ├── color_000017.png │ ├── color_000034.png │ ├── color_000051.png │ ├── color_000068.png │ ├── color_000085.png │ ├── color_000102.png │ └── color_000119.png ├── depth_sparse │ ├── depth_sparse_000000.txt │ ├── depth_sparse_000017.txt │ ├── depth_sparse_000034.txt │ ├── depth_sparse_000051.txt │ ├── depth_sparse_000068.txt │ ├── depth_sparse_000085.txt │ ├── depth_sparse_000102.txt │ └── depth_sparse_000119.txt └── gravity │ ├── gravity_000000.txt │ ├── gravity_000017.txt │ ├── gravity_000034.txt │ ├── gravity_000051.txt │ ├── gravity_000068.txt │ ├── gravity_000085.txt │ ├── gravity_000102.txt │ └── gravity_000119.txt ├── images ├── overview.jpg └── sample_results.jpg ├── main.py ├── network_run.py ├── networks ├── __init__.py ├── depth_completion.py ├── network_utils.py ├── surface_normal.py ├── surface_normal_dorn.py └── warping_2dof_alignment.py ├── normal_utils.py └── plane_mask_detection ├── ABSTRACTIONS.md ├── CODE_OF_CONDUCT.md ├── CONTRIBUTING.md ├── config.yml ├── configs └── R101_bs16_all_plane_normal.yaml ├── demo ├── README.md ├── demo_e2e_mask_rcnn_R_50_FPN_1x.png ├── demo_e2e_mask_rcnn_X_101_32x8d_FPN_1x.png ├── predictor.py └── webcam.py ├── maskrcnn_benchmark ├── __init__.py ├── config │ ├── __init__.py │ ├── defaults.py │ └── paths_catalog.py ├── csrc │ ├── ROIAlign.h │ ├── ROIPool.h │ ├── SigmoidFocalLoss.h │ ├── cpu │ │ ├── ROIAlign_cpu.cpp │ │ ├── nms_cpu.cpp │ │ └── vision.h │ ├── cuda │ │ ├── ROIAlign_cuda.cu │ │ ├── ROIPool_cuda.cu │ │ ├── SigmoidFocalLoss_cuda.cu │ │ ├── deform_conv_cuda.cu │ │ ├── deform_conv_kernel_cuda.cu │ │ ├── deform_pool_cuda.cu │ │ ├── deform_pool_kernel_cuda.cu │ │ ├── nms.cu │ │ └── vision.h │ ├── deform_conv.h │ ├── deform_pool.h │ ├── nms.h │ └── vision.cpp ├── data │ ├── README.md │ ├── __init__.py │ ├── build.py │ ├── collate_batch.py │ ├── datasets │ │ ├── __init__.py │ │ ├── coco.py │ │ ├── concat_dataset.py │ │ ├── evaluation │ │ │ ├── __init__.py │ │ │ ├── __init__backup.py │ │ │ ├── coco │ │ │ │ ├── __init__.py │ │ │ │ └── coco_eval.py │ │ │ └── voc │ │ │ │ ├── __init__.py │ │ │ │ └── voc_eval.py │ │ ├── fmd.py │ │ ├── list_dataset.py │ │ └── voc.py │ ├── samplers │ │ ├── __init__.py │ │ ├── distributed.py │ │ ├── grouped_batch_sampler.py │ │ └── iteration_based_batch_sampler.py │ └── transforms │ │ ├── __init__.py │ │ ├── build.py │ │ └── transforms.py ├── engine │ ├── .inference.py.swp │ ├── __init__.py │ ├── bbox_aug.py │ ├── inference.py │ └── trainer.py ├── layers │ ├── __init__.py │ ├── _utils.py │ ├── batch_norm.py │ ├── dcn │ │ ├── __init__.py │ │ ├── deform_conv_func.py │ │ ├── deform_conv_module.py │ │ ├── deform_pool_func.py │ │ └── deform_pool_module.py │ ├── misc.py │ ├── nms.py │ ├── roi_align.py │ ├── roi_pool.py │ ├── sigmoid_focal_loss.py │ └── smooth_l1_loss.py ├── modeling │ ├── __init__.py │ ├── backbone │ │ ├── __init__.py │ │ ├── backbone.py │ │ ├── fbnet.py │ │ ├── fbnet_builder.py │ │ ├── fbnet_modeldef.py │ │ ├── fpn.py │ │ └── resnet.py │ ├── balanced_positive_negative_sampler.py │ ├── box_coder.py │ ├── detector │ │ ├── __init__.py │ │ ├── detectors.py │ │ └── generalized_rcnn.py │ ├── make_layers.py │ ├── matcher.py │ ├── poolers.py │ ├── registry.py │ ├── roi_heads │ │ ├── __init__.py │ │ ├── box_head │ │ │ ├── __init__.py │ │ │ ├── box_head.py │ │ │ ├── inference.py │ │ │ ├── loss.py │ │ │ ├── roi_box_feature_extractors.py │ │ │ └── roi_box_predictors.py │ │ ├── keypoint_head │ │ │ ├── __init__.py │ │ │ ├── inference.py │ │ │ ├── keypoint_head.py │ │ │ ├── loss.py │ │ │ ├── roi_keypoint_feature_extractors.py │ │ │ └── roi_keypoint_predictors.py │ │ ├── mask_head │ │ │ ├── __init__.py │ │ │ ├── inference.py │ │ │ ├── loss.py │ │ │ ├── mask_head.py │ │ │ ├── mask_upsampling_grad.py │ │ │ ├── roi_mask_feature_extractors.py │ │ │ └── roi_mask_predictors.py │ │ └── roi_heads.py │ ├── rpn │ │ ├── __init__.py │ │ ├── anchor_generator.py │ │ ├── inference.py │ │ ├── loss.py │ │ ├── retinanet │ │ │ ├── __init__.py │ │ │ ├── inference.py │ │ │ ├── loss.py │ │ │ └── retinanet.py │ │ ├── rpn.py │ │ └── utils.py │ ├── upconv │ │ ├── GravityUpConvNet.py │ │ ├── ResGravityUpConvNet.py │ │ ├── UpConvNet.py │ │ ├── __init__.py │ │ └── dorn.py │ └── utils.py ├── solver │ ├── __init__.py │ ├── build.py │ └── lr_scheduler.py ├── structures │ ├── __init__.py │ ├── bounding_box.py │ ├── boxlist_ops.py │ ├── image_list.py │ ├── keypoint.py │ └── segmentation_mask.py └── utils │ ├── README.md │ ├── __init__.py │ ├── c2_model_loading.py │ ├── checkpoint.py │ ├── collect_env.py │ ├── comm.py │ ├── cv2_util.py │ ├── env.py │ ├── imports.py │ ├── logger.py │ ├── metric_logger.py │ ├── miscellaneous.py │ ├── model_serialization.py │ ├── model_zoo.py │ ├── registry.py │ └── timer.py ├── tests ├── checkpoint.py ├── env_tests │ └── env.py ├── test_backbones.py ├── test_box_coder.py ├── test_configs.py ├── test_data_samplers.py ├── test_detectors.py ├── test_fbnet.py ├── test_feature_extractors.py ├── test_metric_logger.py ├── test_nms.py ├── test_predictors.py ├── test_rpn_heads.py ├── test_segmentation_mask.py └── utils.py └── tools ├── cityscapes ├── convert_cityscapes_to_coco.py └── instances2dict_with_polygons.py ├── test_net.py └── train_net.py /data/dataset_creator.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | 3 | import os 4 | import pickle 5 | 6 | ROOT_DIR = '/mars/mnt/dgx/FrameNet' #'/mnt/dgx/nyud_v2' #'/mnt/dgx/FrameNet' 7 | 8 | def create_split(split_name): 9 | if split_name == 'train': 10 | start_scene_id = 0 11 | end_scene_id = 10 12 | elif split_name == 'test': 13 | start_scene_id = 0 14 | end_scene_id = 10 15 | 16 | image_id = 1 17 | frame_gap = 20 18 | scans_dir = 'scannet-small-frames' # "nyud-frames" # 'evaluation_scans' # 'scans' # 'scannet-frames' 19 | final_split = [[], [], []] 20 | num_frames = 0 21 | for scene_id in range(start_scene_id, end_scene_id): 22 | scan_id = 0 23 | while os.path.exists(os.path.join(ROOT_DIR, "%s/scene%04d_%02d" % (scans_dir, scene_id, scan_id))): 24 | frame_folder = os.path.join(ROOT_DIR, "%s/scene%04d_%02d" % (scans_dir, scene_id, scan_id)) 25 | print('Working on folder: ', frame_folder) 26 | if split_name == 'train': 27 | frame_id = 0 28 | else: 29 | frame_id = 10 30 | while os.path.isfile(os.path.join(frame_folder, 'frame-%06d-normal.png' % frame_id)): 31 | color_path = os.path.join(frame_folder, 'frame-%06d-color.png' % frame_id) 32 | orient_path = os.path.join(frame_folder, 'frame-%06d-normal.png' % frame_id) 33 | mask_path = os.path.join(frame_folder, 'frame-%06d-orient-mask.png' % frame_id) 34 | final_split[0].append(color_path) 35 | final_split[1].append(orient_path) 36 | final_split[2].append(mask_path) 37 | frame_id += frame_gap 38 | num_frames += 1 39 | scan_id += 1 40 | print('Number of frames in the %s split: %d' % (split_name, num_frames)) 41 | return final_split 42 | 43 | def main(): 44 | final_dict = {'train': create_split('train'), 'test': create_split('test')} 45 | with open('first10scenes_train_test_split.pkl', 'wb') as f: 46 | pickle.dump(final_dict, f) 47 | 48 | if __name__ == "__main__": 49 | main() 50 | -------------------------------------------------------------------------------- /data/dataset_creator_azure.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | 3 | ### Creates a pickle file for a single Azure Kinect dataset. 4 | ### To create a pickle file for all Azure datasets, run this script on every dataset and 5 | ### merge the pickle files using merge_kinect_pickle_files.py script. 6 | 7 | import argparse 8 | import os 9 | import pickle 10 | 11 | parser = argparse.ArgumentParser(description='Dataset root directory') 12 | parser.add_argument('--root', type=str, required=True) 13 | parser.add_argument('--check_for_percentage', action='store_true') 14 | parser.add_argument('--percentage', type=float, default=8.0) 15 | args = parser.parse_args() 16 | 17 | 18 | def create_split_azure(split_name): 19 | root_dir = os.path.abspath(args.root) 20 | image_id = 0 21 | final_split = [[], [], [], [], []] 22 | while os.path.exists(os.path.join(root_dir, "color/color_%06d.png" % image_id)): 23 | # Make the train and test both contain all the images for now 24 | color_path = os.path.join(root_dir, "color/color_%06d.png" % image_id) 25 | depth_path = os.path.join(root_dir, "depth/depth_%06d.png" % image_id) 26 | gravity_path = os.path.join(root_dir, "gravity/gravity_%06d.txt" % image_id) 27 | mask_path = os.path.join(root_dir, "mask/mask_%06d.png" % image_id) 28 | normal_path = os.path.join(root_dir, "normal/normal_%06d.png" % image_id) 29 | 30 | ptcloud_file = os.path.join(root_dir, 'klt_tracks_triangulation/frame-{0:06d}-observed-pointcloud.txt'.format(image_id)) 31 | if os.path.isfile(ptcloud_file): 32 | final_split[0].append(color_path) 33 | final_split[1].append(depth_path) 34 | final_split[2].append(gravity_path) 35 | final_split[3].append(mask_path) 36 | final_split[4].append(normal_path) 37 | image_id += 1 38 | 39 | print('Number of frames in the %s split: %d' % (split_name, len(final_split[0]))) 40 | return final_split 41 | 42 | 43 | def main(): 44 | # Right now, the train is empty. 45 | final_dict = {'train': [[], [], [], [], []], 'test': create_split_azure('test')} 46 | dataset_name = os.path.basename(os.path.normpath(args.root)) 47 | file_name = dataset_name + '_train_test_split_trinagulated.pkl' 48 | with open(file_name, 'wb') as f: 49 | pickle.dump(final_dict, f) 50 | print('Pickle file saved: ' + file_name) 51 | 52 | 53 | if __name__ == "__main__": 54 | main() 55 | -------------------------------------------------------------------------------- /data/dataset_creator_parallel.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | 3 | import os 4 | import pickle 5 | import numpy as np 6 | from PIL import Image 7 | from multiprocessing import Pool 8 | 9 | ROOT_DIR = '/mars/mnt/dgx/FrameNet' #'/mnt/dgx/nyud_v2' #'/mnt/dgx/FrameNet' 10 | PERCENTAGE = 8.0 11 | 12 | 13 | def inc_depth_has_enough(incomp_depth_filename): 14 | I = np.asarray(Image.open(incomp_depth_filename)) 15 | ratio = 100 * np.sum(I > 0) / I.size 16 | if ratio < PERCENTAGE: 17 | return False 18 | else: 19 | return True 20 | 21 | 22 | def process_single_scene(input_ids): 23 | split_name, scene_id, scan_id = input_ids 24 | scans_dir = 'scannet-small-frames' # "nyud-frames" # 'evaluation_scans' # 'scans' # 'scannet-frames' 25 | final_split = [[], [], []] 26 | frame_gap = 10 27 | 28 | frame_folder = os.path.join(ROOT_DIR, "%s/scene%04d_%02d" % (scans_dir, scene_id, scan_id)) 29 | print('Working on folder: ', frame_folder) 30 | if split_name == 'train': 31 | frame_id = 0 32 | else: 33 | frame_id = 10 34 | while os.path.isfile(os.path.join(frame_folder, 'frame-%06d-normal.png' % frame_id)): 35 | # Only add the frame if the incomplete depth exists 36 | 37 | if scene_id < 707: 38 | # Train scan 39 | di_path = frame_folder.replace('FrameNet/scannet-small-frames', 'ScanNet/scans') 40 | else: 41 | di_path = frame_folder.replace('FrameNet/scannet-small-frames', 'ScanNet/scans_test/scans') 42 | 43 | di_filename = os.path.join(di_path, 'incomplete_depth_abs_bugfix', 'depth-{0:06d}-incomplete.png'.format(frame_id)) 44 | 45 | if os.path.isfile(di_filename) and inc_depth_has_enough(di_filename): 46 | color_path = os.path.join(frame_folder, 'frame-%06d-color.png' % frame_id) 47 | orient_path = os.path.join(frame_folder, 'frame-%06d-normal.png' % frame_id) 48 | mask_path = os.path.join(frame_folder, 'frame-%06d-orient-mask.png' % frame_id) 49 | final_split[0].append(color_path) 50 | final_split[1].append(orient_path) 51 | final_split[2].append(mask_path) 52 | 53 | frame_id += frame_gap 54 | 55 | return final_split 56 | 57 | 58 | def create_split(split_name): 59 | num_parallel_workers = 40 60 | if split_name == 'train': 61 | start_scene_id = 0 62 | end_scene_id = 707 63 | elif split_name == 'test': 64 | start_scene_id = 707 65 | end_scene_id = 807 66 | 67 | scans_dir = 'scannet-small-frames' # "nyud-frames" # 'evaluation_scans' # 'scans' # 'scannet-frames' 68 | 69 | all_pairs = [] 70 | for scene_id in range(start_scene_id, end_scene_id): 71 | scan_id = 0 72 | while os.path.exists(os.path.join(ROOT_DIR, "%s/scene%04d_%02d" % (scans_dir, scene_id, scan_id))): 73 | all_pairs.append((split_name, scene_id, scan_id)) 74 | scan_id += 1 75 | 76 | pool = Pool(num_parallel_workers) 77 | all_outputs = pool.map(process_single_scene, all_pairs) 78 | 79 | final_split = [[], [], []] 80 | for l1, l2, l3 in all_outputs: 81 | final_split[0].extend(l1) 82 | final_split[1].extend(l2) 83 | final_split[2].extend(l3) 84 | 85 | print('Number of frames in the %s split: %d' % (split_name, len(final_split[0]))) 86 | return final_split 87 | 88 | 89 | def main(): 90 | final_dict = {'train': create_split('train'), 'test': create_split('test')} 91 | with open('scannet_inc_depth_train_test_split.pkl', 'wb') as f: 92 | pickle.dump(final_dict, f) 93 | 94 | if __name__ == "__main__": 95 | main() 96 | -------------------------------------------------------------------------------- /data/dataset_creator_scannet_depth_completion.py: -------------------------------------------------------------------------------- 1 | import os 2 | import pickle 3 | from PIL import Image 4 | import numpy as np 5 | 6 | 7 | ROOT_DIR = '/mars/mnt/dgx/FrameNet' 8 | NORMALS_DIR = '/mars/mnt/dgx/FrameNet/DORN_acos_bs16_inference/' #'/mars/mnt/dgx/FrameNet/L2_inference/' 9 | PLANES_DIR = '/mars/mnt/dgx/FrameNet/plane_inference/' 10 | OUTPUT_FILE_NAME = 'scannet_depth_completion_split_acos.pkl' 11 | 12 | INCOMPLETE_DEPTH_DIRECTORY='incomplete_depth_dorn_acos' 13 | 14 | CHECK_FOR_MIN_PERCENT = True 15 | MIN_PERCENT = 8 16 | 17 | 18 | def check_min_perc(scannet_scene_base, scene_name, frame_id): 19 | inc_depth_filename = os.path.join(scannet_scene_base, scene_name, '{0}/depth-{1:06d}-incomplete.png'.format(INCOMPLETE_DEPTH_DIRECTORY, frame_id)) 20 | depths = np.asarray(Image.open(inc_depth_filename)) 21 | valid_depths = (depths > 0).astype(np.float) 22 | num_valid_depths = np.sum(valid_depths) 23 | return ((100.0 * num_valid_depths / depths.size) >= MIN_PERCENT) 24 | 25 | 26 | def create_split(split_name): 27 | if split_name == 'train': 28 | start_scene_id = 0 29 | end_scene_id = 707 30 | elif split_name == 'test': 31 | start_scene_id = 707 32 | end_scene_id = 807 33 | 34 | image_id = 1 35 | frame_gap = 10 36 | scans_dir = 'scannet-small-frames' # "nyud-frames" # 'evaluation_scans' # 'scans' # 'scannet-frames' 37 | final_split = [[], [], []] 38 | num_frames = 0 39 | for scene_id in range(start_scene_id, end_scene_id): 40 | scan_id = 0 41 | while os.path.exists(os.path.join(ROOT_DIR, "%s/scene%04d_%02d" % (scans_dir, scene_id, scan_id))): 42 | frame_folder = os.path.join(ROOT_DIR, "%s/scene%04d_%02d" % (scans_dir, scene_id, scan_id)) 43 | #print('Working on folder: ', frame_folder) 44 | if split_name == 'train': 45 | frame_id = 0 46 | else: 47 | frame_id = 10 48 | while os.path.isfile(os.path.join(frame_folder, 'frame-%06d-normal.png' % frame_id)): 49 | color_path = os.path.join(frame_folder, 'frame-%06d-color.png' % frame_id) 50 | orient_path = os.path.join(frame_folder, 'frame-%06d-normal.png' % frame_id) 51 | mask_path = os.path.join(frame_folder, 'frame-%06d-orient-mask.png' % frame_id) 52 | plane_path = os.path.join(PLANES_DIR, '{0}/scene{1:04d}_{2:02d}/frame-{3:06d}-instance_segmentation.png'.format(PLANES_DIR, scene_id, scan_id, frame_id)) 53 | # Check if this frame has the incomplete depth. If so, then it also has the inference normal and other related. 54 | if scene_id < 707: 55 | scannet_scene_base = '/mars/mnt/dgx/ScanNet/scans/' 56 | else: 57 | scannet_scene_base = '/mars/mnt/dgx/ScanNet/scans_test/scans' 58 | scene_name = 'scene%04d_%02d' % (scene_id, scan_id) 59 | incomplete_depth_path = os.path.join(scannet_scene_base, scene_name, '{0}/depth-{1:06d}-incomplete.png'.format(INCOMPLETE_DEPTH_DIRECTORY, frame_id)) 60 | if os.path.isfile(color_path) and os.path.isfile(orient_path) and os.path.isfile(mask_path) and os.path.isfile(plane_path) and os.path.isfile(incomplete_depth_path): 61 | should_add = False 62 | if not CHECK_FOR_MIN_PERCENT: 63 | should_add = True 64 | elif check_min_perc(scannet_scene_base, scene_name, frame_id): 65 | # In this case CHECK_FOR_MIN_PERCENT is True. 66 | should_add = True 67 | 68 | if should_add: 69 | final_split[0].append(color_path) 70 | final_split[1].append(orient_path) 71 | final_split[2].append(mask_path) 72 | frame_id += frame_gap 73 | num_frames += 1 74 | scan_id += 1 75 | 76 | print('Scene {}: {}'.format(scene_id, len(final_split[0]))) 77 | 78 | print('Number of frames in the %s split: %d' % (split_name, num_frames)) 79 | return final_split 80 | 81 | def main(): 82 | final_dict = {'train': create_split('train'), 'test': create_split('test')} 83 | with open(OUTPUT_FILE_NAME, 'wb') as f: 84 | pickle.dump(final_dict, f) 85 | 86 | if __name__ == "__main__": 87 | main() 88 | 89 | -------------------------------------------------------------------------------- /data/merge_kinect_pickle_files.py: -------------------------------------------------------------------------------- 1 | import os 2 | import pickle 3 | import sys 4 | 5 | ### FORMAT: python merge_kinect_pickle_files.py ... output.pickle 6 | 7 | if __name__ == '__main__': 8 | # assumes that each picle file contains a dictionary of {'train':->list of lists, 'test'->list os lists} 9 | num_sublists = -1 10 | 11 | assert len(sys.argv) > 3 12 | 13 | for index in range(1, len(sys.argv) - 1): 14 | with open(sys.argv[index], 'rb') as fp: 15 | data = pickle.load(fp) 16 | 17 | if num_sublists == -1: 18 | num_sublists = len(data['test']) 19 | final_map = {'test':[[] for i in range(num_sublists)], 'train':[[] for i in range(num_sublists)]} 20 | 21 | for i in range(num_sublists): 22 | final_map['test'][i].extend(data['test'][i]) 23 | final_map['train'][i].extend(data['test'][i]) 24 | 25 | 26 | with open(sys.argv[-1], 'wb') as f: 27 | pickle.dump(final_map, f) 28 | -------------------------------------------------------------------------------- /demo.sh: -------------------------------------------------------------------------------- 1 | #!/bin/sh 2 | 3 | ### Evaluation script. If running for ScanNet or Kinect Azure, please first generate the pickle files 4 | ### by running the corresponding script in the data directory. 5 | ### A description of each command line option is provided the in main.py file. 6 | 7 | python main.py --checkpoint ./checkpoints/depth_completion.ckpt \ 8 | --surface_normal_checkpoint ./checkpoints/surface_normal_fpn_use_gravity.ckpt \ 9 | --plane_detection_config_file ./plane_mask_detection/configs/R101_bs16_all_plane_normal.yaml \ 10 | --train 0 --dataset_type demo --batch_size 1 --use_gravity 1 \ 11 | --root ./demo_dataset --save ./demo_dataset 12 | 13 | 14 | # w/o gravity, using DORN 15 | # ScanNet 16 | #python main.py --checkpoint ./checkpoints/depth_completion.ckpt \ 17 | # --surface_normal_checkpoint ./checkpoints/surface_normal_dorn.ckpt \ 18 | # --plane_detection_config_file ./plane_mask_detection/configs/R101_bs16_all_plane_normal.yaml \ 19 | # --train 0 --dataset_type scannet --batch_size 16 --enriched_samples 100 20 | 21 | # Azure Kinect 22 | #python main.py --checkpoint ./checkpoints/depth_completion.ckpt \ 23 | # --surface_normal_checkpoint ./checkpoints/surface_normal_dorn.ckpt \ 24 | # --plane_detection_config_file ./plane_mask_detection/configs/R101_bs16_all_plane_normal.yaml \ 25 | # --dataset_pickle_file ./data/all_kinect_triangulated_split.pkl \ 26 | # --train 0 --dataset_type azure --batch_size 16 --enriched_samples 100 --skip_every_n_image_test 1 27 | 28 | # 29 | # w gravity, using FPN warping gravity alignment 30 | # ScanNet 31 | #python main.py --checkpoint ./checkpoints/depth_completion.ckpt \ 32 | # --surface_normal_checkpoint ./checkpoints/surface_normal_fpn_use_gravity.ckpt \ 33 | # --plane_detection_config_file ./plane_mask_detection/configs/R101_bs16_all_plane_normal.yaml \ 34 | # --train 0 --dataset_type scannet --batch_size 16 --use_gravity 1 --enriched_samples 100 35 | 36 | # Azure Kienct 37 | #python main.py --checkpoint ./checkpoints/depth_completion.ckpt \ 38 | # --surface_normal_checkpoint ./checkpoints/surface_normal_fpn_use_gravity.ckpt \ 39 | # --plane_detection_config_file ./plane_mask_detection/configs/R101_bs16_all_plane_normal.yaml \ 40 | # --dataset_pickle_file ./data/all_kinect_triangulated_split.pkl \ 41 | # --train 0 --dataset_type azure --batch_size 16 --use_gravity 1 --enriched_samples 100 --skip_every_n_image_test 1 42 | # 43 | -------------------------------------------------------------------------------- /demo_dataset/color/color_000000.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MARSLab-UMN/vi_depth_completion/b8d4bd3f3af674871a9fb8cd6166e49d785859ef/demo_dataset/color/color_000000.png -------------------------------------------------------------------------------- /demo_dataset/color/color_000017.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MARSLab-UMN/vi_depth_completion/b8d4bd3f3af674871a9fb8cd6166e49d785859ef/demo_dataset/color/color_000017.png -------------------------------------------------------------------------------- /demo_dataset/color/color_000034.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MARSLab-UMN/vi_depth_completion/b8d4bd3f3af674871a9fb8cd6166e49d785859ef/demo_dataset/color/color_000034.png -------------------------------------------------------------------------------- /demo_dataset/color/color_000051.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MARSLab-UMN/vi_depth_completion/b8d4bd3f3af674871a9fb8cd6166e49d785859ef/demo_dataset/color/color_000051.png -------------------------------------------------------------------------------- /demo_dataset/color/color_000068.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MARSLab-UMN/vi_depth_completion/b8d4bd3f3af674871a9fb8cd6166e49d785859ef/demo_dataset/color/color_000068.png -------------------------------------------------------------------------------- /demo_dataset/color/color_000085.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MARSLab-UMN/vi_depth_completion/b8d4bd3f3af674871a9fb8cd6166e49d785859ef/demo_dataset/color/color_000085.png -------------------------------------------------------------------------------- /demo_dataset/color/color_000102.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MARSLab-UMN/vi_depth_completion/b8d4bd3f3af674871a9fb8cd6166e49d785859ef/demo_dataset/color/color_000102.png -------------------------------------------------------------------------------- /demo_dataset/color/color_000119.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MARSLab-UMN/vi_depth_completion/b8d4bd3f3af674871a9fb8cd6166e49d785859ef/demo_dataset/color/color_000119.png -------------------------------------------------------------------------------- /demo_dataset/depth_sparse/depth_sparse_000068.txt: -------------------------------------------------------------------------------- 1 | 4 0.733004 0.420492 1.66152 2 | 15 0.20083 0.173569 1.67941 3 | 17 0.946681 0.336118 1.69906 4 | 88 0.721861 0.265034 1.64737 5 | 93 0.310007 0.399798 1.22424 6 | 105 -0.326011 -0.0764085 1.80095 7 | 122 0.648969 0.399174 1.30774 8 | 139 0.124199 0.463365 1.22974 9 | 144 -0.422995 -0.669078 1.71258 10 | 169 -0.0489958 -0.678824 1.67842 11 | 174 -0.128325 -0.0841832 1.63808 12 | 199 -0.363031 -0.561411 1.50373 13 | 204 -0.348672 0.407058 1.48474 14 | 206 0.792083 0.414665 1.7199 15 | 213 -0.37707 0.31559 1.35348 16 | 215 0.0867302 0.2806 1.43306 17 | 221 -0.379506 0.459225 1.35874 18 | 226 -0.293904 -0.216265 1.7305 19 | 296 0.231678 0.323745 2.22302 20 | 321 -0.389072 0.404709 1.31288 21 | 327 -0.38682 0.559708 1.33966 22 | 345 -0.332325 -0.269135 1.78763 23 | 346 0.126421 0.356146 1.39101 24 | 377 -0.0300167 0.38217 1.31752 25 | 380 -0.363 0.173828 1.72222 26 | 388 0.680631 0.323799 1.66114 27 | 391 -0.204842 0.386332 1.28656 28 | 418 -0.0305692 -0.017497 2.08063 29 | 427 0.876303 -0.411215 1.92679 30 | 539 0.127386 0.126247 1.4759 31 | 559 -0.0699734 -0.122456 1.82992 32 | 640 -0.299412 0.372813 1.51903 33 | 886 0.514243 0.480971 1.22612 34 | 1219 0.664442 -0.367896 1.57625 35 | 1524 -0.117774 -0.238288 1.68261 36 | 1558 0.717321 0.387265 1.69249 37 | 1607 0.289285 0.521597 1.22106 38 | 1855 -0.139282 0.468001 1.2439 39 | 1866 0.216461 0.172835 1.66863 40 | 1877 -0.410766 0.349037 1.62994 41 | 1883 0.686364 0.182869 1.61479 42 | 1885 0.13802 0.506988 1.25812 43 | 1886 0.0700733 0.466783 1.24097 44 | 1901 -0.284162 -0.12955 1.78514 45 | 1939 0.0842551 0.22096 1.57686 46 | 1941 0.129156 0.621926 1.212 47 | 1986 -0.124012 -0.18048 1.92213 48 | 2115 -0.38477 0.00201943 1.74175 49 | 2157 -0.204689 -0.242569 1.74016 50 | 2220 0.718284 -0.404384 1.63782 51 | 2254 0.686725 0.34035 1.64525 52 | 2261 0.839985 0.378608 1.57494 53 | 2280 0.813002 -0.443162 1.86254 54 | 2293 1.01914 0.25513 1.72905 55 | 2356 0.847708 -0.33952 1.82322 56 | 2400 0.743597 -0.351628 1.59711 57 | 2408 0.85308 -0.368711 1.82136 58 | 2485 0.848242 -0.367043 1.90821 59 | 2567 1.01541 0.419312 1.69574 60 | 2592 0.211785 0.369347 1.37336 61 | 2602 0.888248 -0.480105 1.95322 62 | 2608 0.940398 -0.457014 1.85526 63 | 2633 -0.0906823 0.436326 1.29922 64 | 2792 0.0718618 0.139842 1.31984 65 | 2849 0.840229 -0.422828 1.74812 66 | 2894 0.125399 0.332524 1.41006 67 | 2914 1.05283 0.354751 1.72169 68 | 2928 1.17786 0.415632 1.75985 69 | 2933 1.18925 0.35121 1.79557 70 | 2941 1.06677 0.304612 1.65775 71 | 2965 1.17058 0.35993 1.74186 72 | 3062 1.36776 0.200131 1.69351 73 | 3103 1.34505 0.248687 1.79035 74 | 3188 1.40369 0.281821 1.80417 75 | 3217 0.140339 1.02814 2.12279 76 | 3261 0.798601 0.363095 1.75413 77 | 3269 0.538914 0.561078 1.28302 78 | 3288 0.73979 -0.346379 1.43783 79 | 3314 0.952445 0.390554 1.6421 80 | 3319 0.79912 -0.364178 1.84671 81 | 3445 0.0991035 0.311528 1.60207 82 | 3448 1.03757 0.418133 1.42428 83 | 3532 0.650325 0.285807 1.6125 84 | 3589 1.24846 0.423432 1.79083 85 | 3628 0.921766 -0.444618 1.87305 86 | 3654 1.16942 0.397448 1.81083 87 | 3670 -0.316149 0.653164 1.34586 88 | 3698 -0.194043 0.511761 1.7262 89 | 3866 1.48193 0.131258 1.86124 90 | 3925 1.28924 0.364411 1.65043 91 | 3943 1.54847 0.127082 1.89145 92 | 3958 0.124362 0.283876 1.47259 93 | 3965 1.19809 0.00174206 1.82317 94 | 3971 -0.237255 -0.223506 1.61947 95 | 3997 1.38691 0.407734 1.70128 96 | 4014 1.26205 0.1625 1.88224 97 | 4036 1.47722 0.20834 1.78576 98 | 4067 0.93203 0.575259 2.44143 99 | 4153 0.201522 0.295175 1.75306 100 | 4156 0.679294 0.415975 1.67675 101 | 4168 1.70787 0.223839 1.92276 102 | 4169 1.50537 0.107081 1.69756 103 | 4173 1.4412 0.300916 1.86257 104 | 4179 -0.281307 -0.141605 1.64245 105 | 4184 -0.0693545 0.560009 1.29313 106 | 4185 -0.212344 0.438165 1.52619 107 | 4190 0.713715 0.2344 1.65088 108 | 4200 0.73502 -0.405394 1.7094 109 | 4203 0.0987847 0.498798 1.23323 110 | 4209 -0.318212 0.840758 1.74713 111 | 4220 -0.130711 0.35236 1.46778 112 | 4226 0.316224 0.682039 1.28439 113 | 4236 -0.26943 -0.185103 1.55082 114 | 4293 -0.148927 0.643378 1.33339 115 | 4304 0.847252 -0.408634 1.88784 116 | 4309 0.914076 0.33819 1.44317 117 | 4356 1.21206 0.18634 1.57278 118 | 4384 -0.0891361 0.341754 1.43349 119 | 4398 0.0834767 0.242179 1.54378 120 | 4403 0.757821 0.369545 1.69442 121 | 4407 1.26454 0.306507 1.62171 122 | 4410 0.593492 0.394851 1.34521 123 | 4414 -0.141424 0.597366 1.40373 124 | 4415 1.06183 0.454617 1.89588 125 | 4419 -0.203292 0.642385 1.32325 126 | 4421 -0.190114 -0.189114 1.85773 127 | 4439 -0.40275 0.000567454 1.87329 128 | 4442 -0.152139 0.393434 0.822804 129 | -------------------------------------------------------------------------------- /demo_dataset/depth_sparse/depth_sparse_000085.txt: -------------------------------------------------------------------------------- 1 | 4 0.358715 0.642163 1.27073 2 | 17 0.553274 0.567592 1.37452 3 | 88 0.346884 0.486349 1.27928 4 | 105 -0.715406 0.161732 1.21317 5 | 144 -0.804479 -0.438549 1.19902 6 | 169 -0.434503 -0.449191 1.26338 7 | 206 0.40086 0.646608 1.3424 8 | 215 -0.212405 0.459237 0.910463 9 | 345 -0.723891 -0.0306625 1.2297 10 | 427 0.405165 -0.133647 1.69334 11 | 436 -0.157987 0.520262 1.33671 12 | 559 -0.476574 0.124098 1.31346 13 | 1756 -0.218128 0.481532 1.20716 14 | 1939 -0.253037 0.42371 1.0565 15 | 1951 0.310605 0.521472 1.26526 16 | 2220 0.325847 -0.175732 1.37639 17 | 2254 0.315699 0.559891 1.25624 18 | 2293 0.613311 0.49343 1.43455 19 | 2356 0.405916 -0.0800689 1.57581 20 | 2400 0.362218 -0.129984 1.33558 21 | 2482 0.362528 0.598022 1.25655 22 | 2485 0.384063 -0.0934196 1.6614 23 | 2567 0.62305 0.649969 1.37557 24 | 2894 -0.167635 0.507201 0.890077 25 | 2914 0.65072 0.590935 1.42019 26 | 2928 0.763754 0.658709 1.47871 27 | 2933 0.763787 0.601084 1.52599 28 | 2965 0.759609 0.600742 1.46862 29 | 2966 0.584889 0.341409 1.30102 30 | 3103 0.912691 0.500963 1.57715 31 | 3188 0.966869 0.536608 1.59998 32 | 3261 0.396943 0.601362 1.38496 33 | 3376 0.64628 0.483609 1.40261 34 | 3445 -0.242367 0.517342 1.06984 35 | 3512 0.723011 0.644153 1.41675 36 | 3654 0.742125 0.648944 1.52809 37 | 3866 1.02355 0.398265 1.69844 38 | 4014 0.806566 0.429843 1.65741 39 | 4304 0.387029 -0.137776 1.64838 40 | 4347 0.951919 0.602183 1.64075 41 | 4377 0.474672 -0.186653 1.8348 42 | 4415 0.61825 0.717856 1.5726 43 | 4455 0.546505 0.359973 1.17677 44 | 4489 0.387724 0.585436 1.50614 45 | 4504 0.459909 -0.18176 2.15247 46 | 4644 0.464458 -0.282848 1.9224 47 | 4769 0.528481 -0.123312 1.8417 48 | 4813 -0.193982 0.519309 1.20991 49 | 4823 1.00028 0.637208 1.62704 50 | 4841 -0.168956 0.458552 1.18587 51 | 4846 1.28882 0.446528 1.8784 52 | 4892 1.05455 0.427098 1.74581 53 | 4893 -0.157261 0.423954 1.29104 54 | 4905 -0.558521 0.0925833 1.29648 55 | 4909 -0.669016 0.069585 1.45763 56 | 4913 1.00449 0.683938 1.6143 57 | 4926 1.38029 0.237048 1.90711 58 | 4927 1.27102 0.634952 1.70206 59 | 4931 1.43263 0.705146 1.89573 60 | 4955 1.44431 0.492126 1.80254 61 | 4958 1.20094 0.620294 1.56291 62 | 4967 0.465629 0.728149 1.67711 63 | 4977 1.13846 0.763278 1.6859 64 | 4980 1.23744 0.547101 1.48633 65 | 4984 0.875853 0.54238 1.88229 66 | 4988 0.715576 0.595026 1.35046 67 | 5053 -0.840401 0.26734 1.23749 68 | 5066 1.97721 0.792346 2.92762 69 | 5074 0.455223 -0.246774 1.75207 70 | 5101 -0.228849 0.510596 1.53933 71 | 5103 0.549686 -0.0898608 1.84573 72 | 5125 -0.257758 0.50069 1.13847 73 | 5130 0.423085 -0.168979 1.91548 74 | 5138 -0.580551 0.0641656 1.39588 75 | 5140 -0.800757 0.13414 1.44216 76 | 5146 0.408678 0.537012 1.4728 77 | 5160 1.08957 0.593441 2.14792 78 | 5162 1.4292 0.766861 1.54105 79 | 5184 0.446792 0.693001 1.50535 80 | 5185 1.78475 -0.687446 2.20421 81 | 5198 0.454479 0.775891 1.66186 82 | 5199 1.94491 1.02126 2.02009 83 | 5200 1.81459 0.448921 1.96773 84 | 5207 -0.180691 0.571506 1.42299 85 | 5212 0.404243 0.554891 1.11867 86 | 5218 1.97658 0.120964 2.19828 87 | 5226 -0.216958 0.55152 1.27134 88 | 5229 0.507999 0.787406 1.74966 89 | 5230 1.31318 0.677955 1.36561 90 | -------------------------------------------------------------------------------- /demo_dataset/depth_sparse/depth_sparse_000102.txt: -------------------------------------------------------------------------------- 1 | 4 0.104625 0.29223 0.935689 2 | 15 -0.376784 0.0609284 0.692286 3 | 17 0.280528 0.198354 1.05582 4 | 88 0.103013 0.139121 0.903346 5 | 206 0.131708 0.279315 1.01336 6 | 296 -0.592537 0.159881 1.20481 7 | 418 -0.759679 -0.158852 0.929185 8 | 436 -0.404553 0.145409 0.862371 9 | 539 -0.351588 0.0346154 0.473833 10 | 1756 -0.435532 0.139449 0.717789 11 | 2254 0.0714625 0.215345 0.893052 12 | 2356 0.145811 -0.482488 1.05833 13 | 2485 0.108924 -0.517653 1.1317 14 | 2567 0.342455 0.279384 1.09124 15 | 2914 0.365324 0.211622 1.12481 16 | 2928 0.459581 0.264955 1.22003 17 | 2933 0.454793 0.197228 1.25078 18 | 3103 0.59808 0.0909397 1.30542 19 | 3188 0.644029 0.120884 1.34689 20 | 3261 0.123042 0.224664 1.04185 21 | 3376 0.372367 0.112208 1.08106 22 | 3448 0.482859 0.302306 0.85898 23 | 3654 0.429628 0.24246 1.26 24 | 4347 0.616654 0.17358 1.39846 25 | 4504 0.0950337 -0.726021 1.59116 26 | 4704 0.155408 -0.890514 1.86333 27 | 4750 0.484257 -0.147402 1.42118 28 | 4823 0.664005 0.212071 1.40394 29 | 4838 0.392876 0.255971 1.30634 30 | 4840 0.698133 0.075019 1.46389 31 | 4846 0.912055 -0.0293434 1.65491 32 | 4882 0.396827 0.283232 1.21047 33 | 4888 0.752846 -0.0143964 1.49409 34 | 4892 0.709861 -0.0199779 1.47639 35 | 4893 -0.387862 0.0639046 0.795746 36 | 4920 0.617032 0.135488 1.41717 37 | 4924 0.921734 0.0712152 1.65386 38 | 4926 1.01163 -0.237027 1.6498 39 | 4927 0.9146 0.197219 1.53009 40 | 4931 1.03014 0.219709 1.76397 41 | 4945 0.618044 0.0878243 1.33884 42 | 4963 0.0688895 0.201301 1.0759 43 | 4977 0.778459 0.322278 1.51886 44 | 4984 0.500117 0.0526057 1.59717 45 | 5066 1.35751 0.0547598 2.87508 46 | 5103 0.235165 -0.557126 1.34129 47 | 5128 0.0586974 0.239075 1.07719 48 | 5184 0.141774 0.283868 1.18855 49 | 5264 0.191291 -0.565124 1.1903 50 | 5322 0.8327 0.329159 1.57979 51 | 5382 -0.546382 0.0381901 0.910414 52 | 5466 0.084889 0.0971284 1.27746 53 | 5564 0.312709 0.452004 1.00846 54 | 5571 0.533791 0.261065 1.52987 55 | 5625 0.134042 0.371108 0.841156 56 | 5681 0.377313 0.46518 1.20728 57 | 5690 1.07702 0.204284 1.68709 58 | 5730 0.0823942 0.25882 1.12006 59 | 5744 0.564597 0.224438 1.27546 60 | 5769 0.235186 -0.785427 1.57291 61 | 5923 0.602795 0.0748768 1.75291 62 | 5931 0.0749608 0.30583 1.17392 63 | 5945 0.0729823 0.220461 1.04365 64 | 5990 0.241094 -0.606848 1.51488 65 | 6034 0.445661 0.308446 1.27509 66 | 6035 -0.413788 0.136282 0.863115 67 | 6058 0.329165 0.259523 1.28037 68 | 6120 0.311015 0.333225 1.401 69 | 6131 0.0770971 -0.576734 1.17288 70 | 6156 0.407702 0.0554939 1.2107 71 | 6157 0.200939 0.358667 0.758642 72 | 6159 0.123987 -0.638569 1.3761 73 | 6163 0.272598 -0.766343 1.62085 74 | 6168 0.368746 0.145043 1.10864 75 | 6188 0.672657 0.256947 1.37587 76 | 6196 0.464341 0.015711 1.30712 77 | 6202 0.913592 0.185942 1.41637 78 | 6203 0.107669 0.215534 1.16471 79 | 6204 0.835018 0.307311 1.53895 80 | 6208 0.0761082 0.302123 1.25861 81 | 6220 0.711411 0.12417 1.49661 82 | 6223 0.653221 0.249799 1.33934 83 | 6226 0.933743 0.265751 1.62379 84 | 6238 1.01353 0.0498975 1.2623 85 | 6277 0.132598 -0.544035 1.1928 86 | 6279 1.01483 -0.0343 1.29504 87 | 6285 0.235481 -0.652239 1.52818 88 | 6297 1.21138 0.14139 1.45238 89 | 6301 1.0848 0.11587 1.28576 90 | 6342 -0.514787 0.0885156 1.1706 91 | 6373 0.905498 0.304281 1.07212 92 | 6402 0.829092 0.270511 0.977044 93 | 6429 0.167806 -0.587323 1.47889 94 | 6438 0.980302 0.104539 1.00818 95 | 6476 1.2025 0.168115 1.43668 96 | 6483 0.18918 -0.493305 1.14916 97 | -------------------------------------------------------------------------------- /demo_dataset/gravity/gravity_000000.txt: -------------------------------------------------------------------------------- 1 | 0.053781 2 | -0.99238 3 | 0.11086 4 | -------------------------------------------------------------------------------- /demo_dataset/gravity/gravity_000017.txt: -------------------------------------------------------------------------------- 1 | 0.0083714 2 | -0.99713 3 | 0.075256 4 | -------------------------------------------------------------------------------- /demo_dataset/gravity/gravity_000034.txt: -------------------------------------------------------------------------------- 1 | 0.0042035 2 | -0.99999 3 | -0.0017274 4 | -------------------------------------------------------------------------------- /demo_dataset/gravity/gravity_000051.txt: -------------------------------------------------------------------------------- 1 | -0.00040261 2 | -0.99997 3 | 0.0070634 4 | -------------------------------------------------------------------------------- /demo_dataset/gravity/gravity_000068.txt: -------------------------------------------------------------------------------- 1 | 0.047565 2 | -0.99797 3 | 0.042364 4 | -------------------------------------------------------------------------------- /demo_dataset/gravity/gravity_000085.txt: -------------------------------------------------------------------------------- 1 | 0.086047 2 | -0.97809 3 | 0.18956 4 | -------------------------------------------------------------------------------- /demo_dataset/gravity/gravity_000102.txt: -------------------------------------------------------------------------------- 1 | 0.041073 2 | -0.99616 3 | -0.077363 4 | -------------------------------------------------------------------------------- /demo_dataset/gravity/gravity_000119.txt: -------------------------------------------------------------------------------- 1 | -0.039726 2 | -0.9988 3 | 0.028591 4 | -------------------------------------------------------------------------------- /images/overview.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MARSLab-UMN/vi_depth_completion/b8d4bd3f3af674871a9fb8cd6166e49d785859ef/images/overview.jpg -------------------------------------------------------------------------------- /images/sample_results.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MARSLab-UMN/vi_depth_completion/b8d4bd3f3af674871a9fb8cd6166e49d785859ef/images/sample_results.jpg -------------------------------------------------------------------------------- /networks/__init__.py: -------------------------------------------------------------------------------- 1 | 2 | __all__ = ['depth_completion', 'network_utils'] 3 | -------------------------------------------------------------------------------- /networks/network_utils.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | 4 | 5 | def weights_init(modules, init_type='xavier'): 6 | assert init_type == 'xavier' or init_type == 'kaiming' 7 | m = modules 8 | if (isinstance(m, nn.Conv1d) or isinstance(m, nn.Conv2d) or isinstance(m, nn.Conv3d) or \ 9 | isinstance(m, nn.ConvTranspose1d) or isinstance(m, nn.ConvTranspose2d) or 10 | isinstance(m, nn.ConvTranspose3d) or isinstance(m, nn.Linear)): 11 | if init_type == 'xavier': 12 | torch.nn.init.xavier_normal_(m.weight) 13 | elif init_type == 'kaiming': 14 | torch.nn.init.kaiming_normal_(m.weight) 15 | 16 | if m.bias is not None: 17 | m.bias.data.zero_() 18 | elif isinstance(m, nn.BatchNorm2d): 19 | m.weight.data.fill_(1.0) 20 | m.bias.data.zero_() 21 | #elif isinstance(m, nn.Sequential) or isinstance(m, nn.ModuleList): 22 | #for m in modules: 23 | #weights_init(m, init_type) 24 | elif isinstance(m, nn.Module): 25 | for _, m in modules.named_children(): 26 | weights_init(m, init_type) 27 | -------------------------------------------------------------------------------- /normal_utils.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn 3 | import torch.nn.functional as F 4 | import numpy as np 5 | 6 | 7 | def compute_normal_vectors_loss_l2(norm_gt, pred_normals, mask): 8 | norm1 = pred_normals[:, 0:3, :, :] 9 | 10 | loss = -torch.sum(F.cosine_similarity(norm1, norm_gt, dim=1)) / torch.sum(mask) 11 | 12 | norm1 = Normalize(norm1) 13 | angle = torch.acos(torch.clamp(torch.sum(norm1 * norm_gt, dim=1), -1, 1)) / np.pi * 180 14 | angle = angle.view(mask.shape[0], 1, mask.shape[2], mask.shape[3]) * mask 15 | angle = torch.sum(angle) 16 | 17 | return loss, angle 18 | 19 | 20 | def compute_normal_vectors_loss_l1(norm_gt, pred_normals, mask, normalize_prediction=True): 21 | mask = mask.float() 22 | loss_func = torch.nn.L1Loss(reduction='sum') 23 | if normalize_prediction: 24 | norms = Normalize(pred_normals[:, 0:3, :, :]) 25 | else: 26 | norms = pred_normals[:, 0:3, :, :] 27 | 28 | angle = torch.acos(torch.clamp(torch.sum(norms * norm_gt, dim=1), -1, 1)) / np.pi * 180 29 | angle = angle.view(mask.shape[0], 1, mask.shape[2], mask.shape[3]) * mask 30 | angle = torch.sum(angle) 31 | 32 | num_elements = torch.sum(mask).item() 33 | loss = loss_func(norms * mask, norm_gt * mask) / num_elements 34 | return loss, angle 35 | 36 | -------------------------------------------------------------------------------- /plane_mask_detection/ABSTRACTIONS.md: -------------------------------------------------------------------------------- 1 | ## Abstractions 2 | The main abstractions introduced by `maskrcnn_benchmark` that are useful to 3 | have in mind are the following: 4 | 5 | ### ImageList 6 | In PyTorch, the first dimension of the input to the network generally represents 7 | the batch dimension, and thus all elements of the same batch have the same 8 | height / width. 9 | In order to support images with different sizes and aspect ratios in the same 10 | batch, we created the `ImageList` class, which holds internally a batch of 11 | images (os possibly different sizes). The images are padded with zeros such that 12 | they have the same final size and batched over the first dimension. The original 13 | sizes of the images before padding are stored in the `image_sizes` attribute, 14 | and the batched tensor in `tensors`. 15 | We provide a convenience function `to_image_list` that accepts a few different 16 | input types, including a list of tensors, and returns an `ImageList` object. 17 | 18 | ```python 19 | from maskrcnn_benchmark.structures.image_list import to_image_list 20 | 21 | images = [torch.rand(3, 100, 200), torch.rand(3, 150, 170)] 22 | batched_images = to_image_list(images) 23 | 24 | # it is also possible to make the final batched image be a multiple of a number 25 | batched_images_32 = to_image_list(images, size_divisible=32) 26 | ``` 27 | 28 | ### BoxList 29 | The `BoxList` class holds a set of bounding boxes (represented as a `Nx4` tensor) for 30 | a specific image, as well as the size of the image as a `(width, height)` tuple. 31 | It also contains a set of methods that allow to perform geometric 32 | transformations to the bounding boxes (such as cropping, scaling and flipping). 33 | The class accepts bounding boxes from two different input formats: 34 | - `xyxy`, where each box is encoded as a `x1`, `y1`, `x2` and `y2` coordinates, and 35 | - `xywh`, where each box is encoded as `x1`, `y1`, `w` and `h`. 36 | 37 | Additionally, each `BoxList` instance can also hold arbitrary additional information 38 | for each bounding box, such as labels, visibility, probability scores etc. 39 | 40 | Here is an example on how to create a `BoxList` from a list of coordinates: 41 | ```python 42 | from maskrcnn_benchmark.structures.bounding_box import BoxList, FLIP_LEFT_RIGHT 43 | 44 | width = 100 45 | height = 200 46 | boxes = [ 47 | [0, 10, 50, 50], 48 | [50, 20, 90, 60], 49 | [10, 10, 50, 50] 50 | ] 51 | # create a BoxList with 3 boxes 52 | bbox = BoxList(boxes, image_size=(width, height), mode='xyxy') 53 | 54 | # perform some box transformations, has similar API as PIL.Image 55 | bbox_scaled = bbox.resize((width * 2, height * 3)) 56 | bbox_flipped = bbox.transpose(FLIP_LEFT_RIGHT) 57 | 58 | # add labels for each bbox 59 | labels = torch.tensor([0, 10, 1]) 60 | bbox.add_field('labels', labels) 61 | 62 | # bbox also support a few operations, like indexing 63 | # here, selects boxes 0 and 2 64 | bbox_subset = bbox[[0, 2]] 65 | ``` 66 | -------------------------------------------------------------------------------- /plane_mask_detection/CODE_OF_CONDUCT.md: -------------------------------------------------------------------------------- 1 | # Code of Conduct 2 | 3 | Facebook has adopted a Code of Conduct that we expect project participants to adhere to. 4 | Please read the [full text](https://code.fb.com/codeofconduct/) 5 | so that you can understand what actions will and will not be tolerated. 6 | -------------------------------------------------------------------------------- /plane_mask_detection/CONTRIBUTING.md: -------------------------------------------------------------------------------- 1 | # Contributing to Mask-RCNN Benchmark 2 | We want to make contributing to this project as easy and transparent as 3 | possible. 4 | 5 | ## Our Development Process 6 | Minor changes and improvements will be released on an ongoing basis. Larger changes (e.g., changesets implementing a new paper) will be released on a more periodic basis. 7 | 8 | ## Pull Requests 9 | We actively welcome your pull requests. 10 | 11 | 1. Fork the repo and create your branch from `master`. 12 | 2. If you've added code that should be tested, add tests. 13 | 3. If you've changed APIs, update the documentation. 14 | 4. Ensure the test suite passes. 15 | 5. Make sure your code lints. 16 | 6. If you haven't already, complete the Contributor License Agreement ("CLA"). 17 | 18 | ## Contributor License Agreement ("CLA") 19 | In order to accept your pull request, we need you to submit a CLA. You only need 20 | to do this once to work on any of Facebook's open source projects. 21 | 22 | Complete your CLA here: 23 | 24 | ## Issues 25 | We use GitHub issues to track public bugs. Please ensure your description is 26 | clear and has sufficient instructions to be able to reproduce the issue. 27 | 28 | Facebook has a [bounty program](https://www.facebook.com/whitehat/) for the safe 29 | disclosure of security bugs. In those cases, please go through the process 30 | outlined on that page and do not file a public issue. 31 | 32 | ## Coding Style 33 | * 4 spaces for indentation rather than tabs 34 | * 80 character line length 35 | * PEP8 formatting following [Black](https://black.readthedocs.io/en/stable/) 36 | 37 | ## License 38 | By contributing to Mask-RCNN Benchmark, you agree that your contributions will be licensed 39 | under the LICENSE file in the root directory of this source tree. 40 | -------------------------------------------------------------------------------- /plane_mask_detection/configs/R101_bs16_all_plane_normal.yaml: -------------------------------------------------------------------------------- 1 | MODEL: 2 | META_ARCHITECTURE: "GeneralizedRCNN" 3 | WEIGHT: "catalog://ImageNetPretrained/MSRA/R-101" 4 | BACKBONE: 5 | CONV_BODY: "R-101-FPN" 6 | FREEZE_CONV_BODY_AT: 0 7 | RESNETS: 8 | BACKBONE_OUT_CHANNELS: 256 9 | RPN: 10 | USE_FPN: True 11 | ANCHOR_STRIDE: (4, 8, 16, 32, 64) 12 | PRE_NMS_TOP_N_TRAIN: 2000 13 | PRE_NMS_TOP_N_TEST: 1000 14 | POST_NMS_TOP_N_TEST: 50 15 | FPN_POST_NMS_TOP_N_TEST: 50 16 | ROI_HEADS: 17 | USE_FPN: True 18 | ROI_BOX_HEAD: 19 | POOLER_RESOLUTION: 7 20 | POOLER_SCALES: (0.25, 0.125, 0.0625, 0.03125) 21 | POOLER_SAMPLING_RATIO: 2 22 | FEATURE_EXTRACTOR: "FPN2MLPFeatureExtractor" 23 | PREDICTOR: "FPNPredictor" 24 | NUM_CLASSES: 2 #81 25 | ROI_MASK_HEAD: 26 | POOLER_SCALES: (0.25, 0.125, 0.0625, 0.03125) 27 | FEATURE_EXTRACTOR: "MaskRCNNFPNFeatureExtractor" 28 | PREDICTOR: "MaskRCNNC4Predictor" 29 | POOLER_RESOLUTION: 14 30 | POOLER_SAMPLING_RATIO: 2 31 | RESOLUTION: 28 32 | SHARE_BOX_FEATURE_EXTRACTOR: False 33 | MASK_ON: True 34 | DATASETS: 35 | TRAIN: ("coco_planemasks_all_train_split",) #("coco_framenet_train_200k_images", "coco_framenet_train_200k_images") ("coco_planenormal_0to1_train",) 36 | TEST: ("first10scenes_coco_plane_masks_test_split",) # ("coco_framenet_val_60k_images",) #("coco_framenet_0to1_test",) # 37 | # Fixed size image during training and testing 38 | INPUT: 39 | MIN_SIZE_TRAIN: (240,) 40 | MAX_SIZE_TRAIN: 320 41 | MIN_SIZE_TEST: 240 42 | MAX_SIZE_TEST: 320 43 | DATALOADER: 44 | SIZE_DIVISIBILITY: 32 45 | NUM_WORKERS: 16 46 | SOLVER: 47 | IMS_PER_BATCH: 16 48 | BASE_LR: 0.002 49 | WEIGHT_DECAY: 0.0001 50 | STEPS: (40000, 60000) 51 | MAX_ITER: 200000 52 | UNSUPERVISED: 53 | PLANE_NORMAL: 0 54 | SUPERVISED: "plane_normal" # "normal_only" 55 | OUTPUT_REPRESENTATION: 5 56 | # 0: estimate directly 3 values of normal vector w/o any particular constraints, follows by normalization, then L1 loss 57 | # 1: estimate 2dof [tt, phi] only, then L1 loss 58 | # 2: estimate 3 value of normal, applying tanh in the last layer (scale up to -1 .. 1), no normalization needed, then MSE loss 59 | # 3: estimate 3 value of normal, applying hard-tanh in the last layer (scale up to -1 .. 1), no normalization needed, then MSE loss 60 | # 4: estimate 2dof, followed by a periodic characteristic, then L1 loss 61 | TEST: 62 | IMS_PER_BATCH: 16 63 | OUTPUT_DIR: ./checkpoints 64 | # '/tiendo/mnt/dgx/PanopticFPN_results/R101_bs64_L2loss_3output_unnormalized/' 65 | #'/tiendo/mnt/dgx/PanopticFPN_results/R101_bs128_Normal_640x480/' 66 | -------------------------------------------------------------------------------- /plane_mask_detection/demo/README.md: -------------------------------------------------------------------------------- 1 | ## Webcam and Jupyter notebook demo 2 | 3 | This folder contains a simple webcam demo that illustrates how you can use `maskrcnn_benchmark` for inference. 4 | 5 | 6 | ### With your preferred environment 7 | 8 | You can start it by running it from this folder, using one of the following commands: 9 | ```bash 10 | # by default, it runs on the GPU 11 | # for best results, use min-image-size 800 12 | python webcam.py --min-image-size 800 13 | # can also run it on the CPU 14 | python webcam.py --min-image-size 300 MODEL.DEVICE cpu 15 | # or change the model that you want to use 16 | python webcam.py --config-file ../configs/caffe2/e2e_mask_rcnn_R_101_FPN_1x_caffe2.yaml --min-image-size 300 MODEL.DEVICE cpu 17 | # in order to see the probability heatmaps, pass --show-mask-heatmaps 18 | python webcam.py --min-image-size 300 --show-mask-heatmaps MODEL.DEVICE cpu 19 | ``` 20 | 21 | ### With Docker 22 | 23 | Build the image with the tag `maskrcnn-benchmark` (check [INSTALL.md](../INSTALL.md) for instructions) 24 | 25 | Adjust permissions of the X server host (be careful with this step, refer to 26 | [here](http://wiki.ros.org/docker/Tutorials/GUI) for alternatives) 27 | 28 | ```bash 29 | xhost + 30 | ``` 31 | 32 | Then run a container with the demo: 33 | 34 | ``` 35 | docker run --rm -it \ 36 | -e DISPLAY=${DISPLAY} \ 37 | --privileged \ 38 | -v /tmp/.X11-unix:/tmp/.X11-unix \ 39 | --device=/dev/video0:/dev/video0 \ 40 | --ipc=host maskrcnn-benchmark \ 41 | python demo/webcam.py --min-image-size 300 \ 42 | --config-file configs/caffe2/e2e_mask_rcnn_R_50_FPN_1x_caffe2.yaml 43 | ``` 44 | 45 | **DISCLAIMER:** *This was tested for an Ubuntu 16.04 machine, 46 | the volume mapping may vary depending on your platform* 47 | -------------------------------------------------------------------------------- /plane_mask_detection/demo/demo_e2e_mask_rcnn_R_50_FPN_1x.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MARSLab-UMN/vi_depth_completion/b8d4bd3f3af674871a9fb8cd6166e49d785859ef/plane_mask_detection/demo/demo_e2e_mask_rcnn_R_50_FPN_1x.png -------------------------------------------------------------------------------- /plane_mask_detection/demo/demo_e2e_mask_rcnn_X_101_32x8d_FPN_1x.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MARSLab-UMN/vi_depth_completion/b8d4bd3f3af674871a9fb8cd6166e49d785859ef/plane_mask_detection/demo/demo_e2e_mask_rcnn_X_101_32x8d_FPN_1x.png -------------------------------------------------------------------------------- /plane_mask_detection/demo/webcam.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. 2 | import argparse 3 | import cv2 4 | 5 | from maskrcnn_benchmark.config import cfg 6 | from predictor import COCODemo 7 | 8 | import time 9 | 10 | 11 | def main(): 12 | parser = argparse.ArgumentParser(description="PyTorch Object Detection Webcam Demo") 13 | parser.add_argument( 14 | "--config-file", 15 | default="../configs/caffe2/e2e_mask_rcnn_R_50_FPN_1x_caffe2.yaml", 16 | metavar="FILE", 17 | help="path to config file", 18 | ) 19 | parser.add_argument( 20 | "--confidence-threshold", 21 | type=float, 22 | default=0.7, 23 | help="Minimum score for the prediction to be shown", 24 | ) 25 | parser.add_argument( 26 | "--min-image-size", 27 | type=int, 28 | default=224, 29 | help="Smallest size of the image to feed to the model. " 30 | "Model was trained with 800, which gives best results", 31 | ) 32 | parser.add_argument( 33 | "--show-mask-heatmaps", 34 | dest="show_mask_heatmaps", 35 | help="Show a heatmap probability for the top masks-per-dim masks", 36 | action="store_true", 37 | ) 38 | parser.add_argument( 39 | "--masks-per-dim", 40 | type=int, 41 | default=2, 42 | help="Number of heatmaps per dimension to show", 43 | ) 44 | parser.add_argument( 45 | "opts", 46 | help="Modify model config options using the command-line", 47 | default=None, 48 | nargs=argparse.REMAINDER, 49 | ) 50 | 51 | args = parser.parse_args() 52 | 53 | # load config from file and command-line arguments 54 | cfg.merge_from_file(args.config_file) 55 | cfg.merge_from_list(args.opts) 56 | cfg.freeze() 57 | 58 | # prepare object that handles inference plus adds predictions on top of image 59 | coco_demo = COCODemo( 60 | cfg, 61 | confidence_threshold=args.confidence_threshold, 62 | show_mask_heatmaps=args.show_mask_heatmaps, 63 | masks_per_dim=args.masks_per_dim, 64 | min_image_size=args.min_image_size, 65 | ) 66 | 67 | cam = cv2.VideoCapture(0) 68 | while True: 69 | start_time = time.time() 70 | ret_val, img = cam.read() 71 | composite = coco_demo.run_on_opencv_image(img) 72 | print("Time: {:.2f} s / img".format(time.time() - start_time)) 73 | cv2.imshow("COCO detections", composite) 74 | if cv2.waitKey(1) == 27: 75 | break # esc to quit 76 | cv2.destroyAllWindows() 77 | 78 | 79 | if __name__ == "__main__": 80 | main() 81 | -------------------------------------------------------------------------------- /plane_mask_detection/maskrcnn_benchmark/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. 2 | -------------------------------------------------------------------------------- /plane_mask_detection/maskrcnn_benchmark/config/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. 2 | from .defaults import _C as cfg 3 | -------------------------------------------------------------------------------- /plane_mask_detection/maskrcnn_benchmark/csrc/ROIAlign.h: -------------------------------------------------------------------------------- 1 | // Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. 2 | #pragma once 3 | 4 | #include "cpu/vision.h" 5 | 6 | #ifdef WITH_CUDA 7 | #include "cuda/vision.h" 8 | #endif 9 | 10 | // Interface for Python 11 | at::Tensor ROIAlign_forward(const at::Tensor& input, 12 | const at::Tensor& rois, 13 | const float spatial_scale, 14 | const int pooled_height, 15 | const int pooled_width, 16 | const int sampling_ratio) { 17 | if (input.type().is_cuda()) { 18 | #ifdef WITH_CUDA 19 | return ROIAlign_forward_cuda(input, rois, spatial_scale, pooled_height, pooled_width, sampling_ratio); 20 | #else 21 | AT_ERROR("Not compiled with GPU support"); 22 | #endif 23 | } 24 | return ROIAlign_forward_cpu(input, rois, spatial_scale, pooled_height, pooled_width, sampling_ratio); 25 | } 26 | 27 | at::Tensor ROIAlign_backward(const at::Tensor& grad, 28 | const at::Tensor& rois, 29 | const float spatial_scale, 30 | const int pooled_height, 31 | const int pooled_width, 32 | const int batch_size, 33 | const int channels, 34 | const int height, 35 | const int width, 36 | const int sampling_ratio) { 37 | if (grad.type().is_cuda()) { 38 | #ifdef WITH_CUDA 39 | return ROIAlign_backward_cuda(grad, rois, spatial_scale, pooled_height, pooled_width, batch_size, channels, height, width, sampling_ratio); 40 | #else 41 | AT_ERROR("Not compiled with GPU support"); 42 | #endif 43 | } 44 | AT_ERROR("Not implemented on the CPU"); 45 | } 46 | 47 | -------------------------------------------------------------------------------- /plane_mask_detection/maskrcnn_benchmark/csrc/ROIPool.h: -------------------------------------------------------------------------------- 1 | // Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. 2 | #pragma once 3 | 4 | #include "cpu/vision.h" 5 | 6 | #ifdef WITH_CUDA 7 | #include "cuda/vision.h" 8 | #endif 9 | 10 | 11 | std::tuple ROIPool_forward(const at::Tensor& input, 12 | const at::Tensor& rois, 13 | const float spatial_scale, 14 | const int pooled_height, 15 | const int pooled_width) { 16 | if (input.type().is_cuda()) { 17 | #ifdef WITH_CUDA 18 | return ROIPool_forward_cuda(input, rois, spatial_scale, pooled_height, pooled_width); 19 | #else 20 | AT_ERROR("Not compiled with GPU support"); 21 | #endif 22 | } 23 | AT_ERROR("Not implemented on the CPU"); 24 | } 25 | 26 | at::Tensor ROIPool_backward(const at::Tensor& grad, 27 | const at::Tensor& input, 28 | const at::Tensor& rois, 29 | const at::Tensor& argmax, 30 | const float spatial_scale, 31 | const int pooled_height, 32 | const int pooled_width, 33 | const int batch_size, 34 | const int channels, 35 | const int height, 36 | const int width) { 37 | if (grad.type().is_cuda()) { 38 | #ifdef WITH_CUDA 39 | return ROIPool_backward_cuda(grad, input, rois, argmax, spatial_scale, pooled_height, pooled_width, batch_size, channels, height, width); 40 | #else 41 | AT_ERROR("Not compiled with GPU support"); 42 | #endif 43 | } 44 | AT_ERROR("Not implemented on the CPU"); 45 | } 46 | 47 | 48 | 49 | -------------------------------------------------------------------------------- /plane_mask_detection/maskrcnn_benchmark/csrc/SigmoidFocalLoss.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include "cpu/vision.h" 4 | 5 | #ifdef WITH_CUDA 6 | #include "cuda/vision.h" 7 | #endif 8 | 9 | // Interface for Python 10 | at::Tensor SigmoidFocalLoss_forward( 11 | const at::Tensor& logits, 12 | const at::Tensor& targets, 13 | const int num_classes, 14 | const float gamma, 15 | const float alpha) { 16 | if (logits.type().is_cuda()) { 17 | #ifdef WITH_CUDA 18 | return SigmoidFocalLoss_forward_cuda(logits, targets, num_classes, gamma, alpha); 19 | #else 20 | AT_ERROR("Not compiled with GPU support"); 21 | #endif 22 | } 23 | AT_ERROR("Not implemented on the CPU"); 24 | } 25 | 26 | at::Tensor SigmoidFocalLoss_backward( 27 | const at::Tensor& logits, 28 | const at::Tensor& targets, 29 | const at::Tensor& d_losses, 30 | const int num_classes, 31 | const float gamma, 32 | const float alpha) { 33 | if (logits.type().is_cuda()) { 34 | #ifdef WITH_CUDA 35 | return SigmoidFocalLoss_backward_cuda(logits, targets, d_losses, num_classes, gamma, alpha); 36 | #else 37 | AT_ERROR("Not compiled with GPU support"); 38 | #endif 39 | } 40 | AT_ERROR("Not implemented on the CPU"); 41 | } 42 | -------------------------------------------------------------------------------- /plane_mask_detection/maskrcnn_benchmark/csrc/cpu/nms_cpu.cpp: -------------------------------------------------------------------------------- 1 | // Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. 2 | #include "cpu/vision.h" 3 | 4 | 5 | template 6 | at::Tensor nms_cpu_kernel(const at::Tensor& dets, 7 | const at::Tensor& scores, 8 | const float threshold) { 9 | AT_ASSERTM(!dets.type().is_cuda(), "dets must be a CPU tensor"); 10 | AT_ASSERTM(!scores.type().is_cuda(), "scores must be a CPU tensor"); 11 | AT_ASSERTM(dets.type() == scores.type(), "dets should have the same type as scores"); 12 | 13 | if (dets.numel() == 0) { 14 | return at::empty({0}, dets.options().dtype(at::kLong).device(at::kCPU)); 15 | } 16 | 17 | auto x1_t = dets.select(1, 0).contiguous(); 18 | auto y1_t = dets.select(1, 1).contiguous(); 19 | auto x2_t = dets.select(1, 2).contiguous(); 20 | auto y2_t = dets.select(1, 3).contiguous(); 21 | 22 | at::Tensor areas_t = (x2_t - x1_t + 1) * (y2_t - y1_t + 1); 23 | 24 | auto order_t = std::get<1>(scores.sort(0, /* descending=*/true)); 25 | 26 | auto ndets = dets.size(0); 27 | at::Tensor suppressed_t = at::zeros({ndets}, dets.options().dtype(at::kByte).device(at::kCPU)); 28 | 29 | auto suppressed = suppressed_t.data(); 30 | auto order = order_t.data(); 31 | auto x1 = x1_t.data(); 32 | auto y1 = y1_t.data(); 33 | auto x2 = x2_t.data(); 34 | auto y2 = y2_t.data(); 35 | auto areas = areas_t.data(); 36 | 37 | for (int64_t _i = 0; _i < ndets; _i++) { 38 | auto i = order[_i]; 39 | if (suppressed[i] == 1) 40 | continue; 41 | auto ix1 = x1[i]; 42 | auto iy1 = y1[i]; 43 | auto ix2 = x2[i]; 44 | auto iy2 = y2[i]; 45 | auto iarea = areas[i]; 46 | 47 | for (int64_t _j = _i + 1; _j < ndets; _j++) { 48 | auto j = order[_j]; 49 | if (suppressed[j] == 1) 50 | continue; 51 | auto xx1 = std::max(ix1, x1[j]); 52 | auto yy1 = std::max(iy1, y1[j]); 53 | auto xx2 = std::min(ix2, x2[j]); 54 | auto yy2 = std::min(iy2, y2[j]); 55 | 56 | auto w = std::max(static_cast(0), xx2 - xx1 + 1); 57 | auto h = std::max(static_cast(0), yy2 - yy1 + 1); 58 | auto inter = w * h; 59 | auto ovr = inter / (iarea + areas[j] - inter); 60 | if (ovr >= threshold) 61 | suppressed[j] = 1; 62 | } 63 | } 64 | return at::nonzero(suppressed_t == 0).squeeze(1); 65 | } 66 | 67 | at::Tensor nms_cpu(const at::Tensor& dets, 68 | const at::Tensor& scores, 69 | const float threshold) { 70 | at::Tensor result; 71 | AT_DISPATCH_FLOATING_TYPES(dets.type(), "nms", [&] { 72 | result = nms_cpu_kernel(dets, scores, threshold); 73 | }); 74 | return result; 75 | } 76 | -------------------------------------------------------------------------------- /plane_mask_detection/maskrcnn_benchmark/csrc/cpu/vision.h: -------------------------------------------------------------------------------- 1 | // Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. 2 | #pragma once 3 | #include 4 | 5 | 6 | at::Tensor ROIAlign_forward_cpu(const at::Tensor& input, 7 | const at::Tensor& rois, 8 | const float spatial_scale, 9 | const int pooled_height, 10 | const int pooled_width, 11 | const int sampling_ratio); 12 | 13 | 14 | at::Tensor nms_cpu(const at::Tensor& dets, 15 | const at::Tensor& scores, 16 | const float threshold); 17 | -------------------------------------------------------------------------------- /plane_mask_detection/maskrcnn_benchmark/csrc/cuda/deform_pool_cuda.cu: -------------------------------------------------------------------------------- 1 | // modify from 2 | // https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/blob/mmdetection/mmdet/ops/dcn/src/modulated_dcn_cuda.c 3 | 4 | // based on 5 | // author: Charles Shang 6 | // https://github.com/torch/cunn/blob/master/lib/THCUNN/generic/SpatialConvolutionMM.cu 7 | 8 | #include 9 | #include 10 | 11 | #include 12 | #include 13 | 14 | #include 15 | #include 16 | #include 17 | 18 | 19 | void DeformablePSROIPoolForward( 20 | const at::Tensor data, const at::Tensor bbox, const at::Tensor trans, 21 | at::Tensor out, at::Tensor top_count, const int batch, const int channels, 22 | const int height, const int width, const int num_bbox, 23 | const int channels_trans, const int no_trans, const float spatial_scale, 24 | const int output_dim, const int group_size, const int pooled_size, 25 | const int part_size, const int sample_per_part, const float trans_std); 26 | 27 | void DeformablePSROIPoolBackwardAcc( 28 | const at::Tensor out_grad, const at::Tensor data, const at::Tensor bbox, 29 | const at::Tensor trans, const at::Tensor top_count, at::Tensor in_grad, 30 | at::Tensor trans_grad, const int batch, const int channels, 31 | const int height, const int width, const int num_bbox, 32 | const int channels_trans, const int no_trans, const float spatial_scale, 33 | const int output_dim, const int group_size, const int pooled_size, 34 | const int part_size, const int sample_per_part, const float trans_std); 35 | 36 | void deform_psroi_pooling_cuda_forward( 37 | at::Tensor input, at::Tensor bbox, at::Tensor trans, at::Tensor out, 38 | at::Tensor top_count, const int no_trans, const float spatial_scale, 39 | const int output_dim, const int group_size, const int pooled_size, 40 | const int part_size, const int sample_per_part, const float trans_std) 41 | { 42 | AT_CHECK(input.is_contiguous(), "input tensor has to be contiguous"); 43 | 44 | const int batch = input.size(0); 45 | const int channels = input.size(1); 46 | const int height = input.size(2); 47 | const int width = input.size(3); 48 | const int channels_trans = no_trans ? 2 : trans.size(1); 49 | 50 | const int num_bbox = bbox.size(0); 51 | if (num_bbox != out.size(0)) 52 | AT_ERROR("Output shape and bbox number wont match: (%d vs %d).", 53 | out.size(0), num_bbox); 54 | 55 | DeformablePSROIPoolForward( 56 | input, bbox, trans, out, top_count, batch, channels, height, width, 57 | num_bbox, channels_trans, no_trans, spatial_scale, output_dim, group_size, 58 | pooled_size, part_size, sample_per_part, trans_std); 59 | } 60 | 61 | void deform_psroi_pooling_cuda_backward( 62 | at::Tensor out_grad, at::Tensor input, at::Tensor bbox, at::Tensor trans, 63 | at::Tensor top_count, at::Tensor input_grad, at::Tensor trans_grad, 64 | const int no_trans, const float spatial_scale, const int output_dim, 65 | const int group_size, const int pooled_size, const int part_size, 66 | const int sample_per_part, const float trans_std) 67 | { 68 | AT_CHECK(out_grad.is_contiguous(), "out_grad tensor has to be contiguous"); 69 | AT_CHECK(input.is_contiguous(), "input tensor has to be contiguous"); 70 | 71 | const int batch = input.size(0); 72 | const int channels = input.size(1); 73 | const int height = input.size(2); 74 | const int width = input.size(3); 75 | const int channels_trans = no_trans ? 2 : trans.size(1); 76 | 77 | const int num_bbox = bbox.size(0); 78 | if (num_bbox != out_grad.size(0)) 79 | AT_ERROR("Output shape and bbox number wont match: (%d vs %d).", 80 | out_grad.size(0), num_bbox); 81 | 82 | DeformablePSROIPoolBackwardAcc( 83 | out_grad, input, bbox, trans, top_count, input_grad, trans_grad, batch, 84 | channels, height, width, num_bbox, channels_trans, no_trans, 85 | spatial_scale, output_dim, group_size, pooled_size, part_size, 86 | sample_per_part, trans_std); 87 | } 88 | -------------------------------------------------------------------------------- /plane_mask_detection/maskrcnn_benchmark/csrc/deform_pool.h: -------------------------------------------------------------------------------- 1 | // Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. 2 | #pragma once 3 | #include "cpu/vision.h" 4 | 5 | #ifdef WITH_CUDA 6 | #include "cuda/vision.h" 7 | #endif 8 | 9 | 10 | // Interface for Python 11 | void deform_psroi_pooling_forward( 12 | at::Tensor input, 13 | at::Tensor bbox, 14 | at::Tensor trans, 15 | at::Tensor out, 16 | at::Tensor top_count, 17 | const int no_trans, 18 | const float spatial_scale, 19 | const int output_dim, 20 | const int group_size, 21 | const int pooled_size, 22 | const int part_size, 23 | const int sample_per_part, 24 | const float trans_std) 25 | { 26 | if (input.type().is_cuda()) { 27 | #ifdef WITH_CUDA 28 | return deform_psroi_pooling_cuda_forward( 29 | input, bbox, trans, out, top_count, 30 | no_trans, spatial_scale, output_dim, group_size, 31 | pooled_size, part_size, sample_per_part, trans_std 32 | ); 33 | #else 34 | AT_ERROR("Not compiled with GPU support"); 35 | #endif 36 | } 37 | AT_ERROR("Not implemented on the CPU"); 38 | } 39 | 40 | 41 | void deform_psroi_pooling_backward( 42 | at::Tensor out_grad, 43 | at::Tensor input, 44 | at::Tensor bbox, 45 | at::Tensor trans, 46 | at::Tensor top_count, 47 | at::Tensor input_grad, 48 | at::Tensor trans_grad, 49 | const int no_trans, 50 | const float spatial_scale, 51 | const int output_dim, 52 | const int group_size, 53 | const int pooled_size, 54 | const int part_size, 55 | const int sample_per_part, 56 | const float trans_std) 57 | { 58 | if (input.type().is_cuda()) { 59 | #ifdef WITH_CUDA 60 | return deform_psroi_pooling_cuda_backward( 61 | out_grad, input, bbox, trans, top_count, input_grad, trans_grad, 62 | no_trans, spatial_scale, output_dim, group_size, pooled_size, 63 | part_size, sample_per_part, trans_std 64 | ); 65 | #else 66 | AT_ERROR("Not compiled with GPU support"); 67 | #endif 68 | } 69 | AT_ERROR("Not implemented on the CPU"); 70 | } 71 | -------------------------------------------------------------------------------- /plane_mask_detection/maskrcnn_benchmark/csrc/nms.h: -------------------------------------------------------------------------------- 1 | // Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. 2 | #pragma once 3 | #include "cpu/vision.h" 4 | 5 | #ifdef WITH_CUDA 6 | #include "cuda/vision.h" 7 | #endif 8 | 9 | 10 | at::Tensor nms(const at::Tensor& dets, 11 | const at::Tensor& scores, 12 | const float threshold) { 13 | 14 | if (dets.type().is_cuda()) { 15 | #ifdef WITH_CUDA 16 | // TODO raise error if not compiled with CUDA 17 | if (dets.numel() == 0) 18 | return at::empty({0}, dets.options().dtype(at::kLong).device(at::kCPU)); 19 | auto b = at::cat({dets, scores.unsqueeze(1)}, 1); 20 | return nms_cuda(b, threshold); 21 | #else 22 | AT_ERROR("Not compiled with GPU support"); 23 | #endif 24 | } 25 | 26 | at::Tensor result = nms_cpu(dets, scores, threshold); 27 | return result; 28 | } 29 | -------------------------------------------------------------------------------- /plane_mask_detection/maskrcnn_benchmark/csrc/vision.cpp: -------------------------------------------------------------------------------- 1 | // Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. 2 | #include "nms.h" 3 | #include "ROIAlign.h" 4 | #include "ROIPool.h" 5 | #include "SigmoidFocalLoss.h" 6 | #include "deform_conv.h" 7 | #include "deform_pool.h" 8 | 9 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 10 | m.def("nms", &nms, "non-maximum suppression"); 11 | m.def("roi_align_forward", &ROIAlign_forward, "ROIAlign_forward"); 12 | m.def("roi_align_backward", &ROIAlign_backward, "ROIAlign_backward"); 13 | m.def("roi_pool_forward", &ROIPool_forward, "ROIPool_forward"); 14 | m.def("roi_pool_backward", &ROIPool_backward, "ROIPool_backward"); 15 | m.def("sigmoid_focalloss_forward", &SigmoidFocalLoss_forward, "SigmoidFocalLoss_forward"); 16 | m.def("sigmoid_focalloss_backward", &SigmoidFocalLoss_backward, "SigmoidFocalLoss_backward"); 17 | // dcn-v2 18 | m.def("deform_conv_forward", &deform_conv_forward, "deform_conv_forward"); 19 | m.def("deform_conv_backward_input", &deform_conv_backward_input, "deform_conv_backward_input"); 20 | m.def("deform_conv_backward_parameters", &deform_conv_backward_parameters, "deform_conv_backward_parameters"); 21 | m.def("modulated_deform_conv_forward", &modulated_deform_conv_forward, "modulated_deform_conv_forward"); 22 | m.def("modulated_deform_conv_backward", &modulated_deform_conv_backward, "modulated_deform_conv_backward"); 23 | m.def("deform_psroi_pooling_forward", &deform_psroi_pooling_forward, "deform_psroi_pooling_forward"); 24 | m.def("deform_psroi_pooling_backward", &deform_psroi_pooling_backward, "deform_psroi_pooling_backward"); 25 | } -------------------------------------------------------------------------------- /plane_mask_detection/maskrcnn_benchmark/data/README.md: -------------------------------------------------------------------------------- 1 | # Setting Up Datasets 2 | This file describes how to perform training on other datasets. 3 | 4 | Only Pascal VOC dataset can be loaded from its original format and be outputted to Pascal style results currently. 5 | 6 | We expect the annotations from other datasets be converted to COCO json format, and 7 | the output will be in COCO-style. (i.e. AP, AP50, AP75, APs, APm, APl for bbox and segm) 8 | 9 | ## Creating Symlinks for PASCAL VOC 10 | 11 | We assume that your symlinked `datasets/voc/VOC` directory has the following structure: 12 | 13 | ``` 14 | VOC 15 | |_ JPEGImages 16 | | |_ .jpg 17 | | |_ ... 18 | | |_ .jpg 19 | |_ Annotations 20 | | |_ pascal_train.json (optional) 21 | | |_ pascal_val.json (optional) 22 | | |_ pascal_test.json (optional) 23 | | |_ .xml 24 | | |_ ... 25 | | |_ .xml 26 | |_ VOCdevkit 27 | ``` 28 | 29 | Create symlinks for `voc/VOC`: 30 | 31 | ``` 32 | cd ~/github/maskrcnn-benchmark 33 | mkdir -p datasets/voc/VOC 34 | ln -s /path/to/VOC /datasets/voc/VOC 35 | ``` 36 | Example configuration files for PASCAL VOC could be found [here](https://github.com/facebookresearch/maskrcnn-benchmark/blob/master/configs/pascal_voc/). 37 | 38 | ### PASCAL VOC Annotations in COCO Format 39 | To output COCO-style evaluation result, PASCAL VOC annotations in COCO json format is required and could be downloaded from [here](https://storage.googleapis.com/coco-dataset/external/PASCAL_VOC.zip) 40 | via http://cocodataset.org/#external. 41 | 42 | ## Creating Symlinks for Cityscapes: 43 | 44 | We assume that your symlinked `datasets/cityscapes` directory has the following structure: 45 | 46 | ``` 47 | cityscapes 48 | |_ images 49 | | |_ .jpg 50 | | |_ ... 51 | | |_ .jpg 52 | |_ annotations 53 | | |_ instanceonly_gtFile_train.json 54 | | |_ ... 55 | |_ raw 56 | |_ gtFine 57 | |_ ... 58 | |_ README.md 59 | ``` 60 | 61 | Create symlinks for `cityscapes`: 62 | 63 | ``` 64 | cd ~/github/maskrcnn-benchmark 65 | mkdir -p datasets/cityscapes 66 | ln -s /path/to/cityscapes datasets/data/cityscapes 67 | ``` 68 | 69 | ### Steps to convert Cityscapes Annotations to COCO Format 70 | 1. Download gtFine_trainvaltest.zip from https://www.cityscapes-dataset.com/downloads/ (login required) 71 | 2. Extract it to /path/to/gtFine_trainvaltest 72 | ``` 73 | cityscapes 74 | |_ gtFine_trainvaltest.zip 75 | |_ gtFine_trainvaltest 76 | |_ gtFine 77 | ``` 78 | 3. Run the below commands to convert the annotations 79 | 80 | ``` 81 | cd ~/github 82 | git clone https://github.com/mcordts/cityscapesScripts.git 83 | cd cityscapesScripts 84 | cp ~/github/maskrcnn-benchmark/tools/cityscapes/instances2dict_with_polygons.py cityscapesscripts/evaluation 85 | python setup.py install 86 | cd ~/github/maskrcnn-benchmark 87 | python tools/cityscapes/convert_cityscapes_to_coco.py --datadir /path/to/cityscapes --outdir /path/to/cityscapes/annotations 88 | ``` 89 | 90 | Example configuration files for Cityscapes could be found [here](https://github.com/facebookresearch/maskrcnn-benchmark/blob/master/configs/cityscapes/). 91 | -------------------------------------------------------------------------------- /plane_mask_detection/maskrcnn_benchmark/data/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. 2 | from .build import make_data_loader 3 | -------------------------------------------------------------------------------- /plane_mask_detection/maskrcnn_benchmark/data/collate_batch.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. 2 | from pytorch_local.maskrcnn_benchmark.structures.image_list import to_image_list 3 | from pytorch_local.maskrcnn_benchmark.structures.image_list import to_gravity_tensor 4 | 5 | 6 | class BatchCollator(object): 7 | """ 8 | From a list of samples from the dataset, 9 | returns the batched images and targets. 10 | This should be passed to the DataLoader 11 | """ 12 | 13 | def __init__(self, size_divisible=0): 14 | self.size_divisible = size_divisible 15 | 16 | def __call__(self, batch): 17 | transposed_batch = list(zip(*batch)) 18 | images = to_image_list(transposed_batch[0], self.size_divisible) 19 | targets = transposed_batch[1] 20 | normal_images = to_image_list(transposed_batch[2], self.size_divisible) 21 | mask_images = to_image_list(transposed_batch[3], self.size_divisible) 22 | img_ids = transposed_batch[4] 23 | gravity_dirs = None 24 | gravity_mask = None 25 | if len(transposed_batch) >= 6: 26 | gravity_dirs = to_gravity_tensor(transposed_batch[5]) 27 | gravity_mask = to_image_list(transposed_batch[6], self.size_divisible) 28 | return images, targets, normal_images, mask_images, img_ids, gravity_dirs, gravity_mask 29 | 30 | 31 | class BBoxAugCollator(object): 32 | """ 33 | From a list of samples from the dataset, 34 | returns the images and targets. 35 | Images should be converted to batched images in `im_detect_bbox_aug` 36 | """ 37 | 38 | def __call__(self, batch): 39 | return list(zip(*batch)) 40 | 41 | -------------------------------------------------------------------------------- /plane_mask_detection/maskrcnn_benchmark/data/datasets/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. 2 | from .coco import COCODataset 3 | from .voc import PascalVOCDataset 4 | from .concat_dataset import ConcatDataset 5 | from .fmd import FastMotionDataset 6 | 7 | __all__ = ["COCODataset", "ConcatDataset", "PascalVOCDataset", "FastMotionDataset"] 8 | -------------------------------------------------------------------------------- /plane_mask_detection/maskrcnn_benchmark/data/datasets/concat_dataset.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. 2 | import bisect 3 | 4 | from torch.utils.data.dataset import ConcatDataset as _ConcatDataset 5 | 6 | 7 | class ConcatDataset(_ConcatDataset): 8 | """ 9 | Same as torch.utils.data.dataset.ConcatDataset, but exposes an extra 10 | method for querying the sizes of the image 11 | """ 12 | 13 | def get_idxs(self, idx): 14 | dataset_idx = bisect.bisect_right(self.cumulative_sizes, idx) 15 | if dataset_idx == 0: 16 | sample_idx = idx 17 | else: 18 | sample_idx = idx - self.cumulative_sizes[dataset_idx - 1] 19 | return dataset_idx, sample_idx 20 | 21 | def get_img_info(self, idx): 22 | dataset_idx, sample_idx = self.get_idxs(idx) 23 | return self.datasets[dataset_idx].get_img_info(sample_idx) 24 | -------------------------------------------------------------------------------- /plane_mask_detection/maskrcnn_benchmark/data/datasets/evaluation/coco/__init__.py: -------------------------------------------------------------------------------- 1 | from .coco_eval import do_coco_evaluation 2 | 3 | 4 | def coco_evaluation( 5 | dataset, 6 | predictions, 7 | output_folder, 8 | box_only, 9 | iou_types, 10 | expected_results, 11 | expected_results_sigma_tol, 12 | ): 13 | return do_coco_evaluation( 14 | dataset=dataset, 15 | predictions=predictions, 16 | box_only=box_only, 17 | output_folder=output_folder, 18 | iou_types=iou_types, 19 | expected_results=expected_results, 20 | expected_results_sigma_tol=expected_results_sigma_tol, 21 | ) 22 | -------------------------------------------------------------------------------- /plane_mask_detection/maskrcnn_benchmark/data/datasets/evaluation/voc/__init__.py: -------------------------------------------------------------------------------- 1 | import logging 2 | 3 | from .voc_eval import do_voc_evaluation 4 | 5 | 6 | def voc_evaluation(dataset, predictions, output_folder, box_only, **_): 7 | logger = logging.getLogger("maskrcnn_benchmark.inference") 8 | if box_only: 9 | logger.warning("voc evaluation doesn't support box_only, ignored.") 10 | logger.info("performing voc evaluation, ignored iou_types.") 11 | return do_voc_evaluation( 12 | dataset=dataset, 13 | predictions=predictions, 14 | output_folder=output_folder, 15 | logger=logger, 16 | ) 17 | -------------------------------------------------------------------------------- /plane_mask_detection/maskrcnn_benchmark/data/datasets/list_dataset.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. 2 | """ 3 | Simple dataset class that wraps a list of path names 4 | """ 5 | 6 | from PIL import Image 7 | 8 | from maskrcnn_benchmark.structures.bounding_box import BoxList 9 | 10 | 11 | class ListDataset(object): 12 | def __init__(self, image_lists, transforms=None): 13 | self.image_lists = image_lists 14 | self.transforms = transforms 15 | 16 | def __getitem__(self, item): 17 | img = Image.open(self.image_lists[item]).convert("RGB") 18 | 19 | # dummy target 20 | w, h = img.size 21 | target = BoxList([[0, 0, w, h]], img.size, mode="xyxy") 22 | 23 | if self.transforms is not None: 24 | img, target = self.transforms(img, target) 25 | 26 | return img, target 27 | 28 | def __len__(self): 29 | return len(self.image_lists) 30 | 31 | def get_img_info(self, item): 32 | """ 33 | Return the image dimensions for the image, without 34 | loading and pre-processing it 35 | """ 36 | pass 37 | -------------------------------------------------------------------------------- /plane_mask_detection/maskrcnn_benchmark/data/datasets/voc.py: -------------------------------------------------------------------------------- 1 | import os 2 | 3 | import torch 4 | import torch.utils.data 5 | from PIL import Image 6 | import sys 7 | 8 | if sys.version_info[0] == 2: 9 | import xml.etree.cElementTree as ET 10 | else: 11 | import xml.etree.ElementTree as ET 12 | 13 | 14 | from pytorch_local.maskrcnn_benchmark.structures.bounding_box import BoxList 15 | 16 | 17 | class PascalVOCDataset(torch.utils.data.Dataset): 18 | 19 | CLASSES = ( 20 | "__background__ ", 21 | "aeroplane", 22 | "bicycle", 23 | "bird", 24 | "boat", 25 | "bottle", 26 | "bus", 27 | "car", 28 | "cat", 29 | "chair", 30 | "cow", 31 | "diningtable", 32 | "dog", 33 | "horse", 34 | "motorbike", 35 | "person", 36 | "pottedplant", 37 | "sheep", 38 | "sofa", 39 | "train", 40 | "tvmonitor", 41 | ) 42 | 43 | def __init__(self, data_dir, split, use_difficult=False, transforms=None): 44 | self.root = data_dir 45 | self.image_set = split 46 | self.keep_difficult = use_difficult 47 | self.transforms = transforms 48 | 49 | self._annopath = os.path.join(self.root, "Annotations", "%s.xml") 50 | self._imgpath = os.path.join(self.root, "JPEGImages", "%s.jpg") 51 | self._imgsetpath = os.path.join(self.root, "ImageSets", "Main", "%s.txt") 52 | 53 | with open(self._imgsetpath % self.image_set) as f: 54 | self.ids = f.readlines() 55 | self.ids = [x.strip("\n") for x in self.ids] 56 | self.id_to_img_map = {k: v for k, v in enumerate(self.ids)} 57 | 58 | cls = PascalVOCDataset.CLASSES 59 | self.class_to_ind = dict(zip(cls, range(len(cls)))) 60 | self.categories = dict(zip(range(len(cls)), cls)) 61 | 62 | def __getitem__(self, index): 63 | img_id = self.ids[index] 64 | img = Image.open(self._imgpath % img_id).convert("RGB") 65 | 66 | target = self.get_groundtruth(index) 67 | target = target.clip_to_image(remove_empty=True) 68 | 69 | if self.transforms is not None: 70 | img, target = self.transforms(img, target) 71 | 72 | return img, target, index 73 | 74 | def __len__(self): 75 | return len(self.ids) 76 | 77 | def get_groundtruth(self, index): 78 | img_id = self.ids[index] 79 | anno = ET.parse(self._annopath % img_id).getroot() 80 | anno = self._preprocess_annotation(anno) 81 | 82 | height, width = anno["im_info"] 83 | target = BoxList(anno["boxes"], (width, height), mode="xyxy") 84 | target.add_field("labels", anno["labels"]) 85 | target.add_field("difficult", anno["difficult"]) 86 | return target 87 | 88 | def _preprocess_annotation(self, target): 89 | boxes = [] 90 | gt_classes = [] 91 | difficult_boxes = [] 92 | TO_REMOVE = 1 93 | 94 | for obj in target.iter("object"): 95 | difficult = int(obj.find("difficult").text) == 1 96 | if not self.keep_difficult and difficult: 97 | continue 98 | name = obj.find("name").text.lower().strip() 99 | bb = obj.find("bndbox") 100 | # Make pixel indexes 0-based 101 | # Refer to "https://github.com/rbgirshick/py-faster-rcnn/blob/master/lib/datasets/pascal_voc.py#L208-L211" 102 | box = [ 103 | bb.find("xmin").text, 104 | bb.find("ymin").text, 105 | bb.find("xmax").text, 106 | bb.find("ymax").text, 107 | ] 108 | bndbox = tuple( 109 | map(lambda x: x - TO_REMOVE, list(map(int, box))) 110 | ) 111 | 112 | boxes.append(bndbox) 113 | gt_classes.append(self.class_to_ind[name]) 114 | difficult_boxes.append(difficult) 115 | 116 | size = target.find("size") 117 | im_info = tuple(map(int, (size.find("height").text, size.find("width").text))) 118 | 119 | res = { 120 | "boxes": torch.tensor(boxes, dtype=torch.float32), 121 | "labels": torch.tensor(gt_classes), 122 | "difficult": torch.tensor(difficult_boxes), 123 | "im_info": im_info, 124 | } 125 | return res 126 | 127 | def get_img_info(self, index): 128 | img_id = self.ids[index] 129 | anno = ET.parse(self._annopath % img_id).getroot() 130 | size = anno.find("size") 131 | im_info = tuple(map(int, (size.find("height").text, size.find("width").text))) 132 | return {"height": im_info[0], "width": im_info[1]} 133 | 134 | def map_class_id_to_class_name(self, class_id): 135 | return PascalVOCDataset.CLASSES[class_id] 136 | -------------------------------------------------------------------------------- /plane_mask_detection/maskrcnn_benchmark/data/samplers/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. 2 | from .distributed import DistributedSampler 3 | from .grouped_batch_sampler import GroupedBatchSampler 4 | from .iteration_based_batch_sampler import IterationBasedBatchSampler 5 | 6 | __all__ = ["DistributedSampler", "GroupedBatchSampler", "IterationBasedBatchSampler"] 7 | -------------------------------------------------------------------------------- /plane_mask_detection/maskrcnn_benchmark/data/samplers/distributed.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. 2 | # Code is copy-pasted exactly as in torch.utils.data.distributed. 3 | # FIXME remove this once c10d fixes the bug it has 4 | import math 5 | import torch 6 | import torch.distributed as dist 7 | from torch.utils.data.sampler import Sampler 8 | 9 | 10 | class DistributedSampler(Sampler): 11 | """Sampler that restricts data loading to a subset of the dataset. 12 | It is especially useful in conjunction with 13 | :class:`torch.nn.parallel.DistributedDataParallel`. In such case, each 14 | process can pass a DistributedSampler instance as a DataLoader sampler, 15 | and load a subset of the original dataset that is exclusive to it. 16 | .. note:: 17 | Dataset is assumed to be of constant size. 18 | Arguments: 19 | dataset: Dataset used for sampling. 20 | num_replicas (optional): Number of processes participating in 21 | distributed training. 22 | rank (optional): Rank of the current process within num_replicas. 23 | """ 24 | 25 | def __init__(self, dataset, num_replicas=None, rank=None, shuffle=True): 26 | if num_replicas is None: 27 | if not dist.is_available(): 28 | raise RuntimeError("Requires distributed package to be available") 29 | num_replicas = dist.get_world_size() 30 | if rank is None: 31 | if not dist.is_available(): 32 | raise RuntimeError("Requires distributed package to be available") 33 | rank = dist.get_rank() 34 | self.dataset = dataset 35 | self.num_replicas = num_replicas 36 | self.rank = rank 37 | self.epoch = 0 38 | self.num_samples = int(math.ceil(len(self.dataset) * 1.0 / self.num_replicas)) 39 | self.total_size = self.num_samples * self.num_replicas 40 | self.shuffle = shuffle 41 | 42 | def __iter__(self): 43 | if self.shuffle: 44 | # deterministically shuffle based on epoch 45 | g = torch.Generator() 46 | g.manual_seed(self.epoch) 47 | indices = torch.randperm(len(self.dataset), generator=g).tolist() 48 | else: 49 | indices = torch.arange(len(self.dataset)).tolist() 50 | 51 | # add extra samples to make it evenly divisible 52 | indices += indices[: (self.total_size - len(indices))] 53 | assert len(indices) == self.total_size 54 | 55 | # subsample 56 | offset = self.num_samples * self.rank 57 | indices = indices[offset : offset + self.num_samples] 58 | assert len(indices) == self.num_samples 59 | 60 | return iter(indices) 61 | 62 | def __len__(self): 63 | return self.num_samples 64 | 65 | def set_epoch(self, epoch): 66 | self.epoch = epoch 67 | -------------------------------------------------------------------------------- /plane_mask_detection/maskrcnn_benchmark/data/samplers/iteration_based_batch_sampler.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. 2 | from torch.utils.data.sampler import BatchSampler 3 | 4 | 5 | class IterationBasedBatchSampler(BatchSampler): 6 | """ 7 | Wraps a BatchSampler, resampling from it until 8 | a specified number of iterations have been sampled 9 | """ 10 | 11 | def __init__(self, batch_sampler, num_iterations, start_iter=0): 12 | self.batch_sampler = batch_sampler 13 | self.num_iterations = num_iterations 14 | self.start_iter = start_iter 15 | 16 | def __iter__(self): 17 | iteration = self.start_iter 18 | while iteration <= self.num_iterations: 19 | # if the underlying sampler has a set_epoch method, like 20 | # DistributedSampler, used for making each process see 21 | # a different split of the dataset, then set it 22 | if hasattr(self.batch_sampler.sampler, "set_epoch"): 23 | self.batch_sampler.sampler.set_epoch(iteration) 24 | for batch in self.batch_sampler: 25 | iteration += 1 26 | if iteration > self.num_iterations: 27 | break 28 | yield batch 29 | 30 | def __len__(self): 31 | return self.num_iterations 32 | -------------------------------------------------------------------------------- /plane_mask_detection/maskrcnn_benchmark/data/transforms/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. 2 | from .transforms import Compose 3 | from .transforms import Resize 4 | from .transforms import RandomHorizontalFlip 5 | from .transforms import ToTensor 6 | from .transforms import Normalize 7 | 8 | from .build import build_transforms 9 | -------------------------------------------------------------------------------- /plane_mask_detection/maskrcnn_benchmark/data/transforms/build.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. 2 | from . import transforms as T 3 | 4 | 5 | def build_transforms(cfg, is_train=True): 6 | if is_train: 7 | min_size = cfg.INPUT.MIN_SIZE_TRAIN 8 | max_size = cfg.INPUT.MAX_SIZE_TRAIN 9 | flip_horizontal_prob = 0.5 # cfg.INPUT.FLIP_PROB_TRAIN 10 | flip_vertical_prob = cfg.INPUT.VERTICAL_FLIP_PROB_TRAIN 11 | brightness = cfg.INPUT.BRIGHTNESS 12 | contrast = cfg.INPUT.CONTRAST 13 | saturation = cfg.INPUT.SATURATION 14 | hue = cfg.INPUT.HUE 15 | else: 16 | min_size = cfg.INPUT.MIN_SIZE_TEST 17 | max_size = cfg.INPUT.MAX_SIZE_TEST 18 | flip_horizontal_prob = 0.0 19 | flip_vertical_prob = 0.0 20 | brightness = 0.0 21 | contrast = 0.0 22 | saturation = 0.0 23 | hue = 0.0 24 | 25 | to_bgr255 = cfg.INPUT.TO_BGR255 26 | normalize_transform = T.Normalize( 27 | mean=cfg.INPUT.PIXEL_MEAN, std=cfg.INPUT.PIXEL_STD, to_bgr255=to_bgr255 28 | ) 29 | color_jitter = T.ColorJitter( 30 | brightness=brightness, 31 | contrast=contrast, 32 | saturation=saturation, 33 | hue=hue, 34 | ) 35 | 36 | transform = T.Compose( 37 | [ 38 | color_jitter, 39 | T.Resize(min_size, max_size), 40 | T.RandomHorizontalFlip(flip_horizontal_prob), 41 | T.RandomVerticalFlip(flip_vertical_prob), 42 | T.ToTensor(), 43 | normalize_transform, 44 | ] 45 | ) 46 | return transform 47 | -------------------------------------------------------------------------------- /plane_mask_detection/maskrcnn_benchmark/engine/.inference.py.swp: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MARSLab-UMN/vi_depth_completion/b8d4bd3f3af674871a9fb8cd6166e49d785859ef/plane_mask_detection/maskrcnn_benchmark/engine/.inference.py.swp -------------------------------------------------------------------------------- /plane_mask_detection/maskrcnn_benchmark/engine/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. 2 | -------------------------------------------------------------------------------- /plane_mask_detection/maskrcnn_benchmark/layers/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. 2 | import torch 3 | 4 | from .batch_norm import FrozenBatchNorm2d 5 | from .misc import Conv2d 6 | from .misc import DFConv2d 7 | from .misc import ConvTranspose2d 8 | from .misc import BatchNorm2d 9 | from .misc import interpolate 10 | from .nms import nms 11 | from .roi_align import ROIAlign 12 | from .roi_align import roi_align 13 | from .roi_pool import ROIPool 14 | from .roi_pool import roi_pool 15 | from .smooth_l1_loss import smooth_l1_loss 16 | from .sigmoid_focal_loss import SigmoidFocalLoss 17 | from .dcn.deform_conv_func import deform_conv, modulated_deform_conv 18 | from .dcn.deform_conv_module import DeformConv, ModulatedDeformConv, ModulatedDeformConvPack 19 | from .dcn.deform_pool_func import deform_roi_pooling 20 | from .dcn.deform_pool_module import DeformRoIPooling, DeformRoIPoolingPack, ModulatedDeformRoIPoolingPack 21 | 22 | 23 | __all__ = [ 24 | "nms", 25 | "roi_align", 26 | "ROIAlign", 27 | "roi_pool", 28 | "ROIPool", 29 | "smooth_l1_loss", 30 | "Conv2d", 31 | "DFConv2d", 32 | "ConvTranspose2d", 33 | "interpolate", 34 | "BatchNorm2d", 35 | "FrozenBatchNorm2d", 36 | "SigmoidFocalLoss", 37 | 'deform_conv', 38 | 'modulated_deform_conv', 39 | 'DeformConv', 40 | 'ModulatedDeformConv', 41 | 'ModulatedDeformConvPack', 42 | 'deform_roi_pooling', 43 | 'DeformRoIPooling', 44 | 'DeformRoIPoolingPack', 45 | 'ModulatedDeformRoIPoolingPack', 46 | ] 47 | 48 | -------------------------------------------------------------------------------- /plane_mask_detection/maskrcnn_benchmark/layers/_utils.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. 2 | import glob 3 | import os.path 4 | 5 | import torch 6 | 7 | try: 8 | from torch.utils.cpp_extension import load as load_ext 9 | from torch.utils.cpp_extension import CUDA_HOME 10 | except ImportError: 11 | raise ImportError("The cpp layer extensions requires PyTorch 0.4 or higher") 12 | 13 | 14 | def _load_C_extensions(): 15 | this_dir = os.path.dirname(os.path.abspath(__file__)) 16 | this_dir = os.path.dirname(this_dir) 17 | this_dir = os.path.join(this_dir, "csrc") 18 | 19 | main_file = glob.glob(os.path.join(this_dir, "*.cpp")) 20 | source_cpu = glob.glob(os.path.join(this_dir, "cpu", "*.cpp")) 21 | source_cuda = glob.glob(os.path.join(this_dir, "cuda", "*.cu")) 22 | 23 | source = main_file + source_cpu 24 | 25 | extra_cflags = [] 26 | if torch.cuda.is_available() and CUDA_HOME is not None: 27 | source.extend(source_cuda) 28 | extra_cflags = ["-DWITH_CUDA"] 29 | source = [os.path.join(this_dir, s) for s in source] 30 | extra_include_paths = [this_dir] 31 | return load_ext( 32 | "torchvision", 33 | source, 34 | extra_cflags=extra_cflags, 35 | extra_include_paths=extra_include_paths, 36 | ) 37 | 38 | 39 | _C = _load_C_extensions() 40 | -------------------------------------------------------------------------------- /plane_mask_detection/maskrcnn_benchmark/layers/batch_norm.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. 2 | import torch 3 | from torch import nn 4 | 5 | 6 | class FrozenBatchNorm2d(nn.Module): 7 | """ 8 | BatchNorm2d where the batch statistics and the affine parameters 9 | are fixed 10 | """ 11 | 12 | def __init__(self, n): 13 | super(FrozenBatchNorm2d, self).__init__() 14 | self.register_buffer("weight", torch.ones(n)) 15 | self.register_buffer("bias", torch.zeros(n)) 16 | self.register_buffer("running_mean", torch.zeros(n)) 17 | self.register_buffer("running_var", torch.ones(n)) 18 | 19 | def forward(self, x): 20 | # Cast all fixed parameters to half() if necessary 21 | if x.dtype == torch.float16: 22 | self.weight = self.weight.half() 23 | self.bias = self.bias.half() 24 | self.running_mean = self.running_mean.half() 25 | self.running_var = self.running_var.half() 26 | 27 | scale = self.weight * self.running_var.rsqrt() 28 | bias = self.bias - self.running_mean * scale 29 | scale = scale.reshape(1, -1, 1, 1) 30 | bias = bias.reshape(1, -1, 1, 1) 31 | return x * scale + bias 32 | -------------------------------------------------------------------------------- /plane_mask_detection/maskrcnn_benchmark/layers/dcn/__init__.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copied From [mmdetection](https://github.com/open-mmlab/mmdetection/tree/master/mmdet/ops/dcn) 3 | # 4 | -------------------------------------------------------------------------------- /plane_mask_detection/maskrcnn_benchmark/layers/dcn/deform_pool_func.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch.autograd import Function 3 | from torch.autograd.function import once_differentiable 4 | 5 | from maskrcnn_benchmark import _C 6 | 7 | 8 | class DeformRoIPoolingFunction(Function): 9 | 10 | @staticmethod 11 | def forward( 12 | ctx, 13 | data, 14 | rois, 15 | offset, 16 | spatial_scale, 17 | out_size, 18 | out_channels, 19 | no_trans, 20 | group_size=1, 21 | part_size=None, 22 | sample_per_part=4, 23 | trans_std=.0 24 | ): 25 | ctx.spatial_scale = spatial_scale 26 | ctx.out_size = out_size 27 | ctx.out_channels = out_channels 28 | ctx.no_trans = no_trans 29 | ctx.group_size = group_size 30 | ctx.part_size = out_size if part_size is None else part_size 31 | ctx.sample_per_part = sample_per_part 32 | ctx.trans_std = trans_std 33 | 34 | assert 0.0 <= ctx.trans_std <= 1.0 35 | if not data.is_cuda: 36 | raise NotImplementedError 37 | 38 | n = rois.shape[0] 39 | output = data.new_empty(n, out_channels, out_size, out_size) 40 | output_count = data.new_empty(n, out_channels, out_size, out_size) 41 | _C.deform_psroi_pooling_forward( 42 | data, 43 | rois, 44 | offset, 45 | output, 46 | output_count, 47 | ctx.no_trans, 48 | ctx.spatial_scale, 49 | ctx.out_channels, 50 | ctx.group_size, 51 | ctx.out_size, 52 | ctx.part_size, 53 | ctx.sample_per_part, 54 | ctx.trans_std 55 | ) 56 | 57 | if data.requires_grad or rois.requires_grad or offset.requires_grad: 58 | ctx.save_for_backward(data, rois, offset) 59 | ctx.output_count = output_count 60 | 61 | return output 62 | 63 | @staticmethod 64 | @once_differentiable 65 | def backward(ctx, grad_output): 66 | if not grad_output.is_cuda: 67 | raise NotImplementedError 68 | 69 | data, rois, offset = ctx.saved_tensors 70 | output_count = ctx.output_count 71 | grad_input = torch.zeros_like(data) 72 | grad_rois = None 73 | grad_offset = torch.zeros_like(offset) 74 | 75 | _C.deform_psroi_pooling_backward( 76 | grad_output, 77 | data, 78 | rois, 79 | offset, 80 | output_count, 81 | grad_input, 82 | grad_offset, 83 | ctx.no_trans, 84 | ctx.spatial_scale, 85 | ctx.out_channels, 86 | ctx.group_size, 87 | ctx.out_size, 88 | ctx.part_size, 89 | ctx.sample_per_part, 90 | ctx.trans_std 91 | ) 92 | return (grad_input, grad_rois, grad_offset, None, None, None, None, None, None, None, None) 93 | 94 | 95 | deform_roi_pooling = DeformRoIPoolingFunction.apply 96 | -------------------------------------------------------------------------------- /plane_mask_detection/maskrcnn_benchmark/layers/nms.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. 2 | # from ._utils import _C 3 | from maskrcnn_benchmark import _C 4 | 5 | from apex import amp 6 | 7 | # Only valid with fp32 inputs - give AMP the hint 8 | nms = amp.float_function(_C.nms) 9 | 10 | # nms.__doc__ = """ 11 | # This function performs Non-maximum suppresion""" 12 | -------------------------------------------------------------------------------- /plane_mask_detection/maskrcnn_benchmark/layers/roi_align.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. 2 | import torch 3 | from torch import nn 4 | from torch.autograd import Function 5 | from torch.autograd.function import once_differentiable 6 | from torch.nn.modules.utils import _pair 7 | 8 | from maskrcnn_benchmark import _C 9 | 10 | from apex import amp 11 | 12 | class _ROIAlign(Function): 13 | @staticmethod 14 | def forward(ctx, input, roi, output_size, spatial_scale, sampling_ratio): 15 | ctx.save_for_backward(roi) 16 | ctx.output_size = _pair(output_size) 17 | ctx.spatial_scale = spatial_scale 18 | ctx.sampling_ratio = sampling_ratio 19 | ctx.input_shape = input.size() 20 | output = _C.roi_align_forward( 21 | input, roi, spatial_scale, output_size[0], output_size[1], sampling_ratio 22 | ) 23 | return output 24 | 25 | @staticmethod 26 | @once_differentiable 27 | def backward(ctx, grad_output): 28 | rois, = ctx.saved_tensors 29 | output_size = ctx.output_size 30 | spatial_scale = ctx.spatial_scale 31 | sampling_ratio = ctx.sampling_ratio 32 | bs, ch, h, w = ctx.input_shape 33 | grad_input = _C.roi_align_backward( 34 | grad_output, 35 | rois, 36 | spatial_scale, 37 | output_size[0], 38 | output_size[1], 39 | bs, 40 | ch, 41 | h, 42 | w, 43 | sampling_ratio, 44 | ) 45 | return grad_input, None, None, None, None 46 | 47 | 48 | roi_align = _ROIAlign.apply 49 | 50 | class ROIAlign(nn.Module): 51 | def __init__(self, output_size, spatial_scale, sampling_ratio): 52 | super(ROIAlign, self).__init__() 53 | self.output_size = output_size 54 | self.spatial_scale = spatial_scale 55 | self.sampling_ratio = sampling_ratio 56 | 57 | @amp.float_function 58 | def forward(self, input, rois): 59 | return roi_align( 60 | input, rois, self.output_size, self.spatial_scale, self.sampling_ratio 61 | ) 62 | 63 | def __repr__(self): 64 | tmpstr = self.__class__.__name__ + "(" 65 | tmpstr += "output_size=" + str(self.output_size) 66 | tmpstr += ", spatial_scale=" + str(self.spatial_scale) 67 | tmpstr += ", sampling_ratio=" + str(self.sampling_ratio) 68 | tmpstr += ")" 69 | return tmpstr 70 | -------------------------------------------------------------------------------- /plane_mask_detection/maskrcnn_benchmark/layers/roi_pool.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. 2 | import torch 3 | from torch import nn 4 | from torch.autograd import Function 5 | from torch.autograd.function import once_differentiable 6 | from torch.nn.modules.utils import _pair 7 | 8 | from maskrcnn_benchmark import _C 9 | 10 | from apex import amp 11 | 12 | class _ROIPool(Function): 13 | @staticmethod 14 | def forward(ctx, input, roi, output_size, spatial_scale): 15 | ctx.output_size = _pair(output_size) 16 | ctx.spatial_scale = spatial_scale 17 | ctx.input_shape = input.size() 18 | output, argmax = _C.roi_pool_forward( 19 | input, roi, spatial_scale, output_size[0], output_size[1] 20 | ) 21 | ctx.save_for_backward(input, roi, argmax) 22 | return output 23 | 24 | @staticmethod 25 | @once_differentiable 26 | def backward(ctx, grad_output): 27 | input, rois, argmax = ctx.saved_tensors 28 | output_size = ctx.output_size 29 | spatial_scale = ctx.spatial_scale 30 | bs, ch, h, w = ctx.input_shape 31 | grad_input = _C.roi_pool_backward( 32 | grad_output, 33 | input, 34 | rois, 35 | argmax, 36 | spatial_scale, 37 | output_size[0], 38 | output_size[1], 39 | bs, 40 | ch, 41 | h, 42 | w, 43 | ) 44 | return grad_input, None, None, None 45 | 46 | 47 | roi_pool = _ROIPool.apply 48 | 49 | 50 | class ROIPool(nn.Module): 51 | def __init__(self, output_size, spatial_scale): 52 | super(ROIPool, self).__init__() 53 | self.output_size = output_size 54 | self.spatial_scale = spatial_scale 55 | 56 | @amp.float_function 57 | def forward(self, input, rois): 58 | return roi_pool(input, rois, self.output_size, self.spatial_scale) 59 | 60 | def __repr__(self): 61 | tmpstr = self.__class__.__name__ + "(" 62 | tmpstr += "output_size=" + str(self.output_size) 63 | tmpstr += ", spatial_scale=" + str(self.spatial_scale) 64 | tmpstr += ")" 65 | return tmpstr 66 | -------------------------------------------------------------------------------- /plane_mask_detection/maskrcnn_benchmark/layers/sigmoid_focal_loss.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch import nn 3 | from torch.autograd import Function 4 | from torch.autograd.function import once_differentiable 5 | 6 | from maskrcnn_benchmark import _C 7 | 8 | # TODO: Use JIT to replace CUDA implementation in the future. 9 | class _SigmoidFocalLoss(Function): 10 | @staticmethod 11 | def forward(ctx, logits, targets, gamma, alpha): 12 | ctx.save_for_backward(logits, targets) 13 | num_classes = logits.shape[1] 14 | ctx.num_classes = num_classes 15 | ctx.gamma = gamma 16 | ctx.alpha = alpha 17 | 18 | losses = _C.sigmoid_focalloss_forward( 19 | logits, targets, num_classes, gamma, alpha 20 | ) 21 | return losses 22 | 23 | @staticmethod 24 | @once_differentiable 25 | def backward(ctx, d_loss): 26 | logits, targets = ctx.saved_tensors 27 | num_classes = ctx.num_classes 28 | gamma = ctx.gamma 29 | alpha = ctx.alpha 30 | d_loss = d_loss.contiguous() 31 | d_logits = _C.sigmoid_focalloss_backward( 32 | logits, targets, d_loss, num_classes, gamma, alpha 33 | ) 34 | return d_logits, None, None, None, None 35 | 36 | 37 | sigmoid_focal_loss_cuda = _SigmoidFocalLoss.apply 38 | 39 | 40 | def sigmoid_focal_loss_cpu(logits, targets, gamma, alpha): 41 | num_classes = logits.shape[1] 42 | gamma = gamma[0] 43 | alpha = alpha[0] 44 | dtype = targets.dtype 45 | device = targets.device 46 | class_range = torch.arange(1, num_classes+1, dtype=dtype, device=device).unsqueeze(0) 47 | 48 | t = targets.unsqueeze(1) 49 | p = torch.sigmoid(logits) 50 | term1 = (1 - p) ** gamma * torch.log(p) 51 | term2 = p ** gamma * torch.log(1 - p) 52 | return -(t == class_range).float() * term1 * alpha - ((t != class_range) * (t >= 0)).float() * term2 * (1 - alpha) 53 | 54 | 55 | class SigmoidFocalLoss(nn.Module): 56 | def __init__(self, gamma, alpha): 57 | super(SigmoidFocalLoss, self).__init__() 58 | self.gamma = gamma 59 | self.alpha = alpha 60 | 61 | def forward(self, logits, targets): 62 | device = logits.device 63 | if logits.is_cuda: 64 | loss_func = sigmoid_focal_loss_cuda 65 | else: 66 | loss_func = sigmoid_focal_loss_cpu 67 | 68 | loss = loss_func(logits, targets, self.gamma, self.alpha) 69 | return loss.sum() 70 | 71 | def __repr__(self): 72 | tmpstr = self.__class__.__name__ + "(" 73 | tmpstr += "gamma=" + str(self.gamma) 74 | tmpstr += ", alpha=" + str(self.alpha) 75 | tmpstr += ")" 76 | return tmpstr 77 | -------------------------------------------------------------------------------- /plane_mask_detection/maskrcnn_benchmark/layers/smooth_l1_loss.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. 2 | import torch 3 | 4 | 5 | # TODO maybe push this to nn? 6 | def smooth_l1_loss(input, target, beta=1. / 9, size_average=True): 7 | """ 8 | very similar to the smooth_l1_loss from pytorch, but with 9 | the extra beta parameter 10 | """ 11 | n = torch.abs(input - target) 12 | cond = n < beta 13 | loss = torch.where(cond, 0.5 * n ** 2 / beta, n - 0.5 * beta) 14 | if size_average: 15 | return loss.mean() 16 | return loss.sum() 17 | -------------------------------------------------------------------------------- /plane_mask_detection/maskrcnn_benchmark/modeling/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MARSLab-UMN/vi_depth_completion/b8d4bd3f3af674871a9fb8cd6166e49d785859ef/plane_mask_detection/maskrcnn_benchmark/modeling/__init__.py -------------------------------------------------------------------------------- /plane_mask_detection/maskrcnn_benchmark/modeling/backbone/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. 2 | from .backbone import build_backbone 3 | from . import fbnet 4 | -------------------------------------------------------------------------------- /plane_mask_detection/maskrcnn_benchmark/modeling/backbone/backbone.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. 2 | from collections import OrderedDict 3 | 4 | from torch import nn 5 | 6 | from plane_mask_detection.maskrcnn_benchmark.modeling import registry 7 | from plane_mask_detection.maskrcnn_benchmark.modeling.make_layers import conv_with_kaiming_uniform 8 | from . import fpn as fpn_module 9 | from . import resnet 10 | 11 | 12 | @registry.BACKBONES.register("R-50-C4") 13 | @registry.BACKBONES.register("R-50-C5") 14 | @registry.BACKBONES.register("R-101-C4") 15 | @registry.BACKBONES.register("R-101-C5") 16 | def build_resnet_backbone(cfg): 17 | body = resnet.ResNet(cfg) 18 | model = nn.Sequential(OrderedDict([("body", body)])) 19 | model.out_channels = cfg.MODEL.RESNETS.BACKBONE_OUT_CHANNELS 20 | return model 21 | 22 | 23 | @registry.BACKBONES.register("R-50-FPN") 24 | @registry.BACKBONES.register("R-101-FPN") 25 | @registry.BACKBONES.register("R-152-FPN") 26 | def build_resnet_fpn_backbone(cfg): 27 | body = resnet.ResNet(cfg) 28 | in_channels_stage2 = cfg.MODEL.RESNETS.RES2_OUT_CHANNELS 29 | out_channels = cfg.MODEL.RESNETS.BACKBONE_OUT_CHANNELS 30 | fpn = fpn_module.FPN( 31 | in_channels_list=[ 32 | in_channels_stage2, 33 | in_channels_stage2 * 2, 34 | in_channels_stage2 * 4, 35 | in_channels_stage2 * 8, 36 | ], 37 | out_channels=out_channels, 38 | conv_block=conv_with_kaiming_uniform( 39 | cfg.MODEL.FPN.USE_GN, cfg.MODEL.FPN.USE_RELU 40 | ), 41 | top_blocks=fpn_module.LastLevelMaxPool(), 42 | ) 43 | model = nn.Sequential(OrderedDict([("body", body), ("fpn", fpn)])) 44 | model.out_channels = out_channels 45 | return model 46 | 47 | 48 | @registry.BACKBONES.register("R-50-FPN-RETINANET") 49 | @registry.BACKBONES.register("R-101-FPN-RETINANET") 50 | def build_resnet_fpn_p3p7_backbone(cfg): 51 | body = resnet.ResNet(cfg) 52 | in_channels_stage2 = cfg.MODEL.RESNETS.RES2_OUT_CHANNELS 53 | out_channels = cfg.MODEL.RESNETS.BACKBONE_OUT_CHANNELS 54 | in_channels_p6p7 = in_channels_stage2 * 8 if cfg.MODEL.RETINANET.USE_C5 \ 55 | else out_channels 56 | fpn = fpn_module.FPN( 57 | in_channels_list=[ 58 | 0, 59 | in_channels_stage2 * 2, 60 | in_channels_stage2 * 4, 61 | in_channels_stage2 * 8, 62 | ], 63 | out_channels=out_channels, 64 | conv_block=conv_with_kaiming_uniform( 65 | cfg.MODEL.FPN.USE_GN, cfg.MODEL.FPN.USE_RELU 66 | ), 67 | top_blocks=fpn_module.LastLevelP6P7(in_channels_p6p7, out_channels), 68 | ) 69 | model = nn.Sequential(OrderedDict([("body", body), ("fpn", fpn)])) 70 | model.out_channels = out_channels 71 | return model 72 | 73 | 74 | def build_backbone(cfg): 75 | assert cfg.MODEL.BACKBONE.CONV_BODY in registry.BACKBONES, \ 76 | "cfg.MODEL.BACKBONE.CONV_BODY: {} are not registered in registry".format( 77 | cfg.MODEL.BACKBONE.CONV_BODY 78 | ) 79 | return registry.BACKBONES[cfg.MODEL.BACKBONE.CONV_BODY](cfg) 80 | -------------------------------------------------------------------------------- /plane_mask_detection/maskrcnn_benchmark/modeling/backbone/fpn.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. 2 | import torch 3 | import torch.nn.functional as F 4 | from torch import nn 5 | 6 | 7 | class FPN(nn.Module): 8 | """ 9 | Module that adds FPN on top of a list of feature maps. 10 | The feature maps are currently supposed to be in increasing depth 11 | order, and must be consecutive 12 | """ 13 | 14 | def __init__( 15 | self, in_channels_list, out_channels, conv_block, top_blocks=None 16 | ): 17 | """ 18 | Arguments: 19 | in_channels_list (list[int]): number of channels for each feature map that 20 | will be fed 21 | out_channels (int): number of channels of the FPN representation 22 | top_blocks (nn.Module or None): if provided, an extra operation will 23 | be performed on the output of the last (smallest resolution) 24 | FPN output, and the result will extend the result list 25 | """ 26 | super(FPN, self).__init__() 27 | self.inner_blocks = [] 28 | self.layer_blocks = [] 29 | for idx, in_channels in enumerate(in_channels_list, 1): 30 | inner_block = "fpn_inner{}".format(idx) 31 | layer_block = "fpn_layer{}".format(idx) 32 | 33 | if in_channels == 0: 34 | continue 35 | inner_block_module = conv_block(in_channels, out_channels, 1) 36 | layer_block_module = conv_block(out_channels, out_channels, 3, 1) 37 | self.add_module(inner_block, inner_block_module) 38 | self.add_module(layer_block, layer_block_module) 39 | self.inner_blocks.append(inner_block) 40 | self.layer_blocks.append(layer_block) 41 | self.top_blocks = top_blocks 42 | 43 | def forward(self, x): 44 | """ 45 | Arguments: 46 | x (list[Tensor]): feature maps for each feature level. 47 | Returns: 48 | results (tuple[Tensor]): feature maps after FPN layers. 49 | They are ordered from highest resolution first. 50 | """ 51 | last_inner = getattr(self, self.inner_blocks[-1])(x[-1]) 52 | results = [] 53 | results.append(getattr(self, self.layer_blocks[-1])(last_inner)) 54 | for feature, inner_block, layer_block in zip( 55 | x[:-1][::-1], self.inner_blocks[:-1][::-1], self.layer_blocks[:-1][::-1] 56 | ): 57 | if not inner_block: 58 | continue 59 | inner_top_down = F.interpolate(last_inner, scale_factor=2, mode="nearest") 60 | inner_lateral = getattr(self, inner_block)(feature) 61 | # TODO use size instead of scale to make it robust to different sizes 62 | # inner_top_down = F.upsample(last_inner, size=inner_lateral.shape[-2:], 63 | # mode='bilinear', align_corners=False) 64 | last_inner = inner_lateral + inner_top_down 65 | results.insert(0, getattr(self, layer_block)(last_inner)) 66 | 67 | if isinstance(self.top_blocks, LastLevelP6P7): 68 | last_results = self.top_blocks(x[-1], results[-1]) 69 | results.extend(last_results) 70 | elif isinstance(self.top_blocks, LastLevelMaxPool): 71 | last_results = self.top_blocks(results[-1]) 72 | results.extend(last_results) 73 | 74 | return tuple(results) 75 | 76 | 77 | class LastLevelMaxPool(nn.Module): 78 | def forward(self, x): 79 | return [F.max_pool2d(x, 1, 2, 0)] 80 | 81 | 82 | class LastLevelP6P7(nn.Module): 83 | """ 84 | This module is used in RetinaNet to generate extra layers, P6 and P7. 85 | """ 86 | def __init__(self, in_channels, out_channels): 87 | super(LastLevelP6P7, self).__init__() 88 | self.p6 = nn.Conv2d(in_channels, out_channels, 3, 2, 1) 89 | self.p7 = nn.Conv2d(out_channels, out_channels, 3, 2, 1) 90 | for module in [self.p6, self.p7]: 91 | nn.init.kaiming_uniform_(module.weight, a=1) 92 | nn.init.constant_(module.bias, 0) 93 | self.use_P5 = in_channels == out_channels 94 | 95 | def forward(self, c5, p5): 96 | x = p5 if self.use_P5 else c5 97 | p6 = self.p6(x) 98 | p7 = self.p7(F.relu(p6)) 99 | return [p6, p7] 100 | -------------------------------------------------------------------------------- /plane_mask_detection/maskrcnn_benchmark/modeling/balanced_positive_negative_sampler.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. 2 | import torch 3 | 4 | 5 | class BalancedPositiveNegativeSampler(object): 6 | """ 7 | This class samples batches, ensuring that they contain a fixed proportion of positives 8 | """ 9 | 10 | def __init__(self, batch_size_per_image, positive_fraction): 11 | """ 12 | Arguments: 13 | batch_size_per_image (int): number of elements to be selected per image 14 | positive_fraction (float): percentage of positive elements per batch 15 | """ 16 | self.batch_size_per_image = batch_size_per_image 17 | self.positive_fraction = positive_fraction 18 | 19 | def __call__(self, matched_idxs): 20 | """ 21 | Arguments: 22 | matched idxs: list of tensors containing -1, 0 or positive values. 23 | Each tensor corresponds to a specific image. 24 | -1 values are ignored, 0 are considered as negatives and > 0 as 25 | positives. 26 | 27 | Returns: 28 | pos_idx (list[tensor]) 29 | neg_idx (list[tensor]) 30 | 31 | Returns two lists of binary masks for each image. 32 | The first list contains the positive elements that were selected, 33 | and the second list the negative example. 34 | """ 35 | pos_idx = [] 36 | neg_idx = [] 37 | for matched_idxs_per_image in matched_idxs: 38 | positive = torch.nonzero(matched_idxs_per_image >= 1).squeeze(1) 39 | negative = torch.nonzero(matched_idxs_per_image == 0).squeeze(1) 40 | 41 | num_pos = int(self.batch_size_per_image * self.positive_fraction) 42 | # protect against not enough positive examples 43 | num_pos = min(positive.numel(), num_pos) 44 | num_neg = self.batch_size_per_image - num_pos 45 | # protect against not enough negative examples 46 | num_neg = min(negative.numel(), num_neg) 47 | 48 | # randomly select positive and negative examples 49 | perm1 = torch.randperm(positive.numel(), device=positive.device)[:num_pos] 50 | perm2 = torch.randperm(negative.numel(), device=negative.device)[:num_neg] 51 | 52 | pos_idx_per_image = positive[perm1] 53 | neg_idx_per_image = negative[perm2] 54 | 55 | # create binary mask from indices 56 | pos_idx_per_image_mask = torch.zeros_like( 57 | matched_idxs_per_image, dtype=torch.uint8 58 | ) 59 | neg_idx_per_image_mask = torch.zeros_like( 60 | matched_idxs_per_image, dtype=torch.uint8 61 | ) 62 | pos_idx_per_image_mask[pos_idx_per_image] = 1 63 | neg_idx_per_image_mask[neg_idx_per_image] = 1 64 | 65 | pos_idx.append(pos_idx_per_image_mask) 66 | neg_idx.append(neg_idx_per_image_mask) 67 | 68 | return pos_idx, neg_idx 69 | -------------------------------------------------------------------------------- /plane_mask_detection/maskrcnn_benchmark/modeling/box_coder.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. 2 | import math 3 | 4 | import torch 5 | 6 | 7 | class BoxCoder(object): 8 | """ 9 | This class encodes and decodes a set of bounding boxes into 10 | the representation used for training the regressors. 11 | """ 12 | 13 | def __init__(self, weights, bbox_xform_clip=math.log(1000. / 16)): 14 | """ 15 | Arguments: 16 | weights (4-element tuple) 17 | bbox_xform_clip (float) 18 | """ 19 | self.weights = weights 20 | self.bbox_xform_clip = bbox_xform_clip 21 | 22 | def encode(self, reference_boxes, proposals): 23 | """ 24 | Encode a set of proposals with respect to some 25 | reference boxes 26 | 27 | Arguments: 28 | reference_boxes (Tensor): reference boxes 29 | proposals (Tensor): boxes to be encoded 30 | """ 31 | 32 | TO_REMOVE = 1 # TODO remove 33 | ex_widths = proposals[:, 2] - proposals[:, 0] + TO_REMOVE 34 | ex_heights = proposals[:, 3] - proposals[:, 1] + TO_REMOVE 35 | ex_ctr_x = proposals[:, 0] + 0.5 * ex_widths 36 | ex_ctr_y = proposals[:, 1] + 0.5 * ex_heights 37 | 38 | gt_widths = reference_boxes[:, 2] - reference_boxes[:, 0] + TO_REMOVE 39 | gt_heights = reference_boxes[:, 3] - reference_boxes[:, 1] + TO_REMOVE 40 | gt_ctr_x = reference_boxes[:, 0] + 0.5 * gt_widths 41 | gt_ctr_y = reference_boxes[:, 1] + 0.5 * gt_heights 42 | 43 | wx, wy, ww, wh = self.weights 44 | targets_dx = wx * (gt_ctr_x - ex_ctr_x) / ex_widths 45 | targets_dy = wy * (gt_ctr_y - ex_ctr_y) / ex_heights 46 | targets_dw = ww * torch.log(gt_widths / ex_widths) 47 | targets_dh = wh * torch.log(gt_heights / ex_heights) 48 | 49 | targets = torch.stack((targets_dx, targets_dy, targets_dw, targets_dh), dim=1) 50 | return targets 51 | 52 | def decode(self, rel_codes, boxes): 53 | """ 54 | From a set of original boxes and encoded relative box offsets, 55 | get the decoded boxes. 56 | 57 | Arguments: 58 | rel_codes (Tensor): encoded boxes 59 | boxes (Tensor): reference boxes. 60 | """ 61 | 62 | boxes = boxes.to(rel_codes.dtype) 63 | 64 | TO_REMOVE = 1 # TODO remove 65 | widths = boxes[:, 2] - boxes[:, 0] + TO_REMOVE 66 | heights = boxes[:, 3] - boxes[:, 1] + TO_REMOVE 67 | ctr_x = boxes[:, 0] + 0.5 * widths 68 | ctr_y = boxes[:, 1] + 0.5 * heights 69 | 70 | wx, wy, ww, wh = self.weights 71 | dx = rel_codes[:, 0::4] / wx 72 | dy = rel_codes[:, 1::4] / wy 73 | dw = rel_codes[:, 2::4] / ww 74 | dh = rel_codes[:, 3::4] / wh 75 | 76 | # Prevent sending too large values into torch.exp() 77 | dw = torch.clamp(dw, max=self.bbox_xform_clip) 78 | dh = torch.clamp(dh, max=self.bbox_xform_clip) 79 | 80 | pred_ctr_x = dx * widths[:, None] + ctr_x[:, None] 81 | pred_ctr_y = dy * heights[:, None] + ctr_y[:, None] 82 | pred_w = torch.exp(dw) * widths[:, None] 83 | pred_h = torch.exp(dh) * heights[:, None] 84 | 85 | pred_boxes = torch.zeros_like(rel_codes) 86 | # x1 87 | pred_boxes[:, 0::4] = pred_ctr_x - 0.5 * pred_w 88 | # y1 89 | pred_boxes[:, 1::4] = pred_ctr_y - 0.5 * pred_h 90 | # x2 (note: "- 1" is correct; don't be fooled by the asymmetry) 91 | pred_boxes[:, 2::4] = pred_ctr_x + 0.5 * pred_w - 1 92 | # y2 (note: "- 1" is correct; don't be fooled by the asymmetry) 93 | pred_boxes[:, 3::4] = pred_ctr_y + 0.5 * pred_h - 1 94 | 95 | return pred_boxes 96 | -------------------------------------------------------------------------------- /plane_mask_detection/maskrcnn_benchmark/modeling/detector/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. 2 | from .detectors import build_detection_model 3 | -------------------------------------------------------------------------------- /plane_mask_detection/maskrcnn_benchmark/modeling/detector/detectors.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. 2 | from .generalized_rcnn import GeneralizedRCNN 3 | 4 | 5 | _DETECTION_META_ARCHITECTURES = {"GeneralizedRCNN": GeneralizedRCNN} 6 | 7 | 8 | def build_detection_model(cfg): 9 | meta_arch = _DETECTION_META_ARCHITECTURES[cfg.MODEL.META_ARCHITECTURE] 10 | return meta_arch(cfg) 11 | -------------------------------------------------------------------------------- /plane_mask_detection/maskrcnn_benchmark/modeling/make_layers.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. 2 | """ 3 | Miscellaneous utility functions 4 | """ 5 | 6 | import torch 7 | from torch import nn 8 | from torch.nn import functional as F 9 | from plane_mask_detection.maskrcnn_benchmark.config import cfg 10 | from plane_mask_detection.maskrcnn_benchmark.layers import Conv2d 11 | from plane_mask_detection.maskrcnn_benchmark.modeling.poolers import Pooler 12 | 13 | 14 | def get_group_gn(dim, dim_per_gp, num_groups): 15 | """get number of groups used by GroupNorm, based on number of channels.""" 16 | assert dim_per_gp == -1 or num_groups == -1, \ 17 | "GroupNorm: can only specify G or C/G." 18 | 19 | if dim_per_gp > 0: 20 | assert dim % dim_per_gp == 0, \ 21 | "dim: {}, dim_per_gp: {}".format(dim, dim_per_gp) 22 | group_gn = dim // dim_per_gp 23 | else: 24 | assert dim % num_groups == 0, \ 25 | "dim: {}, num_groups: {}".format(dim, num_groups) 26 | group_gn = num_groups 27 | 28 | return group_gn 29 | 30 | 31 | def group_norm(out_channels, affine=True, divisor=1): 32 | out_channels = out_channels // divisor 33 | dim_per_gp = cfg.MODEL.GROUP_NORM.DIM_PER_GP // divisor 34 | num_groups = cfg.MODEL.GROUP_NORM.NUM_GROUPS // divisor 35 | eps = cfg.MODEL.GROUP_NORM.EPSILON # default: 1e-5 36 | return torch.nn.GroupNorm( 37 | get_group_gn(out_channels, dim_per_gp, num_groups), 38 | out_channels, 39 | eps, 40 | affine 41 | ) 42 | 43 | 44 | def make_conv3x3( 45 | in_channels, 46 | out_channels, 47 | dilation=1, 48 | stride=1, 49 | use_gn=False, 50 | use_relu=False, 51 | kaiming_init=True 52 | ): 53 | conv = Conv2d( 54 | in_channels, 55 | out_channels, 56 | kernel_size=3, 57 | stride=stride, 58 | padding=dilation, 59 | dilation=dilation, 60 | bias=False if use_gn else True 61 | ) 62 | if kaiming_init: 63 | nn.init.kaiming_normal_( 64 | conv.weight, mode="fan_out", nonlinearity="relu" 65 | ) 66 | else: 67 | torch.nn.init.normal_(conv.weight, std=0.01) 68 | if not use_gn: 69 | nn.init.constant_(conv.bias, 0) 70 | module = [conv,] 71 | if use_gn: 72 | module.append(group_norm(out_channels)) 73 | if use_relu: 74 | module.append(nn.ReLU(inplace=True)) 75 | if len(module) > 1: 76 | return nn.Sequential(*module) 77 | return conv 78 | 79 | 80 | def make_fc(dim_in, hidden_dim, use_gn=False): 81 | ''' 82 | Caffe2 implementation uses XavierFill, which in fact 83 | corresponds to kaiming_uniform_ in PyTorch 84 | ''' 85 | if use_gn: 86 | fc = nn.Linear(dim_in, hidden_dim, bias=False) 87 | nn.init.kaiming_uniform_(fc.weight, a=1) 88 | return nn.Sequential(fc, group_norm(hidden_dim)) 89 | fc = nn.Linear(dim_in, hidden_dim) 90 | nn.init.kaiming_uniform_(fc.weight, a=1) 91 | nn.init.constant_(fc.bias, 0) 92 | return fc 93 | 94 | 95 | def conv_with_kaiming_uniform(use_gn=False, use_relu=False): 96 | def make_conv( 97 | in_channels, out_channels, kernel_size, stride=1, dilation=1 98 | ): 99 | conv = Conv2d( 100 | in_channels, 101 | out_channels, 102 | kernel_size=kernel_size, 103 | stride=stride, 104 | padding=dilation * (kernel_size - 1) // 2, 105 | dilation=dilation, 106 | bias=False if use_gn else True 107 | ) 108 | # Caffe2 implementation uses XavierFill, which in fact 109 | # corresponds to kaiming_uniform_ in PyTorch 110 | nn.init.kaiming_uniform_(conv.weight, a=1) 111 | if not use_gn: 112 | nn.init.constant_(conv.bias, 0) 113 | module = [conv,] 114 | if use_gn: 115 | module.append(group_norm(out_channels)) 116 | if use_relu: 117 | module.append(nn.ReLU(inplace=True)) 118 | if len(module) > 1: 119 | return nn.Sequential(*module) 120 | return conv 121 | 122 | return make_conv 123 | -------------------------------------------------------------------------------- /plane_mask_detection/maskrcnn_benchmark/modeling/registry.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. 2 | 3 | from plane_mask_detection.maskrcnn_benchmark.utils.registry import Registry 4 | 5 | BACKBONES = Registry() 6 | RPN_HEADS = Registry() 7 | ROI_BOX_FEATURE_EXTRACTORS = Registry() 8 | ROI_BOX_PREDICTOR = Registry() 9 | ROI_KEYPOINT_FEATURE_EXTRACTORS = Registry() 10 | ROI_KEYPOINT_PREDICTOR = Registry() 11 | ROI_MASK_FEATURE_EXTRACTORS = Registry() 12 | ROI_MASK_PREDICTOR = Registry() 13 | -------------------------------------------------------------------------------- /plane_mask_detection/maskrcnn_benchmark/modeling/roi_heads/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MARSLab-UMN/vi_depth_completion/b8d4bd3f3af674871a9fb8cd6166e49d785859ef/plane_mask_detection/maskrcnn_benchmark/modeling/roi_heads/__init__.py -------------------------------------------------------------------------------- /plane_mask_detection/maskrcnn_benchmark/modeling/roi_heads/box_head/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MARSLab-UMN/vi_depth_completion/b8d4bd3f3af674871a9fb8cd6166e49d785859ef/plane_mask_detection/maskrcnn_benchmark/modeling/roi_heads/box_head/__init__.py -------------------------------------------------------------------------------- /plane_mask_detection/maskrcnn_benchmark/modeling/roi_heads/box_head/box_head.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. 2 | import torch 3 | from torch import nn 4 | 5 | from .roi_box_feature_extractors import make_roi_box_feature_extractor 6 | from .roi_box_predictors import make_roi_box_predictor 7 | from .inference import make_roi_box_post_processor 8 | from .loss import make_roi_box_loss_evaluator 9 | 10 | 11 | class ROIBoxHead(torch.nn.Module): 12 | """ 13 | Generic Box Head class. 14 | """ 15 | 16 | def __init__(self, cfg, in_channels): 17 | super(ROIBoxHead, self).__init__() 18 | self.feature_extractor = make_roi_box_feature_extractor(cfg, in_channels) 19 | self.predictor = make_roi_box_predictor( 20 | cfg, self.feature_extractor.out_channels) 21 | self.post_processor = make_roi_box_post_processor(cfg) 22 | self.loss_evaluator = make_roi_box_loss_evaluator(cfg) 23 | 24 | def forward(self, features, proposals, targets=None): 25 | """ 26 | Arguments: 27 | features (list[Tensor]): feature-maps from possibly several levels 28 | proposals (list[BoxList]): proposal boxes 29 | targets (list[BoxList], optional): the ground-truth targets. 30 | 31 | Returns: 32 | x (Tensor): the result of the feature extractor 33 | proposals (list[BoxList]): during training, the subsampled proposals 34 | are returned. During testing, the predicted boxlists are returned 35 | losses (dict[Tensor]): During training, returns the losses for the 36 | head. During testing, returns an empty dict. 37 | """ 38 | 39 | if self.training: 40 | # Faster R-CNN subsamples during training the proposals with a fixed 41 | # positive / negative ratio 42 | with torch.no_grad(): 43 | proposals = self.loss_evaluator.subsample(proposals, targets) 44 | 45 | # extract features that will be fed to the final classifier. The 46 | # feature_extractor generally corresponds to the pooler + heads 47 | x = self.feature_extractor(features, proposals) 48 | # final classifier that converts the features into predictions 49 | class_logits, box_regression = self.predictor(x) 50 | 51 | if not self.training: 52 | result = self.post_processor((class_logits, box_regression), proposals) 53 | return x, result, {} 54 | 55 | loss_classifier, loss_box_reg = self.loss_evaluator( 56 | [class_logits], [box_regression] 57 | ) 58 | return ( 59 | x, 60 | proposals, 61 | dict(loss_classifier=loss_classifier, loss_box_reg=loss_box_reg), 62 | ) 63 | 64 | 65 | def build_roi_box_head(cfg, in_channels): 66 | """ 67 | Constructs a new box head. 68 | By default, uses ROIBoxHead, but if it turns out not to be enough, just register a new class 69 | and make it a parameter in the config 70 | """ 71 | return ROIBoxHead(cfg, in_channels) 72 | -------------------------------------------------------------------------------- /plane_mask_detection/maskrcnn_benchmark/modeling/roi_heads/box_head/roi_box_predictors.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. 2 | from plane_mask_detection.maskrcnn_benchmark.modeling import registry 3 | from torch import nn 4 | 5 | 6 | @registry.ROI_BOX_PREDICTOR.register("FastRCNNPredictor") 7 | class FastRCNNPredictor(nn.Module): 8 | def __init__(self, config, in_channels): 9 | super(FastRCNNPredictor, self).__init__() 10 | assert in_channels is not None 11 | 12 | num_inputs = in_channels 13 | 14 | num_classes = config.MODEL.ROI_BOX_HEAD.NUM_CLASSES 15 | self.avgpool = nn.AdaptiveAvgPool2d(1) 16 | self.cls_score = nn.Linear(num_inputs, num_classes) 17 | num_bbox_reg_classes = 2 if config.MODEL.CLS_AGNOSTIC_BBOX_REG else num_classes 18 | self.bbox_pred = nn.Linear(num_inputs, num_bbox_reg_classes * 4) 19 | 20 | nn.init.normal_(self.cls_score.weight, mean=0, std=0.01) 21 | nn.init.constant_(self.cls_score.bias, 0) 22 | 23 | nn.init.normal_(self.bbox_pred.weight, mean=0, std=0.001) 24 | nn.init.constant_(self.bbox_pred.bias, 0) 25 | 26 | def forward(self, x): 27 | x = self.avgpool(x) 28 | x = x.view(x.size(0), -1) 29 | cls_logit = self.cls_score(x) 30 | bbox_pred = self.bbox_pred(x) 31 | return cls_logit, bbox_pred 32 | 33 | 34 | @registry.ROI_BOX_PREDICTOR.register("FPNPredictor") 35 | class FPNPredictor(nn.Module): 36 | def __init__(self, cfg, in_channels): 37 | super(FPNPredictor, self).__init__() 38 | num_classes = cfg.MODEL.ROI_BOX_HEAD.NUM_CLASSES 39 | representation_size = in_channels 40 | 41 | self.cls_score = nn.Linear(representation_size, num_classes) 42 | num_bbox_reg_classes = 2 if cfg.MODEL.CLS_AGNOSTIC_BBOX_REG else num_classes 43 | self.bbox_pred = nn.Linear(representation_size, num_bbox_reg_classes * 4) 44 | 45 | nn.init.normal_(self.cls_score.weight, std=0.01) 46 | nn.init.normal_(self.bbox_pred.weight, std=0.001) 47 | for l in [self.cls_score, self.bbox_pred]: 48 | nn.init.constant_(l.bias, 0) 49 | 50 | def forward(self, x): 51 | if x.ndimension() == 4: 52 | assert list(x.shape[2:]) == [1, 1] 53 | x = x.view(x.size(0), -1) 54 | scores = self.cls_score(x) 55 | bbox_deltas = self.bbox_pred(x) 56 | 57 | return scores, bbox_deltas 58 | 59 | 60 | def make_roi_box_predictor(cfg, in_channels): 61 | func = registry.ROI_BOX_PREDICTOR[cfg.MODEL.ROI_BOX_HEAD.PREDICTOR] 62 | return func(cfg, in_channels) 63 | -------------------------------------------------------------------------------- /plane_mask_detection/maskrcnn_benchmark/modeling/roi_heads/keypoint_head/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MARSLab-UMN/vi_depth_completion/b8d4bd3f3af674871a9fb8cd6166e49d785859ef/plane_mask_detection/maskrcnn_benchmark/modeling/roi_heads/keypoint_head/__init__.py -------------------------------------------------------------------------------- /plane_mask_detection/maskrcnn_benchmark/modeling/roi_heads/keypoint_head/keypoint_head.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | from .roi_keypoint_feature_extractors import make_roi_keypoint_feature_extractor 4 | from .roi_keypoint_predictors import make_roi_keypoint_predictor 5 | from .inference import make_roi_keypoint_post_processor 6 | from .loss import make_roi_keypoint_loss_evaluator 7 | 8 | 9 | class ROIKeypointHead(torch.nn.Module): 10 | def __init__(self, cfg, in_channels): 11 | super(ROIKeypointHead, self).__init__() 12 | self.cfg = cfg.clone() 13 | self.feature_extractor = make_roi_keypoint_feature_extractor(cfg, in_channels) 14 | self.predictor = make_roi_keypoint_predictor( 15 | cfg, self.feature_extractor.out_channels) 16 | self.post_processor = make_roi_keypoint_post_processor(cfg) 17 | self.loss_evaluator = make_roi_keypoint_loss_evaluator(cfg) 18 | 19 | def forward(self, features, proposals, targets=None): 20 | """ 21 | Arguments: 22 | features (list[Tensor]): feature-maps from possibly several levels 23 | proposals (list[BoxList]): proposal boxes 24 | targets (list[BoxList], optional): the ground-truth targets. 25 | 26 | Returns: 27 | x (Tensor): the result of the feature extractor 28 | proposals (list[BoxList]): during training, the original proposals 29 | are returned. During testing, the predicted boxlists are returned 30 | with the `mask` field set 31 | losses (dict[Tensor]): During training, returns the losses for the 32 | head. During testing, returns an empty dict. 33 | """ 34 | if self.training: 35 | with torch.no_grad(): 36 | proposals = self.loss_evaluator.subsample(proposals, targets) 37 | 38 | x = self.feature_extractor(features, proposals) 39 | kp_logits = self.predictor(x) 40 | 41 | if not self.training: 42 | result = self.post_processor(kp_logits, proposals) 43 | return x, result, {} 44 | 45 | loss_kp = self.loss_evaluator(proposals, kp_logits) 46 | 47 | return x, proposals, dict(loss_kp=loss_kp) 48 | 49 | 50 | def build_roi_keypoint_head(cfg, in_channels): 51 | return ROIKeypointHead(cfg, in_channels) 52 | -------------------------------------------------------------------------------- /plane_mask_detection/maskrcnn_benchmark/modeling/roi_heads/keypoint_head/roi_keypoint_feature_extractors.py: -------------------------------------------------------------------------------- 1 | from torch import nn 2 | from torch.nn import functional as F 3 | 4 | from plane_mask_detection.maskrcnn_benchmark.modeling import registry 5 | from plane_mask_detection.maskrcnn_benchmark.modeling.poolers import Pooler 6 | 7 | from plane_mask_detection.maskrcnn_benchmark.layers import Conv2d 8 | 9 | 10 | @registry.ROI_KEYPOINT_FEATURE_EXTRACTORS.register("KeypointRCNNFeatureExtractor") 11 | class KeypointRCNNFeatureExtractor(nn.Module): 12 | def __init__(self, cfg, in_channels): 13 | super(KeypointRCNNFeatureExtractor, self).__init__() 14 | 15 | resolution = cfg.MODEL.ROI_KEYPOINT_HEAD.POOLER_RESOLUTION 16 | scales = cfg.MODEL.ROI_KEYPOINT_HEAD.POOLER_SCALES 17 | sampling_ratio = cfg.MODEL.ROI_KEYPOINT_HEAD.POOLER_SAMPLING_RATIO 18 | pooler = Pooler( 19 | output_size=(resolution, resolution), 20 | scales=scales, 21 | sampling_ratio=sampling_ratio, 22 | ) 23 | self.pooler = pooler 24 | 25 | input_features = in_channels 26 | layers = cfg.MODEL.ROI_KEYPOINT_HEAD.CONV_LAYERS 27 | next_feature = input_features 28 | self.blocks = [] 29 | for layer_idx, layer_features in enumerate(layers, 1): 30 | layer_name = "conv_fcn{}".format(layer_idx) 31 | module = Conv2d(next_feature, layer_features, 3, stride=1, padding=1) 32 | nn.init.kaiming_normal_(module.weight, mode="fan_out", nonlinearity="relu") 33 | nn.init.constant_(module.bias, 0) 34 | self.add_module(layer_name, module) 35 | next_feature = layer_features 36 | self.blocks.append(layer_name) 37 | self.out_channels = layer_features 38 | 39 | def forward(self, x, proposals): 40 | x = self.pooler(x, proposals) 41 | for layer_name in self.blocks: 42 | x = F.relu(getattr(self, layer_name)(x)) 43 | return x 44 | 45 | 46 | def make_roi_keypoint_feature_extractor(cfg, in_channels): 47 | func = registry.ROI_KEYPOINT_FEATURE_EXTRACTORS[ 48 | cfg.MODEL.ROI_KEYPOINT_HEAD.FEATURE_EXTRACTOR 49 | ] 50 | return func(cfg, in_channels) 51 | -------------------------------------------------------------------------------- /plane_mask_detection/maskrcnn_benchmark/modeling/roi_heads/keypoint_head/roi_keypoint_predictors.py: -------------------------------------------------------------------------------- 1 | from torch import nn 2 | 3 | from plane_mask_detection.maskrcnn_benchmark import layers 4 | from plane_mask_detection.maskrcnn_benchmark.modeling import registry 5 | 6 | 7 | @registry.ROI_KEYPOINT_PREDICTOR.register("KeypointRCNNPredictor") 8 | class KeypointRCNNPredictor(nn.Module): 9 | def __init__(self, cfg, in_channels): 10 | super(KeypointRCNNPredictor, self).__init__() 11 | input_features = in_channels 12 | num_keypoints = cfg.MODEL.ROI_KEYPOINT_HEAD.NUM_CLASSES 13 | deconv_kernel = 4 14 | self.kps_score_lowres = layers.ConvTranspose2d( 15 | input_features, 16 | num_keypoints, 17 | deconv_kernel, 18 | stride=2, 19 | padding=deconv_kernel // 2 - 1, 20 | ) 21 | nn.init.kaiming_normal_( 22 | self.kps_score_lowres.weight, mode="fan_out", nonlinearity="relu" 23 | ) 24 | nn.init.constant_(self.kps_score_lowres.bias, 0) 25 | self.up_scale = 2 26 | self.out_channels = num_keypoints 27 | 28 | def forward(self, x): 29 | x = self.kps_score_lowres(x) 30 | x = layers.interpolate( 31 | x, scale_factor=self.up_scale, mode="bilinear", align_corners=False 32 | ) 33 | return x 34 | 35 | 36 | def make_roi_keypoint_predictor(cfg, in_channels): 37 | func = registry.ROI_KEYPOINT_PREDICTOR[cfg.MODEL.ROI_KEYPOINT_HEAD.PREDICTOR] 38 | return func(cfg, in_channels) 39 | -------------------------------------------------------------------------------- /plane_mask_detection/maskrcnn_benchmark/modeling/roi_heads/mask_head/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MARSLab-UMN/vi_depth_completion/b8d4bd3f3af674871a9fb8cd6166e49d785859ef/plane_mask_detection/maskrcnn_benchmark/modeling/roi_heads/mask_head/__init__.py -------------------------------------------------------------------------------- /plane_mask_detection/maskrcnn_benchmark/modeling/roi_heads/mask_head/mask_head.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. 2 | import torch 3 | from torch import nn 4 | 5 | from plane_mask_detection.maskrcnn_benchmark.structures.bounding_box import BoxList 6 | 7 | from .roi_mask_feature_extractors import make_roi_mask_feature_extractor 8 | from .roi_mask_predictors import make_roi_mask_predictor 9 | from .inference import make_roi_mask_post_processor 10 | from .loss import make_roi_mask_loss_evaluator 11 | 12 | 13 | def keep_only_positive_boxes(boxes): 14 | """ 15 | Given a set of BoxList containing the `labels` field, 16 | return a set of BoxList for which `labels > 0`. 17 | 18 | Arguments: 19 | boxes (list of BoxList) 20 | """ 21 | assert isinstance(boxes, (list, tuple)) 22 | assert isinstance(boxes[0], BoxList) 23 | assert boxes[0].has_field("labels") 24 | positive_boxes = [] 25 | positive_inds = [] 26 | num_boxes = 0 27 | for boxes_per_image in boxes: 28 | labels = boxes_per_image.get_field("labels") 29 | inds_mask = labels > 0 30 | inds = inds_mask.nonzero().squeeze(1) 31 | positive_boxes.append(boxes_per_image[inds]) 32 | positive_inds.append(inds_mask) 33 | return positive_boxes, positive_inds 34 | 35 | 36 | class ROIMaskHead(torch.nn.Module): 37 | def __init__(self, cfg, in_channels): 38 | super(ROIMaskHead, self).__init__() 39 | self.cfg = cfg.clone() 40 | self.feature_extractor = make_roi_mask_feature_extractor(cfg, in_channels) 41 | self.predictor = make_roi_mask_predictor( 42 | cfg, self.feature_extractor.out_channels) 43 | self.post_processor = make_roi_mask_post_processor(cfg) 44 | self.loss_evaluator = make_roi_mask_loss_evaluator(cfg) 45 | 46 | def forward(self, features, proposals, targets=None, normal_prediction=None): 47 | # print('in mask_head:', proposals[0]) 48 | """ 49 | Arguments: 50 | features (list[Tensor]): feature-maps from possibly several levels 51 | proposals (list[BoxList]): proposal boxes 52 | targets (list[BoxList], optional): the ground-truth targets. 53 | 54 | Returns: 55 | x (Tensor): the result of the feature extractor 56 | proposals (list[BoxList]): during training, the original proposals 57 | are returned. During testing, the predicted boxlists are returned 58 | with the `mask` field set 59 | losses (dict[Tensor]): During training, returns the losses for the 60 | head. During testing, returns an empty dict. 61 | """ 62 | 63 | if self.training: 64 | # during training, only focus on positive boxes 65 | all_proposals = proposals 66 | proposals, positive_inds = keep_only_positive_boxes(proposals) 67 | if self.training and self.cfg.MODEL.ROI_MASK_HEAD.SHARE_BOX_FEATURE_EXTRACTOR: 68 | x = features 69 | x = x[torch.cat(positive_inds, dim=0)] 70 | else: 71 | x = self.feature_extractor(features, proposals) 72 | mask_logits = self.predictor(x) 73 | 74 | if not self.training: 75 | result = self.post_processor(mask_logits, proposals) 76 | return x, result, {} 77 | 78 | if normal_prediction is None: # or self.cfg.SOLVER.UNSUPERVISED_PLANE_NORMAL is False: 79 | loss_mask = self.loss_evaluator(proposals, mask_logits, targets) 80 | return x, all_proposals, dict(loss_mask=loss_mask) 81 | else: 82 | loss_mask, loss_plane_normal = self.loss_evaluator(proposals, mask_logits, targets, normal_prediction) 83 | return x, all_proposals, dict(loss_mask=loss_mask, loss_plane_normal=loss_plane_normal) 84 | 85 | 86 | def build_roi_mask_head(cfg, in_channels): 87 | return ROIMaskHead(cfg, in_channels) 88 | -------------------------------------------------------------------------------- /plane_mask_detection/maskrcnn_benchmark/modeling/roi_heads/mask_head/roi_mask_feature_extractors.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. 2 | from torch import nn 3 | from torch.nn import functional as F 4 | 5 | from ..box_head.roi_box_feature_extractors import ResNet50Conv5ROIFeatureExtractor 6 | from plane_mask_detection.maskrcnn_benchmark.modeling import registry 7 | from plane_mask_detection.maskrcnn_benchmark.modeling.poolers import Pooler 8 | from plane_mask_detection.maskrcnn_benchmark.modeling.make_layers import make_conv3x3 9 | 10 | 11 | registry.ROI_MASK_FEATURE_EXTRACTORS.register( 12 | "ResNet50Conv5ROIFeatureExtractor", ResNet50Conv5ROIFeatureExtractor 13 | ) 14 | 15 | 16 | @registry.ROI_MASK_FEATURE_EXTRACTORS.register("MaskRCNNFPNFeatureExtractor") 17 | class MaskRCNNFPNFeatureExtractor(nn.Module): 18 | """ 19 | Heads for FPN for classification 20 | """ 21 | 22 | def __init__(self, cfg, in_channels): 23 | """ 24 | Arguments: 25 | num_classes (int): number of output classes 26 | input_size (int): number of channels of the input once it's flattened 27 | representation_size (int): size of the intermediate representation 28 | """ 29 | super(MaskRCNNFPNFeatureExtractor, self).__init__() 30 | 31 | resolution = cfg.MODEL.ROI_MASK_HEAD.POOLER_RESOLUTION 32 | scales = cfg.MODEL.ROI_MASK_HEAD.POOLER_SCALES 33 | sampling_ratio = cfg.MODEL.ROI_MASK_HEAD.POOLER_SAMPLING_RATIO 34 | pooler = Pooler( 35 | output_size=(resolution, resolution), 36 | scales=scales, 37 | sampling_ratio=sampling_ratio, 38 | ) 39 | input_size = in_channels 40 | self.pooler = pooler 41 | 42 | use_gn = cfg.MODEL.ROI_MASK_HEAD.USE_GN 43 | layers = cfg.MODEL.ROI_MASK_HEAD.CONV_LAYERS 44 | dilation = cfg.MODEL.ROI_MASK_HEAD.DILATION 45 | 46 | next_feature = input_size 47 | self.blocks = [] 48 | for layer_idx, layer_features in enumerate(layers, 1): 49 | layer_name = "mask_fcn{}".format(layer_idx) 50 | module = make_conv3x3( 51 | next_feature, layer_features, 52 | dilation=dilation, stride=1, use_gn=use_gn 53 | ) 54 | self.add_module(layer_name, module) 55 | next_feature = layer_features 56 | self.blocks.append(layer_name) 57 | self.out_channels = layer_features 58 | 59 | def forward(self, x, proposals): 60 | x = self.pooler(x, proposals) 61 | 62 | for layer_name in self.blocks: 63 | x = F.relu(getattr(self, layer_name)(x)) 64 | 65 | return x 66 | 67 | 68 | def make_roi_mask_feature_extractor(cfg, in_channels): 69 | func = registry.ROI_MASK_FEATURE_EXTRACTORS[ 70 | cfg.MODEL.ROI_MASK_HEAD.FEATURE_EXTRACTOR 71 | ] 72 | return func(cfg, in_channels) 73 | -------------------------------------------------------------------------------- /plane_mask_detection/maskrcnn_benchmark/modeling/roi_heads/mask_head/roi_mask_predictors.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. 2 | from torch import nn 3 | from torch.nn import functional as F 4 | 5 | from plane_mask_detection.maskrcnn_benchmark.layers import Conv2d 6 | from plane_mask_detection.maskrcnn_benchmark.layers import ConvTranspose2d 7 | from plane_mask_detection.maskrcnn_benchmark.modeling import registry 8 | 9 | 10 | @registry.ROI_MASK_PREDICTOR.register("MaskRCNNC4Predictor") 11 | class MaskRCNNC4Predictor(nn.Module): 12 | def __init__(self, cfg, in_channels): 13 | super(MaskRCNNC4Predictor, self).__init__() 14 | num_classes = cfg.MODEL.ROI_BOX_HEAD.NUM_CLASSES 15 | dim_reduced = cfg.MODEL.ROI_MASK_HEAD.CONV_LAYERS[-1] 16 | num_inputs = in_channels 17 | 18 | self.conv5_mask = ConvTranspose2d(num_inputs, dim_reduced, 2, 2, 0) 19 | self.mask_fcn_logits = Conv2d(dim_reduced, num_classes, 1, 1, 0) 20 | 21 | for name, param in self.named_parameters(): 22 | if "bias" in name: 23 | nn.init.constant_(param, 0) 24 | elif "weight" in name: 25 | # Caffe2 implementation uses MSRAFill, which in fact 26 | # corresponds to kaiming_normal_ in PyTorch 27 | nn.init.kaiming_normal_(param, mode="fan_out", nonlinearity="relu") 28 | 29 | def forward(self, x): 30 | x = F.relu(self.conv5_mask(x)) 31 | return self.mask_fcn_logits(x) 32 | 33 | 34 | @registry.ROI_MASK_PREDICTOR.register("MaskRCNNConv1x1Predictor") 35 | class MaskRCNNConv1x1Predictor(nn.Module): 36 | def __init__(self, cfg, in_channels): 37 | super(MaskRCNNConv1x1Predictor, self).__init__() 38 | num_classes = cfg.MODEL.ROI_BOX_HEAD.NUM_CLASSES 39 | num_inputs = in_channels 40 | 41 | self.mask_fcn_logits = Conv2d(num_inputs, num_classes, 1, 1, 0) 42 | 43 | for name, param in self.named_parameters(): 44 | if "bias" in name: 45 | nn.init.constant_(param, 0) 46 | elif "weight" in name: 47 | # Caffe2 implementation uses MSRAFill, which in fact 48 | # corresponds to kaiming_normal_ in PyTorch 49 | nn.init.kaiming_normal_(param, mode="fan_out", nonlinearity="relu") 50 | 51 | def forward(self, x): 52 | return self.mask_fcn_logits(x) 53 | 54 | 55 | def make_roi_mask_predictor(cfg, in_channels): 56 | func = registry.ROI_MASK_PREDICTOR[cfg.MODEL.ROI_MASK_HEAD.PREDICTOR] 57 | return func(cfg, in_channels) 58 | -------------------------------------------------------------------------------- /plane_mask_detection/maskrcnn_benchmark/modeling/roi_heads/roi_heads.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. 2 | import torch 3 | 4 | from .box_head.box_head import build_roi_box_head 5 | from .mask_head.mask_head import build_roi_mask_head 6 | from .keypoint_head.keypoint_head import build_roi_keypoint_head 7 | 8 | 9 | class CombinedROIHeads(torch.nn.ModuleDict): 10 | """ 11 | Combines a set of individual heads (for box prediction or masks) into a single 12 | head. 13 | """ 14 | 15 | def __init__(self, cfg, heads): 16 | super(CombinedROIHeads, self).__init__(heads) 17 | self.cfg = cfg.clone() 18 | if cfg.MODEL.MASK_ON and cfg.MODEL.ROI_MASK_HEAD.SHARE_BOX_FEATURE_EXTRACTOR: 19 | self.mask.feature_extractor = self.box.feature_extractor 20 | if cfg.MODEL.KEYPOINT_ON and cfg.MODEL.ROI_KEYPOINT_HEAD.SHARE_BOX_FEATURE_EXTRACTOR: 21 | self.keypoint.feature_extractor = self.box.feature_extractor 22 | 23 | def forward(self, features, proposals, targets=None, normal_prediction=None): 24 | # print('in roi_heads:', proposals[0]) 25 | 26 | losses = {} 27 | # TODO rename x to roi_box_features, if it doesn't increase memory consumption 28 | x, detections, loss_box = self.box(features, proposals, targets) 29 | # print('in roi_heads:', detections[0]) 30 | losses.update(loss_box) 31 | if self.cfg.MODEL.MASK_ON: 32 | mask_features = features 33 | # optimization: during training, if we share the feature extractor between 34 | # the box and the mask heads, then we can reuse the features already computed 35 | if ( 36 | self.training 37 | and self.cfg.MODEL.ROI_MASK_HEAD.SHARE_BOX_FEATURE_EXTRACTOR 38 | ): 39 | mask_features = x 40 | # During training, self.box() will return the unaltered proposals as "detections" 41 | # this makes the API consistent during training and testing 42 | if normal_prediction is None: 43 | x, detections, loss_mask = self.mask(mask_features, detections, targets) 44 | else: 45 | x, detections, loss_mask = self.mask(mask_features, detections, targets, normal_prediction) 46 | losses.update(loss_mask) 47 | 48 | if self.cfg.MODEL.KEYPOINT_ON: 49 | keypoint_features = features 50 | # optimization: during training, if we share the feature extractor between 51 | # the box and the mask heads, then we can reuse the features already computed 52 | if ( 53 | self.training 54 | and self.cfg.MODEL.ROI_KEYPOINT_HEAD.SHARE_BOX_FEATURE_EXTRACTOR 55 | ): 56 | keypoint_features = x 57 | # During training, self.box() will return the unaltered proposals as "detections" 58 | # this makes the API consistent during training and testing 59 | x, detections, loss_keypoint = self.keypoint(keypoint_features, detections, targets) 60 | losses.update(loss_keypoint) 61 | return x, detections, losses 62 | 63 | 64 | def build_roi_heads(cfg, in_channels): 65 | # individually create the heads, that will be combined together 66 | # afterwards 67 | roi_heads = [] 68 | if cfg.MODEL.RETINANET_ON: 69 | return [] 70 | 71 | if not cfg.MODEL.RPN_ONLY: 72 | roi_heads.append(("box", build_roi_box_head(cfg, in_channels))) 73 | if cfg.MODEL.MASK_ON: 74 | roi_heads.append(("mask", build_roi_mask_head(cfg, in_channels))) 75 | if cfg.MODEL.KEYPOINT_ON: 76 | roi_heads.append(("keypoint", build_roi_keypoint_head(cfg, in_channels))) 77 | 78 | # combine individual heads in a single module 79 | if roi_heads: 80 | roi_heads = CombinedROIHeads(cfg, roi_heads) 81 | 82 | return roi_heads 83 | -------------------------------------------------------------------------------- /plane_mask_detection/maskrcnn_benchmark/modeling/rpn/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. 2 | # from .rpn import build_rpn 3 | -------------------------------------------------------------------------------- /plane_mask_detection/maskrcnn_benchmark/modeling/rpn/retinanet/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MARSLab-UMN/vi_depth_completion/b8d4bd3f3af674871a9fb8cd6166e49d785859ef/plane_mask_detection/maskrcnn_benchmark/modeling/rpn/retinanet/__init__.py -------------------------------------------------------------------------------- /plane_mask_detection/maskrcnn_benchmark/modeling/rpn/retinanet/loss.py: -------------------------------------------------------------------------------- 1 | """ 2 | This file contains specific functions for computing losses on the RetinaNet 3 | file 4 | """ 5 | 6 | import torch 7 | from torch.nn import functional as F 8 | 9 | from ..utils import concat_box_prediction_layers 10 | 11 | from plane_mask_detection.maskrcnn_benchmark.layers import smooth_l1_loss 12 | from plane_mask_detection.maskrcnn_benchmark.layers import SigmoidFocalLoss 13 | from plane_mask_detection.maskrcnn_benchmark.modeling.matcher import Matcher 14 | from plane_mask_detection.maskrcnn_benchmark.modeling.utils import cat 15 | from plane_mask_detection.maskrcnn_benchmark.structures.boxlist_ops import boxlist_iou 16 | from plane_mask_detection.maskrcnn_benchmark.structures.boxlist_ops import cat_boxlist 17 | from plane_mask_detection.maskrcnn_benchmark.modeling.rpn.loss import RPNLossComputation 18 | 19 | class RetinaNetLossComputation(RPNLossComputation): 20 | """ 21 | This class computes the RetinaNet loss. 22 | """ 23 | 24 | def __init__(self, proposal_matcher, box_coder, 25 | generate_labels_func, 26 | sigmoid_focal_loss, 27 | bbox_reg_beta=0.11, 28 | regress_norm=1.0): 29 | """ 30 | Arguments: 31 | proposal_matcher (Matcher) 32 | box_coder (BoxCoder) 33 | """ 34 | self.proposal_matcher = proposal_matcher 35 | self.box_coder = box_coder 36 | self.box_cls_loss_func = sigmoid_focal_loss 37 | self.bbox_reg_beta = bbox_reg_beta 38 | self.copied_fields = ['labels'] 39 | self.generate_labels_func = generate_labels_func 40 | self.discard_cases = ['between_thresholds'] 41 | self.regress_norm = regress_norm 42 | 43 | def __call__(self, anchors, box_cls, box_regression, targets): 44 | """ 45 | Arguments: 46 | anchors (list[BoxList]) 47 | box_cls (list[Tensor]) 48 | box_regression (list[Tensor]) 49 | targets (list[BoxList]) 50 | 51 | Returns: 52 | retinanet_cls_loss (Tensor) 53 | retinanet_regression_loss (Tensor 54 | """ 55 | anchors = [cat_boxlist(anchors_per_image) for anchors_per_image in anchors] 56 | labels, regression_targets = self.prepare_targets(anchors, targets) 57 | 58 | N = len(labels) 59 | box_cls, box_regression = \ 60 | concat_box_prediction_layers(box_cls, box_regression) 61 | 62 | labels = torch.cat(labels, dim=0) 63 | regression_targets = torch.cat(regression_targets, dim=0) 64 | pos_inds = torch.nonzero(labels > 0).squeeze(1) 65 | 66 | retinanet_regression_loss = smooth_l1_loss( 67 | box_regression[pos_inds], 68 | regression_targets[pos_inds], 69 | beta=self.bbox_reg_beta, 70 | size_average=False, 71 | ) / (max(1, pos_inds.numel() * self.regress_norm)) 72 | 73 | labels = labels.int() 74 | 75 | retinanet_cls_loss = self.box_cls_loss_func( 76 | box_cls, 77 | labels 78 | ) / (pos_inds.numel() + N) 79 | 80 | return retinanet_cls_loss, retinanet_regression_loss 81 | 82 | 83 | def generate_retinanet_labels(matched_targets): 84 | labels_per_image = matched_targets.get_field("labels") 85 | return labels_per_image 86 | 87 | 88 | def make_retinanet_loss_evaluator(cfg, box_coder): 89 | matcher = Matcher( 90 | cfg.MODEL.RETINANET.FG_IOU_THRESHOLD, 91 | cfg.MODEL.RETINANET.BG_IOU_THRESHOLD, 92 | allow_low_quality_matches=True, 93 | ) 94 | sigmoid_focal_loss = SigmoidFocalLoss( 95 | cfg.MODEL.RETINANET.LOSS_GAMMA, 96 | cfg.MODEL.RETINANET.LOSS_ALPHA 97 | ) 98 | 99 | loss_evaluator = RetinaNetLossComputation( 100 | matcher, 101 | box_coder, 102 | generate_retinanet_labels, 103 | sigmoid_focal_loss, 104 | bbox_reg_beta = cfg.MODEL.RETINANET.BBOX_REG_BETA, 105 | regress_norm = cfg.MODEL.RETINANET.BBOX_REG_WEIGHT, 106 | ) 107 | return loss_evaluator 108 | -------------------------------------------------------------------------------- /plane_mask_detection/maskrcnn_benchmark/modeling/rpn/utils.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. 2 | """ 3 | Utility functions minipulating the prediction layers 4 | """ 5 | 6 | from ..utils import cat 7 | 8 | import torch 9 | 10 | def permute_and_flatten(layer, N, A, C, H, W): 11 | layer = layer.view(N, -1, C, H, W) 12 | layer = layer.permute(0, 3, 4, 1, 2) 13 | layer = layer.reshape(N, -1, C) 14 | return layer 15 | 16 | 17 | def concat_box_prediction_layers(box_cls, box_regression): 18 | box_cls_flattened = [] 19 | box_regression_flattened = [] 20 | # for each feature level, permute the outputs to make them be in the 21 | # same format as the labels. Note that the labels are computed for 22 | # all feature levels concatenated, so we keep the same representation 23 | # for the objectness and the box_regression 24 | for box_cls_per_level, box_regression_per_level in zip( 25 | box_cls, box_regression 26 | ): 27 | N, AxC, H, W = box_cls_per_level.shape 28 | Ax4 = box_regression_per_level.shape[1] 29 | A = Ax4 // 4 30 | C = AxC // A 31 | box_cls_per_level = permute_and_flatten( 32 | box_cls_per_level, N, A, C, H, W 33 | ) 34 | box_cls_flattened.append(box_cls_per_level) 35 | 36 | box_regression_per_level = permute_and_flatten( 37 | box_regression_per_level, N, A, 4, H, W 38 | ) 39 | box_regression_flattened.append(box_regression_per_level) 40 | # concatenate on the first dimension (representing the feature levels), to 41 | # take into account the way the labels were generated (with all feature maps 42 | # being concatenated as well) 43 | box_cls = cat(box_cls_flattened, dim=1).reshape(-1, C) 44 | box_regression = cat(box_regression_flattened, dim=1).reshape(-1, 4) 45 | return box_cls, box_regression 46 | -------------------------------------------------------------------------------- /plane_mask_detection/maskrcnn_benchmark/modeling/upconv/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MARSLab-UMN/vi_depth_completion/b8d4bd3f3af674871a9fb8cd6166e49d785859ef/plane_mask_detection/maskrcnn_benchmark/modeling/upconv/__init__.py -------------------------------------------------------------------------------- /plane_mask_detection/maskrcnn_benchmark/modeling/utils.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. 2 | """ 3 | Miscellaneous utility functions 4 | """ 5 | 6 | import torch 7 | 8 | 9 | def cat(tensors, dim=0): 10 | """ 11 | Efficient version of torch.cat that avoids a copy if there is only a single element in a list 12 | """ 13 | assert isinstance(tensors, (list, tuple)) 14 | if len(tensors) == 1: 15 | return tensors[0] 16 | return torch.cat(tensors, dim) 17 | -------------------------------------------------------------------------------- /plane_mask_detection/maskrcnn_benchmark/solver/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. 2 | from .build import make_optimizer 3 | from .build import make_lr_scheduler 4 | from .lr_scheduler import WarmupMultiStepLR 5 | -------------------------------------------------------------------------------- /plane_mask_detection/maskrcnn_benchmark/solver/build.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. 2 | import torch 3 | 4 | from .lr_scheduler import WarmupMultiStepLR 5 | 6 | 7 | def make_optimizer(cfg, model): 8 | params = [] 9 | for key, value in model.named_parameters(): 10 | if not value.requires_grad: 11 | continue 12 | lr = cfg.SOLVER.BASE_LR 13 | weight_decay = cfg.SOLVER.WEIGHT_DECAY 14 | if "bias" in key: 15 | lr = cfg.SOLVER.BASE_LR * cfg.SOLVER.BIAS_LR_FACTOR 16 | weight_decay = cfg.SOLVER.WEIGHT_DECAY_BIAS 17 | params += [{"params": [value], "lr": lr, "weight_decay": weight_decay}] 18 | 19 | optimizer = torch.optim.SGD(params, lr, momentum=cfg.SOLVER.MOMENTUM) 20 | return optimizer 21 | 22 | 23 | def make_lr_scheduler(cfg, optimizer): 24 | return WarmupMultiStepLR( 25 | optimizer, 26 | cfg.SOLVER.STEPS, 27 | cfg.SOLVER.GAMMA, 28 | warmup_factor=cfg.SOLVER.WARMUP_FACTOR, 29 | warmup_iters=cfg.SOLVER.WARMUP_ITERS, 30 | warmup_method=cfg.SOLVER.WARMUP_METHOD, 31 | ) 32 | -------------------------------------------------------------------------------- /plane_mask_detection/maskrcnn_benchmark/solver/lr_scheduler.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. 2 | from bisect import bisect_right 3 | 4 | import torch 5 | 6 | 7 | # FIXME ideally this would be achieved with a CombinedLRScheduler, 8 | # separating MultiStepLR with WarmupLR 9 | # but the current LRScheduler design doesn't allow it 10 | class WarmupMultiStepLR(torch.optim.lr_scheduler._LRScheduler): 11 | def __init__( 12 | self, 13 | optimizer, 14 | milestones, 15 | gamma=0.1, 16 | warmup_factor=1.0 / 3, 17 | warmup_iters=500, 18 | warmup_method="linear", 19 | last_epoch=-1, 20 | ): 21 | if not list(milestones) == sorted(milestones): 22 | raise ValueError( 23 | "Milestones should be a list of" " increasing integers. Got {}", 24 | milestones, 25 | ) 26 | 27 | if warmup_method not in ("constant", "linear"): 28 | raise ValueError( 29 | "Only 'constant' or 'linear' warmup_method accepted" 30 | "got {}".format(warmup_method) 31 | ) 32 | self.milestones = milestones 33 | self.gamma = gamma 34 | self.warmup_factor = warmup_factor 35 | self.warmup_iters = warmup_iters 36 | self.warmup_method = warmup_method 37 | super(WarmupMultiStepLR, self).__init__(optimizer, last_epoch) 38 | 39 | def get_lr(self): 40 | warmup_factor = 1 41 | if self.last_epoch < self.warmup_iters: 42 | if self.warmup_method == "constant": 43 | warmup_factor = self.warmup_factor 44 | elif self.warmup_method == "linear": 45 | alpha = float(self.last_epoch) / self.warmup_iters 46 | warmup_factor = self.warmup_factor * (1 - alpha) + alpha 47 | return [ 48 | base_lr 49 | * warmup_factor 50 | * self.gamma ** bisect_right(self.milestones, self.last_epoch) 51 | for base_lr in self.base_lrs 52 | ] 53 | -------------------------------------------------------------------------------- /plane_mask_detection/maskrcnn_benchmark/structures/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MARSLab-UMN/vi_depth_completion/b8d4bd3f3af674871a9fb8cd6166e49d785859ef/plane_mask_detection/maskrcnn_benchmark/structures/__init__.py -------------------------------------------------------------------------------- /plane_mask_detection/maskrcnn_benchmark/structures/boxlist_ops.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. 2 | import torch 3 | 4 | from .bounding_box import BoxList 5 | 6 | from plane_mask_detection.maskrcnn_benchmark.layers import nms as _box_nms 7 | 8 | 9 | def boxlist_nms(boxlist, nms_thresh, max_proposals=-1, score_field="scores"): 10 | """ 11 | Performs non-maximum suppression on a boxlist, with scores specified 12 | in a boxlist field via score_field. 13 | 14 | Arguments: 15 | boxlist(BoxList) 16 | nms_thresh (float) 17 | max_proposals (int): if > 0, then only the top max_proposals are kept 18 | after non-maximum suppression 19 | score_field (str) 20 | """ 21 | if nms_thresh <= 0: 22 | return boxlist 23 | mode = boxlist.mode 24 | boxlist = boxlist.convert("xyxy") 25 | boxes = boxlist.bbox 26 | score = boxlist.get_field(score_field) 27 | keep = _box_nms(boxes, score, nms_thresh) 28 | if max_proposals > 0: 29 | keep = keep[: max_proposals] 30 | boxlist = boxlist[keep] 31 | return boxlist.convert(mode) 32 | 33 | 34 | def remove_small_boxes(boxlist, min_size): 35 | """ 36 | Only keep boxes with both sides >= min_size 37 | 38 | Arguments: 39 | boxlist (Boxlist) 40 | min_size (int) 41 | """ 42 | # TODO maybe add an API for querying the ws / hs 43 | xywh_boxes = boxlist.convert("xywh").bbox 44 | _, _, ws, hs = xywh_boxes.unbind(dim=1) 45 | keep = ( 46 | (ws >= min_size) & (hs >= min_size) 47 | ).nonzero().squeeze(1) 48 | return boxlist[keep] 49 | 50 | 51 | # implementation from https://github.com/kuangliu/torchcv/blob/master/torchcv/utils/box.py 52 | # with slight modifications 53 | def boxlist_iou(boxlist1, boxlist2): 54 | """Compute the intersection over union of two set of boxes. 55 | The box order must be (xmin, ymin, xmax, ymax). 56 | 57 | Arguments: 58 | box1: (BoxList) bounding boxes, sized [N,4]. 59 | box2: (BoxList) bounding boxes, sized [M,4]. 60 | 61 | Returns: 62 | (tensor) iou, sized [N,M]. 63 | 64 | Reference: 65 | https://github.com/chainer/chainercv/blob/master/chainercv/utils/bbox/bbox_iou.py 66 | """ 67 | if boxlist1.size != boxlist2.size: 68 | raise RuntimeError( 69 | "boxlists should have same image size, got {}, {}".format(boxlist1, boxlist2)) 70 | boxlist1 = boxlist1.convert("xyxy") 71 | boxlist2 = boxlist2.convert("xyxy") 72 | N = len(boxlist1) 73 | M = len(boxlist2) 74 | 75 | area1 = boxlist1.area() 76 | area2 = boxlist2.area() 77 | 78 | box1, box2 = boxlist1.bbox, boxlist2.bbox 79 | 80 | lt = torch.max(box1[:, None, :2], box2[:, :2]) # [N,M,2] 81 | rb = torch.min(box1[:, None, 2:], box2[:, 2:]) # [N,M,2] 82 | 83 | TO_REMOVE = 1 84 | 85 | wh = (rb - lt + TO_REMOVE).clamp(min=0) # [N,M,2] 86 | inter = wh[:, :, 0] * wh[:, :, 1] # [N,M] 87 | 88 | iou = inter / (area1[:, None] + area2 - inter) 89 | return iou 90 | 91 | 92 | # TODO redundant, remove 93 | def _cat(tensors, dim=0): 94 | """ 95 | Efficient version of torch.cat that avoids a copy if there is only a single element in a list 96 | """ 97 | assert isinstance(tensors, (list, tuple)) 98 | if len(tensors) == 1: 99 | return tensors[0] 100 | return torch.cat(tensors, dim) 101 | 102 | 103 | def cat_boxlist(bboxes): 104 | """ 105 | Concatenates a list of BoxList (having the same image size) into a 106 | single BoxList 107 | 108 | Arguments: 109 | bboxes (list[BoxList]) 110 | """ 111 | assert isinstance(bboxes, (list, tuple)) 112 | assert all(isinstance(bbox, BoxList) for bbox in bboxes) 113 | 114 | size = bboxes[0].size 115 | assert all(bbox.size == size for bbox in bboxes) 116 | 117 | mode = bboxes[0].mode 118 | assert all(bbox.mode == mode for bbox in bboxes) 119 | 120 | fields = set(bboxes[0].fields()) 121 | assert all(set(bbox.fields()) == fields for bbox in bboxes) 122 | 123 | cat_boxes = BoxList(_cat([bbox.bbox for bbox in bboxes], dim=0), size, mode) 124 | 125 | for field in fields: 126 | data = _cat([bbox.get_field(field) for bbox in bboxes], dim=0) 127 | cat_boxes.add_field(field, data) 128 | 129 | return cat_boxes 130 | -------------------------------------------------------------------------------- /plane_mask_detection/maskrcnn_benchmark/structures/image_list.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. 2 | from __future__ import division 3 | 4 | import torch 5 | 6 | 7 | class ImageList(object): 8 | """ 9 | Structure that holds a list of images (of possibly 10 | varying sizes) as a single tensor. 11 | This works by padding the images to the same size, 12 | and storing in a field the original sizes of each image 13 | """ 14 | 15 | def __init__(self, tensors, image_sizes): 16 | """ 17 | Arguments: 18 | tensors (tensor) 19 | image_sizes (list[tuple[int, int]]) 20 | """ 21 | self.tensors = tensors 22 | self.image_sizes = image_sizes 23 | 24 | def to(self, *args, **kwargs): 25 | cast_tensor = self.tensors.to(*args, **kwargs) 26 | return ImageList(cast_tensor, self.image_sizes) 27 | 28 | def to_gravity_tensor(gravity_dirs): 29 | gravity_dirs = torch.stack(gravity_dirs, dim=0) 30 | gravity_dirs = gravity_dirs.view(gravity_dirs.shape[0], gravity_dirs.shape[1], 1, 1) 31 | return gravity_dirs 32 | 33 | 34 | def to_image_list(tensors, size_divisible=0): 35 | """ 36 | tensors can be an ImageList, a torch.Tensor or 37 | an iterable of Tensors. It can't be a numpy array. 38 | When tensors is an iterable of Tensors, it pads 39 | the Tensors with zeros so that they have the same 40 | shape 41 | """ 42 | if isinstance(tensors, torch.Tensor) and size_divisible > 0: 43 | tensors = [tensors] 44 | 45 | if isinstance(tensors, ImageList): 46 | return tensors 47 | elif isinstance(tensors, torch.Tensor): 48 | # single tensor shape can be inferred 49 | if tensors.dim() == 3: 50 | tensors = tensors[None] 51 | assert tensors.dim() == 4 52 | image_sizes = [tensor.shape[-2:] for tensor in tensors] 53 | return ImageList(tensors, image_sizes) 54 | elif isinstance(tensors, (tuple, list)): 55 | max_size = tuple(max(s) for s in zip(*[img.shape for img in tensors])) 56 | 57 | # TODO Ideally, just remove this and let me model handle arbitrary 58 | # input sizs 59 | if size_divisible > 0: 60 | import math 61 | 62 | stride = size_divisible 63 | max_size = list(max_size) 64 | max_size[1] = int(math.ceil(max_size[1] / stride) * stride) 65 | max_size[2] = int(math.ceil(max_size[2] / stride) * stride) 66 | max_size = tuple(max_size) 67 | 68 | batch_shape = (len(tensors),) + max_size 69 | batched_imgs = tensors[0].new(*batch_shape).zero_() 70 | for img, pad_img in zip(tensors, batched_imgs): 71 | pad_img[: img.shape[0], : img.shape[1], : img.shape[2]].copy_(img) 72 | 73 | image_sizes = [im.shape[-2:] for im in tensors] 74 | 75 | return ImageList(batched_imgs, image_sizes) 76 | else: 77 | raise TypeError("Unsupported type for to_image_list: {}".format(type(tensors))) 78 | -------------------------------------------------------------------------------- /plane_mask_detection/maskrcnn_benchmark/utils/README.md: -------------------------------------------------------------------------------- 1 | # Utility functions 2 | 3 | This folder contain utility functions that are not used in the 4 | core library, but are useful for building models or training 5 | code using the config system. 6 | -------------------------------------------------------------------------------- /plane_mask_detection/maskrcnn_benchmark/utils/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/MARSLab-UMN/vi_depth_completion/b8d4bd3f3af674871a9fb8cd6166e49d785859ef/plane_mask_detection/maskrcnn_benchmark/utils/__init__.py -------------------------------------------------------------------------------- /plane_mask_detection/maskrcnn_benchmark/utils/collect_env.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. 2 | import PIL 3 | 4 | from torch.utils.collect_env import get_pretty_env_info 5 | 6 | 7 | def get_pil_version(): 8 | return "\n Pillow ({})".format(PIL.__version__) 9 | 10 | 11 | def collect_env_info(): 12 | env_str = get_pretty_env_info() 13 | env_str += get_pil_version() 14 | return env_str 15 | -------------------------------------------------------------------------------- /plane_mask_detection/maskrcnn_benchmark/utils/comm.py: -------------------------------------------------------------------------------- 1 | """ 2 | This file contains primitives for multi-gpu communication. 3 | This is useful when doing distributed training. 4 | """ 5 | 6 | import pickle 7 | import time 8 | 9 | import torch 10 | import torch.distributed as dist 11 | 12 | 13 | def get_world_size(): 14 | if not dist.is_available(): 15 | return 1 16 | if not dist.is_initialized(): 17 | return 1 18 | return dist.get_world_size() 19 | 20 | 21 | def get_rank(): 22 | if not dist.is_available(): 23 | return 0 24 | if not dist.is_initialized(): 25 | return 0 26 | return dist.get_rank() 27 | 28 | 29 | def is_main_process(): 30 | return get_rank() == 0 31 | 32 | 33 | def synchronize(): 34 | """ 35 | Helper function to synchronize (barrier) among all processes when 36 | using distributed training 37 | """ 38 | if not dist.is_available(): 39 | return 40 | if not dist.is_initialized(): 41 | return 42 | world_size = dist.get_world_size() 43 | if world_size == 1: 44 | return 45 | dist.barrier() 46 | 47 | 48 | def all_gather(data): 49 | """ 50 | Run all_gather on arbitrary picklable data (not necessarily tensors) 51 | Args: 52 | data: any picklable object 53 | Returns: 54 | list[data]: list of data gathered from each rank 55 | """ 56 | world_size = get_world_size() 57 | if world_size == 1: 58 | return [data] 59 | 60 | # serialized to a Tensor 61 | buffer = pickle.dumps(data) 62 | storage = torch.ByteStorage.from_buffer(buffer) 63 | tensor = torch.ByteTensor(storage).to("cuda") 64 | 65 | # obtain Tensor size of each rank 66 | local_size = torch.LongTensor([tensor.numel()]).to("cuda") 67 | size_list = [torch.LongTensor([0]).to("cuda") for _ in range(world_size)] 68 | dist.all_gather(size_list, local_size) 69 | size_list = [int(size.item()) for size in size_list] 70 | max_size = max(size_list) 71 | 72 | # receiving Tensor from all ranks 73 | # we pad the tensor because torch all_gather does not support 74 | # gathering tensors of different shapes 75 | tensor_list = [] 76 | for _ in size_list: 77 | tensor_list.append(torch.ByteTensor(size=(max_size,)).to("cuda")) 78 | if local_size != max_size: 79 | padding = torch.ByteTensor(size=(max_size - local_size,)).to("cuda") 80 | tensor = torch.cat((tensor, padding), dim=0) 81 | dist.all_gather(tensor_list, tensor) 82 | 83 | data_list = [] 84 | for size, tensor in zip(size_list, tensor_list): 85 | buffer = tensor.cpu().numpy().tobytes()[:size] 86 | data_list.append(pickle.loads(buffer)) 87 | 88 | return data_list 89 | 90 | 91 | def reduce_dict(input_dict, average=True): 92 | """ 93 | Args: 94 | input_dict (dict): all the values will be reduced 95 | average (bool): whether to do average or sum 96 | Reduce the values in the dictionary from all processes so that process with rank 97 | 0 has the averaged results. Returns a dict with the same fields as 98 | input_dict, after reduction. 99 | """ 100 | world_size = get_world_size() 101 | if world_size < 2: 102 | return input_dict 103 | with torch.no_grad(): 104 | names = [] 105 | values = [] 106 | # sort the keys so that they are consistent across processes 107 | for k in sorted(input_dict.keys()): 108 | names.append(k) 109 | values.append(input_dict[k]) 110 | values = torch.stack(values, dim=0) 111 | dist.reduce(values, dst=0) 112 | if dist.get_rank() == 0 and average: 113 | # only main process gets accumulated, so only divide by 114 | # world_size in this case 115 | values /= world_size 116 | reduced_dict = {k: v for k, v in zip(names, values)} 117 | return reduced_dict 118 | -------------------------------------------------------------------------------- /plane_mask_detection/maskrcnn_benchmark/utils/cv2_util.py: -------------------------------------------------------------------------------- 1 | """ 2 | Module for cv2 utility functions and maintaining version compatibility 3 | between 3.x and 4.x 4 | """ 5 | import cv2 6 | 7 | 8 | def findContours(*args, **kwargs): 9 | """ 10 | Wraps cv2.findContours to maintain compatiblity between versions 11 | 3 and 4 12 | 13 | Returns: 14 | contours, hierarchy 15 | """ 16 | if cv2.__version__.startswith('4'): 17 | contours, hierarchy = cv2.findContours(*args, **kwargs) 18 | elif cv2.__version__.startswith('3'): 19 | _, contours, hierarchy = cv2.findContours(*args, **kwargs) 20 | else: 21 | raise AssertionError( 22 | 'cv2 must be either version 3 or 4 to call this method') 23 | 24 | return contours, hierarchy 25 | -------------------------------------------------------------------------------- /plane_mask_detection/maskrcnn_benchmark/utils/env.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. 2 | import os 3 | 4 | from plane_mask_detection.maskrcnn_benchmark.utils.imports import import_file 5 | 6 | 7 | def setup_environment(): 8 | """Perform environment setup work. The default setup is a no-op, but this 9 | function allows the user to specify a Python source file that performs 10 | custom setup work that may be necessary to their computing environment. 11 | """ 12 | custom_module_path = os.environ.get("TORCH_DETECTRON_ENV_MODULE") 13 | if custom_module_path: 14 | setup_custom_environment(custom_module_path) 15 | else: 16 | # The default setup is a no-op 17 | pass 18 | 19 | 20 | def setup_custom_environment(custom_module_path): 21 | """Load custom environment setup from a Python source file and run the setup 22 | function. 23 | """ 24 | module = import_file("maskrcnn_benchmark.utils.env.custom_module", custom_module_path) 25 | assert hasattr(module, "setup_environment") and callable( 26 | module.setup_environment 27 | ), ( 28 | "Custom environment module defined in {} does not have the " 29 | "required callable attribute 'setup_environment'." 30 | ).format( 31 | custom_module_path 32 | ) 33 | module.setup_environment() 34 | 35 | 36 | # Force environment setup when this module is imported 37 | setup_environment() 38 | -------------------------------------------------------------------------------- /plane_mask_detection/maskrcnn_benchmark/utils/imports.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. 2 | import torch 3 | 4 | if torch._six.PY3: 5 | import importlib 6 | import importlib.util 7 | import sys 8 | 9 | 10 | # from https://stackoverflow.com/questions/67631/how-to-import-a-module-given-the-full-path?utm_medium=organic&utm_source=google_rich_qa&utm_campaign=google_rich_qa 11 | def import_file(module_name, file_path, make_importable=False): 12 | spec = importlib.util.spec_from_file_location(module_name, file_path) 13 | module = importlib.util.module_from_spec(spec) 14 | spec.loader.exec_module(module) 15 | if make_importable: 16 | sys.modules[module_name] = module 17 | return module 18 | else: 19 | import imp 20 | 21 | def import_file(module_name, file_path, make_importable=None): 22 | module = imp.load_source(module_name, file_path) 23 | return module 24 | -------------------------------------------------------------------------------- /plane_mask_detection/maskrcnn_benchmark/utils/logger.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. 2 | import logging 3 | import os 4 | import sys 5 | 6 | 7 | def setup_logger(name, save_dir, distributed_rank, filename="log.txt"): 8 | logger = logging.getLogger(name) 9 | logger.setLevel(logging.DEBUG) 10 | # don't log results for the non-master process 11 | if distributed_rank > 0: 12 | return logger 13 | ch = logging.StreamHandler(stream=sys.stdout) 14 | ch.setLevel(logging.DEBUG) 15 | formatter = logging.Formatter("%(asctime)s %(name)s %(levelname)s: %(message)s") 16 | ch.setFormatter(formatter) 17 | logger.addHandler(ch) 18 | 19 | if save_dir: 20 | fh = logging.FileHandler(os.path.join(save_dir, filename)) 21 | fh.setLevel(logging.DEBUG) 22 | fh.setFormatter(formatter) 23 | logger.addHandler(fh) 24 | 25 | return logger 26 | -------------------------------------------------------------------------------- /plane_mask_detection/maskrcnn_benchmark/utils/metric_logger.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. 2 | from collections import defaultdict 3 | from collections import deque 4 | 5 | import torch 6 | 7 | 8 | class SmoothedValue(object): 9 | """Track a series of values and provide access to smoothed values over a 10 | window or the global series average. 11 | """ 12 | 13 | def __init__(self, window_size=20): 14 | self.deque = deque(maxlen=window_size) 15 | self.series = [] 16 | self.total = 0.0 17 | self.count = 0 18 | 19 | def update(self, value): 20 | self.deque.append(value) 21 | self.series.append(value) 22 | self.count += 1 23 | self.total += value 24 | 25 | @property 26 | def median(self): 27 | d = torch.tensor(list(self.deque)) 28 | return d.median().item() 29 | 30 | @property 31 | def avg(self): 32 | d = torch.tensor(list(self.deque)) 33 | return d.mean().item() 34 | 35 | @property 36 | def global_avg(self): 37 | return self.total / self.count 38 | 39 | 40 | class MetricLogger(object): 41 | def __init__(self, delimiter="\t"): 42 | self.meters = defaultdict(SmoothedValue) 43 | self.delimiter = delimiter 44 | 45 | def update(self, **kwargs): 46 | for k, v in kwargs.items(): 47 | if isinstance(v, torch.Tensor): 48 | v = v.item() 49 | assert isinstance(v, (float, int)) 50 | self.meters[k].update(v) 51 | 52 | def __getattr__(self, attr): 53 | if attr in self.meters: 54 | return self.meters[attr] 55 | if attr in self.__dict__: 56 | return self.__dict__[attr] 57 | raise AttributeError("'{}' object has no attribute '{}'".format( 58 | type(self).__name__, attr)) 59 | 60 | def __str__(self): 61 | loss_str = [] 62 | for name, meter in self.meters.items(): 63 | loss_str.append( 64 | "{}: {:.4f} ({:.4f})".format(name, meter.median, meter.global_avg) 65 | ) 66 | return self.delimiter.join(loss_str) 67 | -------------------------------------------------------------------------------- /plane_mask_detection/maskrcnn_benchmark/utils/miscellaneous.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. 2 | import errno 3 | import json 4 | import logging 5 | import os 6 | from .comm import is_main_process 7 | 8 | 9 | def mkdir(path): 10 | try: 11 | os.makedirs(path) 12 | except OSError as e: 13 | if e.errno != errno.EEXIST: 14 | raise 15 | 16 | 17 | def save_labels(dataset_list, output_dir): 18 | if is_main_process(): 19 | logger = logging.getLogger(__name__) 20 | 21 | ids_to_labels = {} 22 | for dataset in dataset_list: 23 | if hasattr(dataset, 'categories'): 24 | ids_to_labels.update(dataset.categories) 25 | else: 26 | logger.warning("Dataset [{}] has no categories attribute, labels.json file won't be created".format( 27 | dataset.__class__.__name__)) 28 | 29 | if ids_to_labels: 30 | labels_file = os.path.join(output_dir, 'labels.json') 31 | logger.info("Saving labels mapping into {}".format(labels_file)) 32 | with open(labels_file, 'w') as f: 33 | json.dump(ids_to_labels, f, indent=2) 34 | 35 | 36 | def save_config(cfg, path): 37 | if is_main_process(): 38 | with open(path, 'w') as f: 39 | f.write(cfg.dump()) 40 | -------------------------------------------------------------------------------- /plane_mask_detection/maskrcnn_benchmark/utils/model_serialization.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. 2 | from collections import OrderedDict 3 | import logging 4 | 5 | import torch 6 | 7 | from plane_mask_detection.maskrcnn_benchmark.utils.imports import import_file 8 | 9 | 10 | def align_and_update_state_dicts(model_state_dict, loaded_state_dict): 11 | """ 12 | Strategy: suppose that the models that we will create will have prefixes appended 13 | to each of its keys, for example due to an extra level of nesting that the original 14 | pre-trained weights from ImageNet won't contain. For example, model.state_dict() 15 | might return backbone[0].body.res2.conv1.weight, while the pre-trained model contains 16 | res2.conv1.weight. We thus want to match both parameters together. 17 | For that, we look for each model weight, look among all loaded keys if there is one 18 | that is a suffix of the current weight name, and use it if that's the case. 19 | If multiple matches exist, take the one with longest size 20 | of the corresponding name. For example, for the same model as before, the pretrained 21 | weight file can contain both res2.conv1.weight, as well as conv1.weight. In this case, 22 | we want to match backbone[0].body.conv1.weight to conv1.weight, and 23 | backbone[0].body.res2.conv1.weight to res2.conv1.weight. 24 | """ 25 | current_keys = sorted(list(model_state_dict.keys())) 26 | loaded_keys = sorted(list(loaded_state_dict.keys())) 27 | # get a matrix of string matches, where each (i, j) entry correspond to the size of the 28 | # loaded_key string, if it matches 29 | match_matrix = [ 30 | len(j) if i.endswith(j) else 0 for i in current_keys for j in loaded_keys 31 | ] 32 | match_matrix = torch.as_tensor(match_matrix).view( 33 | len(current_keys), len(loaded_keys) 34 | ) 35 | max_match_size, idxs = match_matrix.max(1) 36 | # remove indices that correspond to no-match 37 | idxs[max_match_size == 0] = -1 38 | 39 | # used for logging 40 | max_size = max([len(key) for key in current_keys]) if current_keys else 1 41 | max_size_loaded = max([len(key) for key in loaded_keys]) if loaded_keys else 1 42 | log_str_template = "{: <{}} loaded from {: <{}} of shape {}" 43 | logger = logging.getLogger(__name__) 44 | for idx_new, idx_old in enumerate(idxs.tolist()): 45 | if idx_old == -1: 46 | continue 47 | key = current_keys[idx_new] 48 | key_old = loaded_keys[idx_old] 49 | model_state_dict[key] = loaded_state_dict[key_old] 50 | # logger.info( 51 | # log_str_template.format( 52 | # key, 53 | # max_size, 54 | # key_old, 55 | # max_size_loaded, 56 | # tuple(loaded_state_dict[key_old].shape), 57 | # ) 58 | # ) 59 | 60 | 61 | def strip_prefix_if_present(state_dict, prefix): 62 | keys = sorted(state_dict.keys()) 63 | if not all(key.startswith(prefix) for key in keys): 64 | return state_dict 65 | stripped_state_dict = OrderedDict() 66 | for key, value in state_dict.items(): 67 | stripped_state_dict[key.replace(prefix, "")] = value 68 | return stripped_state_dict 69 | 70 | 71 | def load_state_dict(model, loaded_state_dict): 72 | model_state_dict = model.state_dict() 73 | # if the state_dict comes from a model that was wrapped in a 74 | # DataParallel or DistributedDataParallel during serialization, 75 | # remove the "module" prefix before performing the matching 76 | loaded_state_dict = strip_prefix_if_present(loaded_state_dict, prefix="module.") 77 | align_and_update_state_dicts(model_state_dict, loaded_state_dict) 78 | 79 | # use strict loading 80 | model.load_state_dict(model_state_dict) 81 | -------------------------------------------------------------------------------- /plane_mask_detection/maskrcnn_benchmark/utils/model_zoo.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. 2 | import os 3 | import sys 4 | 5 | try: 6 | from torch.hub import _download_url_to_file 7 | from torch.hub import urlparse 8 | from torch.hub import HASH_REGEX 9 | except ImportError: 10 | from torch.utils.model_zoo import _download_url_to_file 11 | from torch.utils.model_zoo import urlparse 12 | from torch.utils.model_zoo import HASH_REGEX 13 | 14 | from plane_mask_detection.maskrcnn_benchmark.utils.comm import is_main_process 15 | from plane_mask_detection.maskrcnn_benchmark.utils.comm import synchronize 16 | 17 | 18 | # very similar to https://github.com/pytorch/pytorch/blob/master/torch/utils/model_zoo.py 19 | # but with a few improvements and modifications 20 | def cache_url(url, model_dir=None, progress=True): 21 | r"""Loads the Torch serialized object at the given URL. 22 | If the object is already present in `model_dir`, it's deserialized and 23 | returned. The filename part of the URL should follow the naming convention 24 | ``filename-.ext`` where ```` is the first eight or more 25 | digits of the SHA256 hash of the contents of the file. The hash is used to 26 | ensure unique names and to verify the contents of the file. 27 | The default value of `model_dir` is ``$TORCH_HOME/models`` where 28 | ``$TORCH_HOME`` defaults to ``~/.torch``. The default directory can be 29 | overridden with the ``$TORCH_MODEL_ZOO`` environment variable. 30 | Args: 31 | url (string): URL of the object to download 32 | model_dir (string, optional): directory in which to save the object 33 | progress (bool, optional): whether or not to display a progress bar to stderr 34 | Example: 35 | >>> cached_file = maskrcnn_benchmark.utils.model_zoo.cache_url('https://s3.amazonaws.com/pytorch/models/resnet18-5c106cde.pth') 36 | """ 37 | if model_dir is None: 38 | torch_home = os.path.expanduser(os.getenv("TORCH_HOME", "~/.torch")) 39 | model_dir = os.getenv("TORCH_MODEL_ZOO", os.path.join(torch_home, "models")) 40 | if not os.path.exists(model_dir): 41 | os.makedirs(model_dir) 42 | parts = urlparse(url) 43 | filename = os.path.basename(parts.path) 44 | if filename == "model_final.pkl": 45 | # workaround as pre-trained Caffe2 models from Detectron have all the same filename 46 | # so make the full path the filename by replacing / with _ 47 | filename = parts.path.replace("/", "_") 48 | cached_file = os.path.join(model_dir, filename) 49 | if not os.path.exists(cached_file) and is_main_process(): 50 | sys.stderr.write('Downloading: "{}" to {}\n'.format(url, cached_file)) 51 | hash_prefix = HASH_REGEX.search(filename) 52 | if hash_prefix is not None: 53 | hash_prefix = hash_prefix.group(1) 54 | # workaround: Caffe2 models don't have a hash, but follow the R-50 convention, 55 | # which matches the hash PyTorch uses. So we skip the hash matching 56 | # if the hash_prefix is less than 6 characters 57 | if len(hash_prefix) < 6: 58 | hash_prefix = None 59 | _download_url_to_file(url, cached_file, hash_prefix, progress=progress) 60 | synchronize() 61 | return cached_file 62 | -------------------------------------------------------------------------------- /plane_mask_detection/maskrcnn_benchmark/utils/registry.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. 2 | 3 | 4 | def _register_generic(module_dict, module_name, module): 5 | assert module_name not in module_dict 6 | module_dict[module_name] = module 7 | 8 | 9 | class Registry(dict): 10 | ''' 11 | A helper class for managing registering modules, it extends a dictionary 12 | and provides a register functions. 13 | 14 | Eg. creeting a registry: 15 | some_registry = Registry({"default": default_module}) 16 | 17 | There're two ways of registering new modules: 18 | 1): normal way is just calling register function: 19 | def foo(): 20 | ... 21 | some_registry.register("foo_module", foo) 22 | 2): used as decorator when declaring the module: 23 | @some_registry.register("foo_module") 24 | @some_registry.register("foo_modeul_nickname") 25 | def foo(): 26 | ... 27 | 28 | Access of module is just like using a dictionary, eg: 29 | f = some_registry["foo_modeul"] 30 | ''' 31 | def __init__(self, *args, **kwargs): 32 | super(Registry, self).__init__(*args, **kwargs) 33 | 34 | def register(self, module_name, module=None): 35 | # used as function call 36 | if module is not None: 37 | _register_generic(self, module_name, module) 38 | return 39 | 40 | # used as decorator 41 | def register_fn(fn): 42 | _register_generic(self, module_name, fn) 43 | return fn 44 | 45 | return register_fn 46 | -------------------------------------------------------------------------------- /plane_mask_detection/maskrcnn_benchmark/utils/timer.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. 2 | 3 | 4 | import time 5 | import datetime 6 | 7 | 8 | class Timer(object): 9 | def __init__(self): 10 | self.reset() 11 | 12 | @property 13 | def average_time(self): 14 | return self.total_time / self.calls if self.calls > 0 else 0.0 15 | 16 | def tic(self): 17 | # using time.time instead of time.clock because time time.clock 18 | # does not normalize for multithreading 19 | self.start_time = time.time() 20 | 21 | def toc(self, average=True): 22 | self.add(time.time() - self.start_time) 23 | if average: 24 | return self.average_time 25 | else: 26 | return self.diff 27 | 28 | def add(self, time_diff): 29 | self.diff = time_diff 30 | self.total_time += self.diff 31 | self.calls += 1 32 | 33 | def reset(self): 34 | self.total_time = 0.0 35 | self.calls = 0 36 | self.start_time = 0.0 37 | self.diff = 0.0 38 | 39 | def avg_time_str(self): 40 | time_str = str(datetime.timedelta(seconds=self.average_time)) 41 | return time_str 42 | 43 | 44 | def get_time_str(time_diff): 45 | time_str = str(datetime.timedelta(seconds=time_diff)) 46 | return time_str 47 | -------------------------------------------------------------------------------- /plane_mask_detection/tests/env_tests/env.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. 2 | 3 | import os 4 | 5 | 6 | def get_config_root_path(): 7 | ''' Path to configs for unit tests ''' 8 | # cur_file_dir is root/tests/env_tests 9 | cur_file_dir = os.path.dirname(os.path.abspath(os.path.realpath(__file__))) 10 | ret = os.path.dirname(os.path.dirname(cur_file_dir)) 11 | ret = os.path.join(ret, "configs") 12 | return ret 13 | -------------------------------------------------------------------------------- /plane_mask_detection/tests/test_backbones.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. 2 | 3 | import unittest 4 | import copy 5 | import torch 6 | # import modules to to register backbones 7 | from pytorch_local.maskrcnn_benchmark.modeling.backbone import build_backbone # NoQA 8 | from pytorch_local.maskrcnn_benchmark.modeling import registry 9 | from pytorch_local.maskrcnn_benchmark.config import cfg as g_cfg 10 | from utils import load_config 11 | 12 | 13 | # overwrite configs if specified, otherwise default config is used 14 | BACKBONE_CFGS = { 15 | "R-50-FPN": "e2e_faster_rcnn_R_50_FPN_1x.yaml", 16 | "R-101-FPN": "e2e_faster_rcnn_R_101_FPN_1x.yaml", 17 | "R-152-FPN": "e2e_faster_rcnn_R_101_FPN_1x.yaml", 18 | "R-50-FPN-RETINANET": "retinanet/retinanet_R-50-FPN_1x.yaml", 19 | "R-101-FPN-RETINANET": "retinanet/retinanet_R-101-FPN_1x.yaml", 20 | } 21 | 22 | 23 | class TestBackbones(unittest.TestCase): 24 | def test_build_backbones(self): 25 | ''' Make sure backbones run ''' 26 | 27 | self.assertGreater(len(registry.BACKBONES), 0) 28 | 29 | for name, backbone_builder in registry.BACKBONES.items(): 30 | print('Testing {}...'.format(name)) 31 | if name in BACKBONE_CFGS: 32 | cfg = load_config(BACKBONE_CFGS[name]) 33 | else: 34 | # Use default config if config file is not specified 35 | cfg = copy.deepcopy(g_cfg) 36 | backbone = backbone_builder(cfg) 37 | 38 | # make sures the backbone has `out_channels` 39 | self.assertIsNotNone( 40 | getattr(backbone, 'out_channels', None), 41 | 'Need to provide out_channels for backbone {}'.format(name) 42 | ) 43 | 44 | N, C_in, H, W = 2, 3, 224, 256 45 | input = torch.rand([N, C_in, H, W], dtype=torch.float32) 46 | out = backbone(input) 47 | for cur_out in out: 48 | self.assertEqual( 49 | cur_out.shape[:2], 50 | torch.Size([N, backbone.out_channels]) 51 | ) 52 | 53 | 54 | if __name__ == "__main__": 55 | unittest.main() 56 | -------------------------------------------------------------------------------- /plane_mask_detection/tests/test_box_coder.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. 2 | 3 | import unittest 4 | 5 | import numpy as np 6 | import torch 7 | from pytorch_local.maskrcnn_benchmark.modeling.box_coder import BoxCoder 8 | 9 | 10 | class TestBoxCoder(unittest.TestCase): 11 | def test_box_decoder(self): 12 | """ Match unit test UtilsBoxesTest.TestBboxTransformRandom in 13 | caffe2/operators/generate_proposals_op_util_boxes_test.cc 14 | """ 15 | box_coder = BoxCoder(weights=(1.0, 1.0, 1.0, 1.0)) 16 | bbox = torch.from_numpy( 17 | np.array( 18 | [ 19 | 175.62031555, 20 | 20.91103172, 21 | 253.352005, 22 | 155.0145874, 23 | 169.24636841, 24 | 4.85241556, 25 | 228.8605957, 26 | 105.02092743, 27 | 181.77426147, 28 | 199.82876587, 29 | 192.88427734, 30 | 214.0255127, 31 | 174.36262512, 32 | 186.75761414, 33 | 296.19091797, 34 | 231.27906799, 35 | 22.73153877, 36 | 92.02596283, 37 | 135.5695343, 38 | 208.80291748, 39 | ] 40 | ) 41 | .astype(np.float32) 42 | .reshape(-1, 4) 43 | ) 44 | 45 | deltas = torch.from_numpy( 46 | np.array( 47 | [ 48 | 0.47861834, 49 | 0.13992102, 50 | 0.14961673, 51 | 0.71495209, 52 | 0.29915856, 53 | -0.35664671, 54 | 0.89018666, 55 | 0.70815367, 56 | -0.03852064, 57 | 0.44466892, 58 | 0.49492538, 59 | 0.71409376, 60 | 0.28052918, 61 | 0.02184832, 62 | 0.65289006, 63 | 1.05060139, 64 | -0.38172557, 65 | -0.08533806, 66 | -0.60335309, 67 | 0.79052375, 68 | ] 69 | ) 70 | .astype(np.float32) 71 | .reshape(-1, 4) 72 | ) 73 | 74 | gt_bbox = ( 75 | np.array( 76 | [ 77 | 206.949539, 78 | -30.715202, 79 | 297.387665, 80 | 244.448486, 81 | 143.871216, 82 | -83.342888, 83 | 290.502289, 84 | 121.053398, 85 | 177.430283, 86 | 198.666245, 87 | 196.295273, 88 | 228.703079, 89 | 152.251892, 90 | 145.431564, 91 | 387.215454, 92 | 274.594238, 93 | 5.062420, 94 | 11.040955, 95 | 66.328903, 96 | 269.686218, 97 | ] 98 | ) 99 | .astype(np.float32) 100 | .reshape(-1, 4) 101 | ) 102 | 103 | results = box_coder.decode(deltas, bbox) 104 | 105 | np.testing.assert_allclose(results.detach().numpy(), gt_bbox, atol=1e-4) 106 | 107 | 108 | if __name__ == "__main__": 109 | unittest.main() 110 | -------------------------------------------------------------------------------- /plane_mask_detection/tests/test_configs.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. 2 | 3 | import unittest 4 | import glob 5 | import os 6 | import utils 7 | 8 | 9 | class TestConfigs(unittest.TestCase): 10 | def test_configs_load(self): 11 | ''' Make sure configs are loadable ''' 12 | 13 | cfg_root_path = utils.get_config_root_path() 14 | files = glob.glob( 15 | os.path.join(cfg_root_path, "./**/*.yaml"), recursive=True) 16 | self.assertGreater(len(files), 0) 17 | 18 | for fn in files: 19 | print('Loading {}...'.format(fn)) 20 | utils.load_config_from_file(fn) 21 | 22 | 23 | if __name__ == "__main__": 24 | unittest.main() 25 | -------------------------------------------------------------------------------- /plane_mask_detection/tests/test_fbnet.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. 2 | 3 | import unittest 4 | 5 | import numpy as np 6 | import torch 7 | import pytorch_local.maskrcnn_benchmark.modeling.backbone.fbnet_builder as fbnet_builder 8 | 9 | 10 | TEST_CUDA = torch.cuda.is_available() 11 | 12 | 13 | def _test_primitive(self, device, op_name, op_func, N, C_in, C_out, expand, stride): 14 | op = op_func(C_in, C_out, expand, stride).to(device) 15 | input = torch.rand([N, C_in, 7, 7], dtype=torch.float32).to(device) 16 | output = op(input) 17 | self.assertEqual( 18 | output.shape[:2], torch.Size([N, C_out]), 19 | 'Primitive {} failed for shape {}.'.format(op_name, input.shape) 20 | ) 21 | 22 | 23 | class TestFBNetBuilder(unittest.TestCase): 24 | def test_identity(self): 25 | id_op = fbnet_builder.Identity(20, 20, 1) 26 | input = torch.rand([10, 20, 7, 7], dtype=torch.float32) 27 | output = id_op(input) 28 | np.testing.assert_array_equal(np.array(input), np.array(output)) 29 | 30 | id_op = fbnet_builder.Identity(20, 40, 2) 31 | input = torch.rand([10, 20, 7, 7], dtype=torch.float32) 32 | output = id_op(input) 33 | np.testing.assert_array_equal(output.shape, [10, 40, 4, 4]) 34 | 35 | def test_primitives(self): 36 | ''' Make sures the primitives runs ''' 37 | for op_name, op_func in fbnet_builder.PRIMITIVES.items(): 38 | print('Testing {}'.format(op_name)) 39 | 40 | _test_primitive( 41 | self, "cpu", 42 | op_name, op_func, 43 | N=20, C_in=16, C_out=32, expand=4, stride=1 44 | ) 45 | 46 | @unittest.skipIf(not TEST_CUDA, "no CUDA detected") 47 | def test_primitives_cuda(self): 48 | ''' Make sures the primitives runs on cuda ''' 49 | for op_name, op_func in fbnet_builder.PRIMITIVES.items(): 50 | print('Testing {}'.format(op_name)) 51 | 52 | _test_primitive( 53 | self, "cuda", 54 | op_name, op_func, 55 | N=20, C_in=16, C_out=32, expand=4, stride=1 56 | ) 57 | 58 | def test_primitives_empty_batch(self): 59 | ''' Make sures the primitives runs ''' 60 | for op_name, op_func in fbnet_builder.PRIMITIVES.items(): 61 | print('Testing {}'.format(op_name)) 62 | 63 | # test empty batch size 64 | _test_primitive( 65 | self, "cpu", 66 | op_name, op_func, 67 | N=0, C_in=16, C_out=32, expand=4, stride=1 68 | ) 69 | 70 | @unittest.skipIf(not TEST_CUDA, "no CUDA detected") 71 | def test_primitives_cuda_empty_batch(self): 72 | ''' Make sures the primitives runs ''' 73 | for op_name, op_func in fbnet_builder.PRIMITIVES.items(): 74 | print('Testing {}'.format(op_name)) 75 | 76 | # test empty batch size 77 | _test_primitive( 78 | self, "cuda", 79 | op_name, op_func, 80 | N=0, C_in=16, C_out=32, expand=4, stride=1 81 | ) 82 | 83 | if __name__ == "__main__": 84 | unittest.main() 85 | -------------------------------------------------------------------------------- /plane_mask_detection/tests/test_feature_extractors.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. 2 | 3 | import unittest 4 | import copy 5 | import torch 6 | # import modules to to register feature extractors 7 | from pytorch_local.maskrcnn_benchmark.modeling.backbone import build_backbone # NoQA 8 | from pytorch_local.maskrcnn_benchmark.modeling.roi_heads.roi_heads import build_roi_heads # NoQA 9 | from pytorch_local.maskrcnn_benchmark.modeling import registry 10 | from pytorch_local.maskrcnn_benchmark.structures.bounding_box import BoxList 11 | from pytorch_local.maskrcnn_benchmark.config import cfg as g_cfg 12 | from utils import load_config 13 | 14 | # overwrite configs if specified, otherwise default config is used 15 | FEATURE_EXTRACTORS_CFGS = { 16 | } 17 | 18 | # overwrite configs if specified, otherwise default config is used 19 | FEATURE_EXTRACTORS_INPUT_CHANNELS = { 20 | # in_channels was not used, load through config 21 | "ResNet50Conv5ROIFeatureExtractor": 1024, 22 | } 23 | 24 | 25 | def _test_feature_extractors( 26 | self, extractors, overwrite_cfgs, overwrite_in_channels 27 | ): 28 | ''' Make sure roi box feature extractors run ''' 29 | 30 | self.assertGreater(len(extractors), 0) 31 | 32 | in_channels_default = 64 33 | 34 | for name, builder in extractors.items(): 35 | print('Testing {}...'.format(name)) 36 | if name in overwrite_cfgs: 37 | cfg = load_config(overwrite_cfgs[name]) 38 | else: 39 | # Use default config if config file is not specified 40 | cfg = copy.deepcopy(g_cfg) 41 | 42 | in_channels = overwrite_in_channels.get( 43 | name, in_channels_default) 44 | 45 | fe = builder(cfg, in_channels) 46 | self.assertIsNotNone( 47 | getattr(fe, 'out_channels', None), 48 | 'Need to provide out_channels for feature extractor {}'.format(name) 49 | ) 50 | 51 | N, C_in, H, W = 2, in_channels, 24, 32 52 | input = torch.rand([N, C_in, H, W], dtype=torch.float32) 53 | bboxes = [[1, 1, 10, 10], [5, 5, 8, 8], [2, 2, 3, 4]] 54 | img_size = [384, 512] 55 | box_list = BoxList(bboxes, img_size, "xyxy") 56 | out = fe([input], [box_list] * N) 57 | self.assertEqual( 58 | out.shape[:2], 59 | torch.Size([N * len(bboxes), fe.out_channels]) 60 | ) 61 | 62 | 63 | class TestFeatureExtractors(unittest.TestCase): 64 | def test_roi_box_feature_extractors(self): 65 | ''' Make sure roi box feature extractors run ''' 66 | _test_feature_extractors( 67 | self, 68 | registry.ROI_BOX_FEATURE_EXTRACTORS, 69 | FEATURE_EXTRACTORS_CFGS, 70 | FEATURE_EXTRACTORS_INPUT_CHANNELS, 71 | ) 72 | 73 | def test_roi_keypoints_feature_extractors(self): 74 | ''' Make sure roi keypoints feature extractors run ''' 75 | _test_feature_extractors( 76 | self, 77 | registry.ROI_KEYPOINT_FEATURE_EXTRACTORS, 78 | FEATURE_EXTRACTORS_CFGS, 79 | FEATURE_EXTRACTORS_INPUT_CHANNELS, 80 | ) 81 | 82 | def test_roi_mask_feature_extractors(self): 83 | ''' Make sure roi mask feature extractors run ''' 84 | _test_feature_extractors( 85 | self, 86 | registry.ROI_MASK_FEATURE_EXTRACTORS, 87 | FEATURE_EXTRACTORS_CFGS, 88 | FEATURE_EXTRACTORS_INPUT_CHANNELS, 89 | ) 90 | 91 | 92 | if __name__ == "__main__": 93 | unittest.main() 94 | -------------------------------------------------------------------------------- /plane_mask_detection/tests/test_metric_logger.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. 2 | import unittest 3 | 4 | from pytorch_local.maskrcnn_benchmark.utils.metric_logger import MetricLogger 5 | 6 | 7 | class TestMetricLogger(unittest.TestCase): 8 | def test_update(self): 9 | meter = MetricLogger() 10 | for i in range(10): 11 | meter.update(metric=float(i)) 12 | 13 | m = meter.meters["metric"] 14 | self.assertEqual(m.count, 10) 15 | self.assertEqual(m.total, 45) 16 | self.assertEqual(m.median, 4) 17 | self.assertEqual(m.avg, 4.5) 18 | 19 | def test_no_attr(self): 20 | meter = MetricLogger() 21 | _ = meter.meters 22 | _ = meter.delimiter 23 | def broken(): 24 | _ = meter.not_existent 25 | self.assertRaises(AttributeError, broken) 26 | 27 | if __name__ == "__main__": 28 | unittest.main() 29 | -------------------------------------------------------------------------------- /plane_mask_detection/tests/test_predictors.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. 2 | 3 | import unittest 4 | import copy 5 | import torch 6 | # import modules to to register predictors 7 | from pytorch_local.maskrcnn_benchmark.modeling.backbone import build_backbone # NoQA 8 | from pytorch_local.maskrcnn_benchmark.modeling.roi_heads.roi_heads import build_roi_heads # NoQA 9 | from pytorch_local.maskrcnn_benchmark.modeling import registry 10 | from pytorch_local.maskrcnn_benchmark.config import cfg as g_cfg 11 | from utils import load_config 12 | 13 | 14 | # overwrite configs if specified, otherwise default config is used 15 | PREDICTOR_CFGS = { 16 | } 17 | 18 | # overwrite configs if specified, otherwise default config is used 19 | PREDICTOR_INPUT_CHANNELS = { 20 | } 21 | 22 | 23 | def _test_predictors( 24 | self, predictors, overwrite_cfgs, overwrite_in_channels, 25 | hwsize, 26 | ): 27 | ''' Make sure predictors run ''' 28 | 29 | self.assertGreater(len(predictors), 0) 30 | 31 | in_channels_default = 64 32 | 33 | for name, builder in predictors.items(): 34 | print('Testing {}...'.format(name)) 35 | if name in overwrite_cfgs: 36 | cfg = load_config(overwrite_cfgs[name]) 37 | else: 38 | # Use default config if config file is not specified 39 | cfg = copy.deepcopy(g_cfg) 40 | 41 | in_channels = overwrite_in_channels.get( 42 | name, in_channels_default) 43 | 44 | fe = builder(cfg, in_channels) 45 | 46 | N, C_in, H, W = 2, in_channels, hwsize, hwsize 47 | input = torch.rand([N, C_in, H, W], dtype=torch.float32) 48 | out = fe(input) 49 | yield input, out, cfg 50 | 51 | 52 | class TestPredictors(unittest.TestCase): 53 | def test_roi_box_predictors(self): 54 | ''' Make sure roi box predictors run ''' 55 | for cur_in, cur_out, cur_cfg in _test_predictors( 56 | self, 57 | registry.ROI_BOX_PREDICTOR, 58 | PREDICTOR_CFGS, 59 | PREDICTOR_INPUT_CHANNELS, 60 | hwsize=1, 61 | ): 62 | self.assertEqual(len(cur_out), 2) 63 | scores, bbox_deltas = cur_out[0], cur_out[1] 64 | self.assertEqual( 65 | scores.shape[1], cur_cfg.MODEL.ROI_BOX_HEAD.NUM_CLASSES) 66 | self.assertEqual(scores.shape[0], cur_in.shape[0]) 67 | self.assertEqual(scores.shape[0], bbox_deltas.shape[0]) 68 | self.assertEqual(scores.shape[1] * 4, bbox_deltas.shape[1]) 69 | 70 | def test_roi_keypoints_predictors(self): 71 | ''' Make sure roi keypoint predictors run ''' 72 | for cur_in, cur_out, cur_cfg in _test_predictors( 73 | self, 74 | registry.ROI_KEYPOINT_PREDICTOR, 75 | PREDICTOR_CFGS, 76 | PREDICTOR_INPUT_CHANNELS, 77 | hwsize=14, 78 | ): 79 | self.assertEqual(cur_out.shape[0], cur_in.shape[0]) 80 | self.assertEqual( 81 | cur_out.shape[1], cur_cfg.MODEL.ROI_KEYPOINT_HEAD.NUM_CLASSES) 82 | 83 | def test_roi_mask_predictors(self): 84 | ''' Make sure roi mask predictors run ''' 85 | for cur_in, cur_out, cur_cfg in _test_predictors( 86 | self, 87 | registry.ROI_MASK_PREDICTOR, 88 | PREDICTOR_CFGS, 89 | PREDICTOR_INPUT_CHANNELS, 90 | hwsize=14, 91 | ): 92 | self.assertEqual(cur_out.shape[0], cur_in.shape[0]) 93 | self.assertEqual( 94 | cur_out.shape[1], cur_cfg.MODEL.ROI_BOX_HEAD.NUM_CLASSES) 95 | 96 | 97 | if __name__ == "__main__": 98 | unittest.main() 99 | -------------------------------------------------------------------------------- /plane_mask_detection/tests/test_rpn_heads.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. 2 | 3 | import unittest 4 | import copy 5 | import torch 6 | # import modules to to register rpn heads 7 | from pytorch_local.maskrcnn_benchmark.modeling.backbone import build_backbone # NoQA 8 | from pytorch_local.maskrcnn_benchmark.modeling.rpn.rpn import build_rpn # NoQA 9 | from pytorch_local.maskrcnn_benchmark.modeling import registry 10 | from pytorch_local.maskrcnn_benchmark.config import cfg as g_cfg 11 | from utils import load_config 12 | 13 | 14 | # overwrite configs if specified, otherwise default config is used 15 | RPN_CFGS = { 16 | } 17 | 18 | 19 | class TestRPNHeads(unittest.TestCase): 20 | def test_build_rpn_heads(self): 21 | ''' Make sure rpn heads run ''' 22 | 23 | self.assertGreater(len(registry.RPN_HEADS), 0) 24 | 25 | in_channels = 64 26 | num_anchors = 10 27 | 28 | for name, builder in registry.RPN_HEADS.items(): 29 | print('Testing {}...'.format(name)) 30 | if name in RPN_CFGS: 31 | cfg = load_config(RPN_CFGS[name]) 32 | else: 33 | # Use default config if config file is not specified 34 | cfg = copy.deepcopy(g_cfg) 35 | 36 | rpn = builder(cfg, in_channels, num_anchors) 37 | 38 | N, C_in, H, W = 2, in_channels, 24, 32 39 | input = torch.rand([N, C_in, H, W], dtype=torch.float32) 40 | LAYERS = 3 41 | out = rpn([input] * LAYERS) 42 | self.assertEqual(len(out), 2) 43 | logits, bbox_reg = out 44 | for idx in range(LAYERS): 45 | self.assertEqual( 46 | logits[idx].shape, 47 | torch.Size([ 48 | input.shape[0], num_anchors, 49 | input.shape[2], input.shape[3], 50 | ]) 51 | ) 52 | self.assertEqual( 53 | bbox_reg[idx].shape, 54 | torch.Size([ 55 | logits[idx].shape[0], num_anchors * 4, 56 | logits[idx].shape[2], logits[idx].shape[3], 57 | ]), 58 | ) 59 | 60 | 61 | if __name__ == "__main__": 62 | unittest.main() 63 | -------------------------------------------------------------------------------- /plane_mask_detection/tests/test_segmentation_mask.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. 2 | import unittest 3 | import torch 4 | from pytorch_local.maskrcnn_benchmark.structures.segmentation_mask import SegmentationMask 5 | 6 | 7 | class TestSegmentationMask(unittest.TestCase): 8 | def __init__(self, method_name='runTest'): 9 | super(TestSegmentationMask, self).__init__(method_name) 10 | poly = [[[423.0, 306.5, 406.5, 277.0, 400.0, 271.5, 389.5, 277.0, 11 | 387.5, 292.0, 384.5, 295.0, 374.5, 220.0, 378.5, 210.0, 12 | 391.0, 200.5, 404.0, 199.5, 414.0, 203.5, 425.5, 221.0, 13 | 438.5, 297.0, 423.0, 306.5], 14 | [100, 100, 200, 100, 200, 200, 100, 200], 15 | ]] 16 | width = 640 17 | height = 480 18 | size = width, height 19 | 20 | self.P = SegmentationMask(poly, size, 'poly') 21 | self.M = SegmentationMask(poly, size, 'poly').convert('mask') 22 | 23 | def L1(self, A, B): 24 | diff = A.get_mask_tensor() - B.get_mask_tensor() 25 | diff = torch.sum(torch.abs(diff.float())).item() 26 | return diff 27 | 28 | def test_convert(self): 29 | M_hat = self.M.convert('poly').convert('mask') 30 | P_hat = self.P.convert('mask').convert('poly') 31 | 32 | diff_mask = self.L1(self.M, M_hat) 33 | diff_poly = self.L1(self.P, P_hat) 34 | self.assertTrue(diff_mask == diff_poly) 35 | self.assertTrue(diff_mask <= 8169.) 36 | self.assertTrue(diff_poly <= 8169.) 37 | 38 | def test_crop(self): 39 | box = [400, 250, 500, 300] # xyxy 40 | diff = self.L1(self.M.crop(box), self.P.crop(box)) 41 | self.assertTrue(diff <= 1.) 42 | 43 | def test_resize(self): 44 | new_size = 50, 25 45 | M_hat = self.M.resize(new_size) 46 | P_hat = self.P.resize(new_size) 47 | diff = self.L1(M_hat, P_hat) 48 | 49 | self.assertTrue(self.M.size == self.P.size) 50 | self.assertTrue(M_hat.size == P_hat.size) 51 | self.assertTrue(self.M.size != M_hat.size) 52 | self.assertTrue(diff <= 255.) 53 | 54 | def test_transpose(self): 55 | FLIP_LEFT_RIGHT = 0 56 | FLIP_TOP_BOTTOM = 1 57 | diff_hor = self.L1(self.M.transpose(FLIP_LEFT_RIGHT), 58 | self.P.transpose(FLIP_LEFT_RIGHT)) 59 | 60 | diff_ver = self.L1(self.M.transpose(FLIP_TOP_BOTTOM), 61 | self.P.transpose(FLIP_TOP_BOTTOM)) 62 | 63 | self.assertTrue(diff_hor <= 53250.) 64 | self.assertTrue(diff_ver <= 42494.) 65 | 66 | 67 | if __name__ == "__main__": 68 | 69 | unittest.main() 70 | -------------------------------------------------------------------------------- /plane_mask_detection/tests/utils.py: -------------------------------------------------------------------------------- 1 | from __future__ import absolute_import, division, print_function, unicode_literals 2 | 3 | # Set up custom environment before nearly anything else is imported 4 | # NOTE: this should be the first import (no not reorder) 5 | from pytorch_local.maskrcnn_benchmark.utils.env import setup_environment # noqa F401 isort:skip 6 | import env_tests.env as env_tests 7 | 8 | import os 9 | import copy 10 | 11 | from maskrcnn_benchmark.config import cfg as g_cfg 12 | 13 | 14 | def get_config_root_path(): 15 | return env_tests.get_config_root_path() 16 | 17 | 18 | def load_config(rel_path): 19 | ''' Load config from file path specified as path relative to config_root ''' 20 | cfg_path = os.path.join(env_tests.get_config_root_path(), rel_path) 21 | return load_config_from_file(cfg_path) 22 | 23 | 24 | def load_config_from_file(file_path): 25 | ''' Load config from file path specified as absolute path ''' 26 | ret = copy.deepcopy(g_cfg) 27 | ret.merge_from_file(file_path) 28 | return ret 29 | -------------------------------------------------------------------------------- /plane_mask_detection/tools/cityscapes/instances2dict_with_polygons.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/python 2 | # 3 | # Convert instances from png files to a dictionary 4 | # This files is created according to https://github.com/facebookresearch/Detectron/issues/111 5 | 6 | from __future__ import print_function, absolute_import, division 7 | import os, sys 8 | 9 | sys.path.append( os.path.normpath( os.path.join( os.path.dirname( __file__ ) , '..' , 'helpers' ) ) ) 10 | from csHelpers import * 11 | 12 | # Cityscapes imports 13 | from cityscapesscripts.evaluation.instance import * 14 | from cityscapesscripts.helpers.csHelpers import * 15 | import cv2 16 | from maskrcnn_benchmark.utils import cv2_util 17 | 18 | 19 | def instances2dict_with_polygons(imageFileList, verbose=False): 20 | imgCount = 0 21 | instanceDict = {} 22 | 23 | if not isinstance(imageFileList, list): 24 | imageFileList = [imageFileList] 25 | 26 | if verbose: 27 | print("Processing {} images...".format(len(imageFileList))) 28 | 29 | for imageFileName in imageFileList: 30 | # Load image 31 | img = Image.open(imageFileName) 32 | 33 | # Image as numpy array 34 | imgNp = np.array(img) 35 | 36 | # Initialize label categories 37 | instances = {} 38 | for label in labels: 39 | instances[label.name] = [] 40 | 41 | # Loop through all instance ids in instance image 42 | for instanceId in np.unique(imgNp): 43 | if instanceId < 1000: 44 | continue 45 | instanceObj = Instance(imgNp, instanceId) 46 | instanceObj_dict = instanceObj.toDict() 47 | 48 | #instances[id2label[instanceObj.labelID].name].append(instanceObj.toDict()) 49 | if id2label[instanceObj.labelID].hasInstances: 50 | mask = (imgNp == instanceId).astype(np.uint8) 51 | contour, hier = cv2_util.findContours( 52 | mask.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE) 53 | 54 | polygons = [c.reshape(-1).tolist() for c in contour] 55 | instanceObj_dict['contours'] = polygons 56 | 57 | instances[id2label[instanceObj.labelID].name].append(instanceObj_dict) 58 | 59 | imgKey = os.path.abspath(imageFileName) 60 | instanceDict[imgKey] = instances 61 | imgCount += 1 62 | 63 | if verbose: 64 | print("\rImages Processed: {}".format(imgCount), end=' ') 65 | sys.stdout.flush() 66 | 67 | if verbose: 68 | print("") 69 | 70 | return instanceDict 71 | 72 | def main(argv): 73 | fileList = [] 74 | if (len(argv) > 2): 75 | for arg in argv: 76 | if ("png" in arg): 77 | fileList.append(arg) 78 | instances2dict_with_polygons(fileList, True) 79 | 80 | if __name__ == "__main__": 81 | main(sys.argv[1:]) 82 | -------------------------------------------------------------------------------- /plane_mask_detection/tools/test_net.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. 2 | # Set up custom environment before nearly anything else is imported 3 | # NOTE: this should be the first import (no not reorder) 4 | from maskrcnn_benchmark.utils.env import setup_environment # noqa F401 isort:skip 5 | 6 | import argparse 7 | import os 8 | 9 | import torch 10 | from maskrcnn_benchmark.config import cfg 11 | from maskrcnn_benchmark.data import make_data_loader 12 | from maskrcnn_benchmark.engine.inference import inference 13 | from maskrcnn_benchmark.modeling.detector import build_detection_model 14 | from maskrcnn_benchmark.utils.checkpoint import DetectronCheckpointer 15 | from maskrcnn_benchmark.utils.collect_env import collect_env_info 16 | from maskrcnn_benchmark.utils.comm import synchronize, get_rank 17 | from maskrcnn_benchmark.utils.logger import setup_logger 18 | from maskrcnn_benchmark.utils.miscellaneous import mkdir 19 | 20 | # Check if we can enable mixed-precision via apex.amp 21 | try: 22 | from apex import amp 23 | except ImportError: 24 | raise ImportError('Use APEX for mixed precision via apex.amp') 25 | 26 | 27 | def main(): 28 | parser = argparse.ArgumentParser(description="PyTorch Object Detection Inference") 29 | parser.add_argument( 30 | "--config-file", 31 | default="/private/home/fmassa/github/detectron.pytorch_v2/configs/e2e_faster_rcnn_R_50_C4_1x_caffe2.yaml", 32 | metavar="FILE", 33 | help="path to config file", 34 | ) 35 | parser.add_argument("--local_rank", type=int, default=0) 36 | parser.add_argument( 37 | "--ckpt", 38 | help="The path to the checkpoint for test, default is the latest checkpoint.", 39 | default=None, 40 | ) 41 | parser.add_argument( 42 | "opts", 43 | help="Modify config options using the command-line", 44 | default=None, 45 | nargs=argparse.REMAINDER, 46 | ) 47 | 48 | args = parser.parse_args() 49 | 50 | num_gpus = int(os.environ["WORLD_SIZE"]) if "WORLD_SIZE" in os.environ else 1 51 | distributed = num_gpus > 1 52 | 53 | if distributed: 54 | torch.cuda.set_device(args.local_rank) 55 | torch.distributed.init_process_group( 56 | backend="nccl", init_method="env://" 57 | ) 58 | synchronize() 59 | 60 | cfg.merge_from_file(args.config_file) 61 | cfg.merge_from_list(args.opts) 62 | cfg.freeze() 63 | 64 | save_dir = "" 65 | logger = setup_logger("maskrcnn_benchmark", save_dir, get_rank()) 66 | logger.info("Using {} GPUs".format(num_gpus)) 67 | logger.info(cfg) 68 | 69 | logger.info("Collecting env info (might take some time)") 70 | logger.info("\n" + collect_env_info()) 71 | 72 | model = build_detection_model(cfg) 73 | model.to(cfg.MODEL.DEVICE) 74 | 75 | # Initialize mixed-precision if necessary 76 | use_mixed_precision = cfg.DTYPE == 'float16' 77 | amp_handle = amp.init(enabled=use_mixed_precision, verbose=cfg.AMP_VERBOSE) 78 | 79 | output_dir = cfg.OUTPUT_DIR 80 | checkpointer = DetectronCheckpointer(cfg, model, save_dir=output_dir) 81 | ckpt = cfg.MODEL.WEIGHT if args.ckpt is None else args.ckpt 82 | _ = checkpointer.load(ckpt, use_latest=args.ckpt is None) 83 | 84 | iou_types = ("bbox",) 85 | if cfg.MODEL.MASK_ON: 86 | iou_types = iou_types + ("segm",) 87 | if cfg.MODEL.KEYPOINT_ON: 88 | iou_types = iou_types + ("keypoints",) 89 | output_folders = [None] * len(cfg.DATASETS.TEST) 90 | dataset_names = cfg.DATASETS.TEST 91 | if cfg.OUTPUT_DIR: 92 | for idx, dataset_name in enumerate(dataset_names): 93 | output_folder = os.path.join(cfg.OUTPUT_DIR, "inference", dataset_name) 94 | mkdir(output_folder) 95 | output_folders[idx] = output_folder 96 | data_loaders_val = make_data_loader(cfg, is_train=False, is_distributed=distributed) 97 | for output_folder, dataset_name, data_loader_val in zip(output_folders, dataset_names, data_loaders_val): 98 | inference( 99 | model, 100 | data_loader_val, 101 | dataset_name=dataset_name, 102 | iou_types=iou_types, 103 | box_only=False if cfg.MODEL.RETINANET_ON else cfg.MODEL.RPN_ONLY, 104 | device=cfg.MODEL.DEVICE, 105 | expected_results=cfg.TEST.EXPECTED_RESULTS, 106 | expected_results_sigma_tol=cfg.TEST.EXPECTED_RESULTS_SIGMA_TOL, 107 | output_folder=output_folder, 108 | ) 109 | synchronize() 110 | 111 | 112 | if __name__ == "__main__": 113 | main() 114 | --------------------------------------------------------------------------------