def __init__(self, cfg, is_train=True): if cfg.INPUT.CROP.ENABLED and is_train: self.crop_gen = [ T.ResizeShortestEdge([400, 500, 600], sample_style="choice"), T.RandomCrop(cfg.INPUT.CROP.TYPE, cfg.INPUT.CROP.SIZE), ] else: self.crop_gen = None assert not cfg.MODEL.MASK_ON, "Mask is not supported" self.tfm_gens = build_transform_gen(cfg, is_train) logging.getLogger("detectron2").info( "Full TransformGens used in pre-process: {}, crop: {}".format( str(self.tfm_gens), str(self.crop_gen))) self.img_format = cfg.INPUT.FORMAT self.zip_read = cfg.INPUT.ZIP_READ self.is_train = is_train self.max_frame_dist = cfg.MODEL.DETR.MAX_FRAME_DIST self.ignore_same = cfg.MODEL.DETR.IGNORE_SAME # video2img --> 爆内存问题可能是出在这边,每次得释放掉 / 考虑放到builtin试试. if self.is_train: self.dataset_dicts = get_detection_dataset_dicts( cfg.DATASETS.TRAIN, filter_empty=cfg.DATALOADER.FILTER_EMPTY_ANNOTATIONS, min_keypoints=cfg.MODEL.ROI_KEYPOINT_HEAD. MIN_KEYPOINTS_PER_IMAGE if cfg.MODEL.KEYPOINT_ON else 0, proposal_files=cfg.DATASETS.PROPOSAL_FILES_TRAIN if cfg.MODEL.LOAD_PROPOSALS else None, ) assert len(cfg.DATASETS.TRAIN) == 1, logging.\ info("Only support ONE dataset each time, however, " "there are {} datasets now.".format(len(cfg.DATASETS.TRAIN))) self.video_to_images = MetadataCatalog.get( cfg.DATASETS.TRAIN[0]).video_to_images self.image_to_index = MetadataCatalog.get( cfg.DATASETS.TRAIN[0]).image_to_index else: self.dataset_dicts = get_detection_dataset_dicts( cfg.DATASETS.TEST, filter_empty=False, min_keypoints=0, proposal_files=None) assert len(cfg.DATASETS.TEST) == 1, logging. \ info("Only support ONE dataset each time, however, " "there are {} datasets now.".format(len(cfg.DATASETS.TEST))) self.video_to_images = MetadataCatalog.get( cfg.DATASETS.TEST[0]).video_to_images self.image_to_index = MetadataCatalog.get( cfg.DATASETS.TEST[0]).image_to_index
def get_coco_test(batch_size, opts, cfg_file): cfg = setup(opts, cfg_file) dataset_name = 'coco_2017_val_panoptic_separated' dataset_dicts = get_detection_dataset_dicts( [dataset_name], filter_empty=False, proposal_files=None, ) dataset = DatasetFromList(dataset_dicts) mapper = DatasetMapper(cfg, False) dataset = MapDataset(dataset, mapper) sampler = samplers.InferenceSampler(len(dataset)) batch_sampler = torch.utils.data.sampler.BatchSampler(sampler, batch_size, drop_last=False) data_loader = torch.utils.data.DataLoader( dataset, num_workers=4, batch_sampler=batch_sampler, collate_fn=trivial_batch_collator, ) return data_loader
def build_detection_test_loader(cfg, dataset_name, batch_size, mapper=None): dataset_dicts = get_detection_dataset_dicts( [dataset_name], filter_empty=False, proposal_files=[ cfg.DATASETS.PROPOSAL_FILES_TEST[list( cfg.DATASETS.TEST).index(dataset_name)] ] if cfg.MODEL.LOAD_PROPOSALS else None, ) dataset = DatasetFromList(dataset_dicts) if mapper is None: mapper = DatasetMapper(cfg, False) dataset = MapDataset(dataset, mapper) sampler = samplers.InferenceSampler(len(dataset)) # Always use 1 image per worker during inference since this is the # standard when reporting inference time in papers. batch_sampler = torch.utils.data.sampler.BatchSampler(sampler, batch_size, drop_last=False) data_loader = torch.utils.data.DataLoader( dataset, num_workers=cfg.DATALOADER.NUM_WORKERS, batch_sampler=batch_sampler, collate_fn=_trivial_batch_collator, ) return data_loader
def get_coco_train(batch_size, opts, cfg_file): cfg = setup(opts, cfg_file) dataset_dicts = get_detection_dataset_dicts( cfg.DATASETS.TRAIN, filter_empty=cfg.DATALOADER.FILTER_EMPTY_ANNOTATIONS, min_keypoints=cfg.MODEL.ROI_KEYPOINT_HEAD.MIN_KEYPOINTS_PER_IMAGE if cfg.MODEL.KEYPOINT_ON else 0, proposal_files=cfg.DATASETS.PROPOSAL_FILES_TRAIN if cfg.MODEL.LOAD_PROPOSALS else None, ) dataset = DatasetFromList(dataset_dicts, copy=False) mapper = DatasetMapper(cfg, True) dataset = MapDataset(dataset, mapper) sampler = samplers.TrainingSampler(len(dataset)) batch_sampler = torch.utils.data.sampler.BatchSampler(sampler, batch_size, drop_last=True) # drop_last so the batch always have the same size train_loader = torch.utils.data.DataLoader( dataset, num_workers=4, batch_sampler=batch_sampler, collate_fn=trivial_batch_collator, worker_init_fn=worker_init_reset_seed, ) return train_loader
def __init__(self, dataset_dir=None, model_pth_path=None, score_threshold=0.15, overlap_threshold=0.95, model_config=None) -> None: self.dataset_dir = dataset_dir self.score_threshold = score_threshold self.overlap_threshold = overlap_threshold if model_config is None: model_config = "COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml" dataset_name = Path(self.dataset_dir).stem self.subject_queue = queue.PriorityQueue(3) self.left_object_queue = queue.PriorityQueue(3) self.right_object_queue = queue.PriorityQueue(3) self.right_interact_queue = queue.PriorityQueue(3) self.left_interact_queue = queue.PriorityQueue(3) self.subject_instance_name = 'Mouse' self.left_object_name = 'LeftTeaball' self.right_object_name = 'RightTeaball' self.left_interact_name = 'LeftInteract' self.right_interact_name = 'RightInteract' self.tracking_results = [] try: register_coco_instances( f"{dataset_name}_train", {}, f"{self.dataset_dir}/train/annotations.json", f"{self.dataset_dir}/train/") register_coco_instances( f"{dataset_name}_valid", {}, f"{self.dataset_dir}/valid/annotations.json", f"{self.dataset_dir}/valid/") except AssertionError as e: print(e) dataset_dicts = get_detection_dataset_dicts([f"{dataset_name}_train"]) _dataset_metadata = MetadataCatalog.get(f"{dataset_name}_train") _dataset_metadata.thing_colors = [ cc['color'] for cc in builtin_meta.COCO_CATEGORIES ] num_classes = len(_dataset_metadata.thing_classes) self.class_names = _dataset_metadata.thing_classes self.cfg = get_cfg() # load model config and pretrained model self.cfg.merge_from_file(model_zoo.get_config_file(model_config)) self.cfg.MODEL.WEIGHTS = model_pth_path self.cfg.DATASETS.TRAIN = (f"{dataset_name}_train", ) self.cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = self.score_threshold self.cfg.MODEL.DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu' self.cfg.MODEL.ROI_HEADS.NUM_CLASSES = num_classes self.cfg.MODEL.ROI_HEADS.NMS_THRESH_TEST = self.overlap_threshold # NMS threshold used on RPN proposals self.cfg.MODEL.RPN.NMS_THRESH = self.overlap_threshold self.predictor = DefaultPredictor(self.cfg)
def build_gdrn_test_loader(cfg, dataset_name, train_objs=None): """Similar to `build_detection_train_loader`. But this function uses the given `dataset_name` argument (instead of the names in cfg), and uses batch size 1. Args: cfg: a detectron2 CfgNode dataset_name (str): a name of the dataset that's available in the DatasetCatalog Returns: DataLoader: a torch DataLoader, that loads the given detection dataset, with test-time transformation and batching. """ dataset_dicts = get_detection_dataset_dicts( [dataset_name], filter_empty=False, proposal_files=[cfg.DATASETS.PROPOSAL_FILES_TEST[list(cfg.DATASETS.TEST).index(dataset_name)]] if cfg.MODEL.LOAD_PROPOSALS else None, ) # load test detection results if cfg.MODEL.LOAD_DETS_TEST: det_files = cfg.DATASETS.DET_FILES_TEST assert len(cfg.DATASETS.TEST) == len(det_files) load_detections_into_dataset( dataset_name, dataset_dicts, det_file=det_files[cfg.DATASETS.TEST.index(dataset_name)], top_k_per_obj=cfg.DATASETS.DET_TOPK_PER_OBJ, score_thr=cfg.DATASETS.DET_THR, train_objs=train_objs, ) if cfg.DATALOADER.FILTER_EMPTY_DETS: dataset_dicts = filter_empty_dets(dataset_dicts) dataset = GDRN_DatasetFromList(cfg, split="test", lst=dataset_dicts, flatten=False) sampler = InferenceSampler(len(dataset)) # Always use 1 image per worker during inference since this is the # standard when reporting inference time in papers. batch_sampler = torch.utils.data.sampler.BatchSampler(sampler, 1, drop_last=False) num_workers = cfg.DATALOADER.NUM_WORKERS # Horovod: limit # of CPU threads to be used per worker. # if num_workers > 0: # torch.set_num_threads(num_workers) kwargs = {"num_workers": num_workers} # When supported, use 'forkserver' to spawn dataloader workers instead of 'fork' to prevent # issues with Infiniband implementations that are not fork-safe # https://github.com/horovod/horovod/blob/master/examples/pytorch/pytorch_imagenet_resnet50.py # if (num_workers > 0 and hasattr(mp, '_supports_context') and # mp._supports_context and 'forkserver' in mp.get_all_start_methods()): # kwargs['multiprocessing_context'] = 'forkserver' data_loader = torch.utils.data.DataLoader( dataset, batch_sampler=batch_sampler, collate_fn=trivial_batch_collator, **kwargs ) return data_loader
def build_gdrn_train_loader(cfg, dataset_names): """A data loader is created by the following steps: 1. Use the dataset names in config to query :class:`DatasetCatalog`, and obtain a list of dicts. 2. Coordinate a random shuffle order shared among all processes (all GPUs) 3. Each process spawn another few workers to process the dicts. Each worker will: * Map each metadata dict into another format to be consumed by the model. * Batch them by simply putting dicts into a list. The batched ``list[mapped_dict]`` is what this dataloader will yield. Args: cfg (CfgNode): the config Returns: an infinite iterator of training data """ dataset_dicts = get_detection_dataset_dicts( dataset_names, filter_empty=cfg.DATALOADER.FILTER_EMPTY_ANNOTATIONS, min_keypoints=cfg.MODEL.ROI_KEYPOINT_HEAD.MIN_KEYPOINTS_PER_IMAGE if cfg.MODEL.KEYPOINT_ON else 0, proposal_files=cfg.DATASETS.PROPOSAL_FILES_TRAIN if cfg.MODEL.LOAD_PROPOSALS else None, ) dataset_dicts = filter_invalid_in_dataset_dicts( dataset_dicts, visib_thr=cfg.DATALOADER.FILTER_VISIB_THR) dataset = GDRN_DatasetFromList(cfg, split="train", lst=dataset_dicts, copy=False) sampler_name = cfg.DATALOADER.SAMPLER_TRAIN logger = logging.getLogger(__name__) logger.info("Using training sampler {}".format(sampler_name)) # TODO avoid if-else? if sampler_name == "TrainingSampler": sampler = TrainingSampler(len(dataset)) elif sampler_name == "RepeatFactorTrainingSampler": repeat_factors = RepeatFactorTrainingSampler.repeat_factors_from_category_frequency( dataset_dicts, cfg.DATALOADER.REPEAT_THRESHOLD) sampler = RepeatFactorTrainingSampler(repeat_factors) else: raise ValueError("Unknown training sampler: {}".format(sampler_name)) return my_build_batch_data_loader( dataset, sampler, cfg.SOLVER.IMS_PER_BATCH, aspect_ratio_grouping=cfg.DATALOADER.ASPECT_RATIO_GROUPING, num_workers=cfg.DATALOADER.NUM_WORKERS, )
def build_simple_dataloader(dataset_name: list, batch_size): dataset_dicts = get_detection_dataset_dicts(dataset_name) dataset = DatasetFromList(dataset_dicts, copy=False) cfg = get_cfg() cfg["aug_kwargs"] = {} dataset = MapDataset(dataset, AlbumentationsMapper(cfg, False)) # set the shuffle to False in debugging mode sampler = TrainingSampler(len(dataset), shuffle=False, seed=42) dataloader = build_batch_data_loader(dataset=dataset, sampler=sampler, total_batch_size=batch_size) return dataloader
def build_detection_train_loader_with_train_sampler(cfg, mapper, seed=42, shuffle=True): dataset_dicts = get_detection_dataset_dicts(cfg.DATASETS.TRAIN) dataset = DatasetFromList(dataset_dicts, copy=False) dataset = MapDataset(dataset, mapper) logger = logging.getLogger(__name__) logger.info("Using training sampler TrainingSampler with shuffle=False") sampler = TrainingSampler(len(dataset), shuffle=shuffle, seed=seed) return build_batch_data_loader( dataset, sampler, cfg.SOLVER.IMS_PER_BATCH, aspect_ratio_grouping=cfg.DATALOADER.ASPECT_RATIO_GROUPING, num_workers=cfg.DATALOADER.NUM_WORKERS, )
def build_weighted_detection_train_loader(cfg: CfgNode, mapper=None): dataset_repeat_factors = get_train_datasets_repeat_factors(cfg) # OrderedDict to guarantee order of values() consistent with repeat factors dataset_name_to_dicts = OrderedDict( { name: get_detection_dataset_dicts( [name], filter_empty=cfg.DATALOADER.FILTER_EMPTY_ANNOTATIONS, min_keypoints=cfg.MODEL.ROI_KEYPOINT_HEAD.MIN_KEYPOINTS_PER_IMAGE if cfg.MODEL.KEYPOINT_ON else 0, proposal_files=cfg.DATASETS.PROPOSAL_FILES_TRAIN if cfg.MODEL.LOAD_PROPOSALS else None, ) for name in cfg.DATASETS.TRAIN } ) # Repeat factor for every sample in the dataset repeat_factors = [ [dataset_repeat_factors[dsname]] * len(dataset_name_to_dicts[dsname]) for dsname in cfg.DATASETS.TRAIN ] repeat_factors = list(itertools.chain.from_iterable(repeat_factors)) dataset_dicts = dataset_name_to_dicts.values() dataset_dicts = list(itertools.chain.from_iterable(dataset_dicts)) dataset = DatasetFromList(dataset_dicts, copy=False) if mapper is None: mapper = DatasetMapper(cfg, True) dataset = MapDataset(dataset, mapper) logger.info( "Using WeightedTrainingSampler with repeat_factors={}".format( cfg.DATASETS.TRAIN_REPEAT_FACTOR ) ) sampler = RepeatFactorTrainingSampler(torch.tensor(repeat_factors)) return build_batch_data_loader( dataset, sampler, cfg.SOLVER.IMS_PER_BATCH, aspect_ratio_grouping=cfg.DATALOADER.ASPECT_RATIO_GROUPING, num_workers=cfg.DATALOADER.NUM_WORKERS, )
def build_detection_test_loader(cfg, dataset_name, batch_size, mapper=None): """ Similar to `build_detection_train_loader`. But this function uses the given `dataset_name` argument (instead of the names in cfg), and uses batch size 1. Args: cfg: a detectron2 CfgNode dataset_name (str): a name of the dataset that's available in the DatasetCatalog mapper (callable): a callable which takes a sample (dict) from dataset and returns the format to be consumed by the model. By default it will be `DatasetMapper(cfg, False)`. Returns: DataLoader: a torch DataLoader, that loads the given detection dataset, with test-time transformation and batching. """ dataset_dicts = get_detection_dataset_dicts( [dataset_name], filter_empty=False, proposal_files=[ cfg.DATASETS.PROPOSAL_FILES_TEST[list( cfg.DATASETS.TEST).index(dataset_name)] ] if cfg.MODEL.LOAD_PROPOSALS else None, ) dataset = DatasetFromList(dataset_dicts) if mapper is None: mapper = DatasetMapper(cfg, False) dataset = MapDataset(dataset, mapper) sampler = samplers.InferenceSampler(len(dataset)) # Always use 1 image per worker during inference since this is the # standard when reporting inference time in papers. batch_sampler = torch.utils.data.sampler.BatchSampler(sampler, batch_size, drop_last=False) data_loader = torch.utils.data.DataLoader( dataset, num_workers=cfg.DATALOADER.NUM_WORKERS, batch_sampler=batch_sampler, collate_fn=trivial_batch_collator, ) return data_loader
def main(args): cfg = setup(args) predictor = DefaultPredictor(cfg) if args.img != '': show_prediction(args.img, predictor, args.scale) elif args.dir != '': files = [] for ext in ['/*.jpg', '/*.png']: files.extend(glob.glob(args.dir + ext)) for file in files: key = show_prediction(file, predictor, args.scale) if key == ord('q'): cv2.destroyAllWindows() break else: dataset_dicts = get_detection_dataset_dicts(['bizcard_val']) for d in random.sample(dataset_dicts, 300): key = show_prediction(d['file_name'], predictor, args.scale) if key == ord('q'): cv2.destroyAllWindows() break
def build_train_loader(cfg, mapper=None): if mapper is None: mapper = get_dataset_mapper(cfg.DATASETS.TRAIN[0])(cfg, True) dataset_dicts = get_detection_dataset_dicts( cfg.DATASETS.TRAIN, filter_empty=cfg.DATALOADER.FILTER_EMPTY_ANNOTATIONS, min_keypoints=cfg.MODEL.ROI_KEYPOINT_HEAD.MIN_KEYPOINTS_PER_IMAGE if cfg.MODEL.KEYPOINT_ON else 0, proposal_files=cfg.DATASETS.PROPOSAL_FILES_TRAIN if cfg.MODEL.LOAD_PROPOSALS else None, ) dataset = DatasetFromList(dataset_dicts, copy=False) dataset = MapDataset(dataset, mapper) sampler_name = cfg.DATALOADER.SAMPLER_TRAIN logger = logging.getLogger(__name__) logger.info("Using training sampler {}".format(sampler_name)) # TODO avoid if-else? if sampler_name == "TrainingSampler": sampler = TrainingSampler(len(dataset)) elif sampler_name == "RepeatFactorTrainingSampler": repeat_factors = RepeatFactorTrainingSampler.repeat_factors_from_category_frequency( dataset_dicts, cfg.DATALOADER.REPEAT_THRESHOLD) sampler = RepeatFactorTrainingSampler(repeat_factors) elif sampler_name == "RatioFactorTrainingSampler": repeat_factors = repeat_factors_from_ratios(dataset_dicts) sampler = RepeatFactorTrainingSampler(repeat_factors) else: raise ValueError("Unknown training sampler: {}".format(sampler_name)) return build_batch_data_loader( dataset, sampler, cfg.SOLVER.IMS_PER_BATCH, aspect_ratio_grouping=cfg.DATALOADER.ASPECT_RATIO_GROUPING, num_workers=cfg.DATALOADER.NUM_WORKERS, )
# # In[80]: # In[81]: register_coco_instances(f"{DATASET_NAME}_train", {}, f"{DATASET_DIR}/train/annotations.json", f"{DATASET_DIR}/train/") register_coco_instances(f"{DATASET_NAME}_valid", {}, f"{DATASET_DIR}/valid/annotations.json", f"{DATASET_DIR}/valid/") # In[25]: dataset_dicts = get_detection_dataset_dicts([f"{DATASET_NAME}_train"]) # In[26]: _dataset_metadata = MetadataCatalog.get(f"{DATASET_NAME}_train") _dataset_metadata.thing_colors = [ cc['color'] for cc in builtin_meta.COCO_CATEGORIES ] # In[27]: _dataset_metadata # In[28]: NUM_CLASSES = len(_dataset_metadata.thing_classes)
def build_detection_train_loader(cfg, mapper=None): """ A data loader is created by the following steps: 1. Use the dataset names in config to query :class:`DatasetCatalog`, and obtain a list of dicts. 2. Coordinate a random shuffle order shared among all processes (all GPUs) 3. Each process spawn another few workers to process the dicts. Each worker will: * Map each metadata dict into another format to be consumed by the model. * Batch them by simply putting dicts into a list. The batched ``list[mapped_dict]`` is what this dataloader will yield. Args: cfg (CfgNode): the config mapper (callable): a callable which takes a sample (dict) from dataset and returns the format to be consumed by the model. By default it will be ``DatasetMapper(cfg, True)``. Returns: an infinite iterator of training data """ num_workers = get_world_size() images_per_batch = cfg.SOLVER.IMS_PER_BATCH assert ( images_per_batch % num_workers == 0 ), "SOLVER.IMS_PER_BATCH ({}) must be divisible by the number of workers ({}).".format( images_per_batch, num_workers) assert ( images_per_batch >= num_workers ), "SOLVER.IMS_PER_BATCH ({}) must be larger than the number of workers ({}).".format( images_per_batch, num_workers) images_per_worker = images_per_batch // num_workers dataset_dicts = get_detection_dataset_dicts( cfg.DATASETS.TRAIN, filter_empty=cfg.DATALOADER.FILTER_EMPTY_ANNOTATIONS, min_keypoints=cfg.MODEL.ROI_KEYPOINT_HEAD.MIN_KEYPOINTS_PER_IMAGE if cfg.MODEL.KEYPOINT_ON else 0, proposal_files=cfg.DATASETS.PROPOSAL_FILES_TRAIN if cfg.MODEL.LOAD_PROPOSALS else None, ) dataset = DatasetFromList(dataset_dicts, copy=False) if mapper is None: mapper = DatasetMapper(cfg, True) dataset = MapDataset(dataset, mapper) # d=dataset[0] # print(d) # print(d['image'].size()) sampler_name = cfg.DATALOADER.SAMPLER_TRAIN logger = logging.getLogger(__name__) logger.info("Using training sampler {}".format(sampler_name)) # TODO avoid if-else? if sampler_name == "TrainingSampler": sampler = TrainingSampler(len(dataset)) elif sampler_name == "RepeatFactorTrainingSampler": repeat_factors = ( RepeatFactorTrainingSampler.repeat_factors_from_category_frequency( dataset_dicts, cfg.DATALOADER.REPEAT_THRESHOLD)) sampler = RepeatFactorTrainingSampler(repeat_factors) elif sampler_name == "PairTrainingSampler": sampler = PairTrainingSampler(cfg, dataset_dicts, images_per_worker) data_loader = torch.utils.data.DataLoader( dataset, num_workers=cfg.DATALOADER.NUM_WORKERS, batch_sampler=sampler, collate_fn=trivial_batch_collator, worker_init_fn=worker_init_reset_seed, ) return PairDataLoader(cfg, data_loader) else: raise ValueError("Unknown training sampler: {}".format(sampler_name)) return build_batch_data_loader( dataset, sampler, cfg.SOLVER.IMS_PER_BATCH, aspect_ratio_grouping=cfg.DATALOADER.ASPECT_RATIO_GROUPING, num_workers=cfg.DATALOADER.NUM_WORKERS, )
cfg, args ) # if you don't like any of the default setup, write your own setup code # if 'build_model(cfg)': # meta_arch = cfg.MODEL.META_ARCHITECTURE # model = META_ARCH_REGISTRY.get(meta_arch)(cfg) # model.to(torch.device(cfg.MODEL.DEVICE)) if 'do-train': ... if 'build_detection_train_loader': # all dataset_dicts w.r.t cfg.DATASETS.TRAIN will be flattened. dataset_dicts = get_detection_dataset_dicts( cfg.DATASETS.TRAIN, filter_empty=cfg.DATALOADER.FILTER_EMPTY_ANNOTATIONS, min_keypoints=cfg.MODEL.ROI_KEYPOINT_HEAD.MIN_KEYPOINTS_PER_IMAGE if cfg.MODEL.KEYPOINT_ON else 0, proposal_files=cfg.DATASETS.PROPOSAL_FILES_TRAIN if cfg.MODEL.LOAD_PROPOSALS else None, ) dataset = DatasetFromList(dataset_dicts, copy=False) dataset = MapDataset(dataset, DatasetMapper(cfg, True)) sampler_name = cfg.DATALOADER.SAMPLER_TRAIN logger = logging.getLogger(__name__) logger.info("Using training sampler {}".format(sampler_name)) # TODO avoid if-else? if sampler_name == "TrainingSampler": sampler = TrainingSampler(len(dataset)) else: raise ValueError("Unknown training sampler: {}".format(sampler_name)) data_loader = build_batch_data_loader(
def __init__(self, dataset_dir=None, out_put_dir=None, score_threshold=0.15, overlap_threshold=0.7, max_iterations=3000, batch_size=8, model_config=None) -> None: self.dataset_dir = dataset_dir self.batch_size = batch_size self.logger = logging.getLogger(__name__) if model_config is None: model_config = "COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml" if out_put_dir is None: self.out_put_dir = str( Path(self.dataset_dir).parent / 'Annolid_training_outputs') else: self.out_put_dir = out_put_dir self.score_threshold = score_threshold self.overlap_threshold = overlap_threshold self.dataset_name = Path(self.dataset_dir).stem try: register_coco_instances( f"{self.dataset_name}_train", {}, f"{self.dataset_dir}/train/annotations.json", f"{self.dataset_dir}/train/") register_coco_instances( f"{self.dataset_name}_valid", {}, f"{self.dataset_dir}/valid/annotations.json", f"{self.dataset_dir}/valid/") except AssertionError as e: self.logger.info(e) dataset_dicts = get_detection_dataset_dicts( [f"{self.dataset_name}_train"]) _dataset_metadata = MetadataCatalog.get(f"{self.dataset_name}_train") _dataset_metadata.thing_colors = [ cc['color'] for cc in builtin_meta.COCO_CATEGORIES ] num_classes = len(_dataset_metadata.thing_classes) self.class_names = _dataset_metadata.thing_classes self.cfg = get_cfg() # load model config and pretrained model self.cfg.merge_from_file(model_zoo.get_config_file(model_config)) self.cfg.DATASETS.TRAIN = (f"{self.dataset_name}_train", ) self.cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = self.score_threshold self.cfg.MODEL.DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu' self.cfg.MODEL.ROI_HEADS.NUM_CLASSES = num_classes self.cfg.MODEL.ROI_HEADS.NMS_THRESH_TEST = self.overlap_threshold # NMS threshold used on RPN proposals self.cfg.MODEL.RPN.NMS_THRESH = self.overlap_threshold self.cfg.DATASETS.TEST = () self.cfg.DATALOADER.NUM_WORKERS = 2 # @param self.cfg.DATALOADER.SAMPLER_TRAIN = "RepeatFactorTrainingSampler" self.cfg.DATALOADER.REPEAT_THRESHOLD = 0.3 self.cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url( model_config) # Let training initialize from model zoo self.cfg.SOLVER.IMS_PER_BATCH = self.batch_size # @param self.cfg.SOLVER.BASE_LR = 0.0025 # @param # pick a good LR self.logger.info(f"Max iterations {max_iterations}") self.logger.info(f"Batch size is: {batch_size}") self.logger.info(f"Dataset dir is : {dataset_dir}") self.logger.info(f"Model config file is : {model_config}") # @param # 3000 iterations seems good enough for 100 frames dataset; # you will need to train longer for a practical dataset self.cfg.SOLVER.MAX_ITER = max_iterations self.cfg.SOLVER.CHECKPOINT_PERIOD = 1000 # @param # @param # faster, and good enough for this toy dataset (default: 512) self.cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 128 # (see https://detectron2.readthedocs.io/tutorials/datasets.html#update-the-config-for-new-datasets) self.cfg.OUTPUT_DIR = self.out_put_dir os.makedirs(self.cfg.OUTPUT_DIR, exist_ok=True) self.trainer = DefaultTrainer(self.cfg) self.trainer.resume_or_load(resume=True)
VizUtil.viz_mask(mask), mask], axis=1) cv2.imshow('view', view.astype(np.uint8)) key = cv2.waitKey(0) if key == ord('q'): cv2.destroyAllWindows() return key if __name__ == '__main__': parser = argparse.ArgumentParser() parser.add_argument('--coco', action='store_true') args = parser.parse_args() if args.coco: for d in random.sample(get_detection_dataset_dicts(['bizcard_train']), 300): img = cv2.imread(d["file_name"]) visualizer = Visualizer( img[:, :, ::-1], metadata=MetadataCatalog.get('bizcard_train'), scale=0.2) vis = visualizer.draw_dataset_dict(d) cv2.imshow('', vis.get_image()[:, :, ::-1]) key = cv2.waitKey(0) if key == ord('q'): cv2.destroyAllWindows() break else: view_gt_dir( Path('/data/training/business_card/input/source_images'),
segmentations_dilated_flat, bbox_dilated_flat = process_mask_and_bbox(multi_poly_dilated_flat) segmentations_dilated_square, bbox_dilated_square = process_mask_and_bbox(multi_poly_dilated_square) save_file_with_new_masks(data, ann, img_id, segmentations_scaled, bbox_scaled, 'coco_dataset_train_visualizer_scale.json') save_file_with_new_masks(data, ann, img_id, segmentations_dilated_round, bbox_dilated_round, 'coco_dataset_train_visualizer_dilated_round.json') save_file_with_new_masks(data, ann, img_id, segmentations_dilated_flat, bbox_dilated_flat, 'coco_dataset_train_visualizer_dilated_flat.json') save_file_with_new_masks(data, ann, img_id, segmentations_dilated_square, bbox_dilated_square, 'coco_dataset_train_visualizer_dilated_square.json') root = Tk() root.title("Image Viewer") register_coco_instances("UCSP Fondecyt Dataset train", {}, "coco_dataset_train.json", "blurred_images") register_coco_instances("UCSP Fondecyt Dataset train 2", {}, "coco_dataset_train.json", "new_images") dataset_dicts = get_detection_dataset_dicts(["UCSP Fondecyt Dataset train"]) dataset_dicts2 = get_detection_dataset_dicts(["UCSP Fondecyt Dataset train 2"]) dataset_metadata = MetadataCatalog.get("UCSP Fondecyt Dataset train") print(dataset_dicts[img_id-1]["file_name"]) my_img0 = ImageTk.PhotoImage(Image.open(dataset_dicts2[img_id-1]["file_name"][:-4] + "_w.png")) my_img1 = ImageTk.PhotoImage(Image.open(dataset_dicts[img_id-1]["file_name"])) #img = cv2.imread(dataset_dicts[img_id-1]["file_name"]) #visualizer = Visualizer(img[:, :, ::-1], metadata=dataset_metadata) #vis = visualizer.draw_dataset_dict(dataset_dicts[img_id-1]) #my_img2 = ImageTk.PhotoImage(Image.fromarray(vis.get_image())) img_orig = Image.open(dataset_dicts[img_id-1]["file_name"]) draw = ImageDraw.Draw(img_orig) for list_coords in dataset_dicts[img_id-1]["annotations"][0]["segmentation"]: draw.polygon(list_coords) my_img2 = ImageTk.PhotoImage(img_orig)