def main(): charts_cfg = cfg.clone() tables_cfg = cfg.clone() charts_cfg.merge_from_file(charts_path) tables_cfg.merge_from_file(tables_path) model_charts = ChataDemo(charts_cfg, "charts") model_tables = ChataDemo(tables_cfg, "tables") client = ApiClient() client.login() newest_ids = client.get_new_publications() for pub_id in newest_ids: pages = client.get_pages(pub_id) print(pub_id) for id, url in pages: print(id) print(url) resp = urllib.request.urlopen(url) image = np.asarray(bytearray(resp.read()), dtype="uint8") image = cv2.imdecode(image, cv2.IMREAD_COLOR) predictions_charts = model_charts.get_predictions(image) predictions_tables = model_tables.get_predictions(image) annotations_charts = annotate(predictions_charts) annotations_tables = annotate(predictions_tables) annotations = annotations_charts+annotations_tables print(annotations) if annotations: client.add_annotation(id, annotations)
def __init__(self, model_path, config_file, categories, show_mask_heatmaps=False, iou_thr=0.5, score_thr=0, device=None): super(MbDetector, self).__init__(categories=categories, iou_thr=iou_thr, score_thr=score_thr, device=device) assert is_file(model_path), "model path does not exist!" assert is_file(config_file), "config path does not exist!" self.model_path = model_path self.config_file = config_file cfg.merge_from_file(config_file) self.cfg = cfg.clone() self.cpu_device = torch.device("cpu") mask_threshold = -1 if show_mask_heatmaps else 0.5 self.masker = Masker(threshold=mask_threshold, padding=1) self.show_mask_heatmaps = show_mask_heatmaps self._build_detector()
def __init__(self, cfg, confidence_threshold=-1, show_mask_heatmaps=False, show_mask_montage=True, masks_per_dim=2, categories=None): self.categories = categories self.cfg = cfg.clone() mask_threshold = -1 if show_mask_montage else 0.5 self.masker = Masker(threshold=mask_threshold, padding=1) if confidence_threshold >= 0: self.confidence_threshold = confidence_threshold else: self.confidence_threshold = cfg.MODEL.ROI_HEADS.SCORE_THRESH_VISUALIZATION self.show_heatmap = show_mask_heatmaps self.show_mask_montage = show_mask_montage self.masks_per_dim = masks_per_dim # used to make colors for each class self.palette = torch.tensor([2**25 - 1, 2**15 - 1, 2**21 - 1]) self.my_palette = [] colors = create_palette() self.my_palette = [] for color in colors: self.my_palette.append(color['rgb'])
def __init__(self): self.outputdir = cfg.outputdir + '/maskrcnn' os.system('mkdir -p ' + self.outputdir) confidence_threshold = 0.8 show_mask_heatmaps = False masks_per_dim = 1 min_image_size = 224 mcfg.merge_from_file(cfg.maskrcnnpath) mcfg.freeze() self.cfg = mcfg.clone() self.model = build_detection_model(mcfg) self.model.eval() self.device = torch.device(mcfg.MODEL.DEVICE) self.model.to(self.device) self.min_image_size = min_image_size checkpointer = DetectronCheckpointer(mcfg, self.model) _ = checkpointer.load(mcfg.MODEL.WEIGHT) self.transforms = self.build_transform() mask_threshold = -1 if show_mask_heatmaps else 0.4 self.masker = Masker(threshold=mask_threshold, padding=100) # used to make colors for each class self.palette = torch.tensor([2**25 - 1, 2**15 - 1, 2**21 - 1]) self.cpu_device = torch.device("cpu") self.confidence_threshold = confidence_threshold self.show_mask_heatmaps = show_mask_heatmaps self.masks_per_dim = masks_per_dim
def __init__(self, cfg): super(REID_Baseline, self).__init__() self.cfg = cfg.clone() self.base = ResNet(cfg.REID.MODEL.LAST_STRIDE) self.base.load_param(cfg.REID.MODEL.PRETRAIN_PATH) self.gap = nn.AdaptiveAvgPool2d(1) self.num_classes = cfg.REID.DATASETS.ID_NUM # if cfg.MODEL.MASK_ON: # self.in_planes = 4096 # else: # self.in_planes = 2048 self.in_planes = 2048 self.reid_feat_dim = 256 self.bottleneck = nn.BatchNorm1d(self.in_planes) self.bottleneck.bias.requires_grad_(False) # no shift self.classifier = nn.Linear(self.in_planes, self.num_classes, bias=False) self.bottleneck.apply(weights_init_kaiming) self.classifier.apply(weights_init_classifier) self.queue_size = 5000 self.margin = self.cfg.TRIPLET.MARGIN self.lut_momentum = 0.0 ### self.reid_feat_dim = self.in_planes # self.register_buffer('lut', torch.zeros(self.num_pid, self.reid_feat_dim).cuda()) self.register_buffer('lut1', torch.zeros(self.num_classes, self.reid_feat_dim).cuda()) self.register_buffer('lut2', torch.zeros(self.num_classes, self.reid_feat_dim).cuda()) self.register_buffer('queue', torch.zeros(self.queue_size, self.reid_feat_dim).cuda())
def __init__(self, cfg, confidence_threshold=0.7, show_mask_heatmaps=False, masks_per_dim=2, min_image_size=224): self.cfg = cfg.clone() self.model = build_detection_model(cfg) self.model.eval() self.device = torch.device(cfg.MODEL.DEVICE) self.model.to(self.device) for n, m in self.model.named_modules(): if n == "roi_heads": m.register_forward_hook(hook) self.min_image_size = min_image_size save_dir = cfg.OUTPUT_DIR checkpointer = DetectronCheckpointer(cfg, self.model, save_dir=save_dir) _ = checkpointer.load(cfg.MODEL.WEIGHT) self.transforms = self.build_transform() mask_threshold = -1 if show_mask_heatmaps else 0.5 self.masker = Masker(threshold=mask_threshold, padding=1) # used to make colors for each class self.palette = torch.tensor([2**25 - 1, 2**15 - 1, 2**21 - 1]) self.cpu_device = torch.device("cpu") self.confidence_threshold = confidence_threshold self.show_mask_heatmaps = show_mask_heatmaps self.masks_per_dim = masks_per_dim
def __init__( self, cfg, confidence_threshold=0.7, show_mask_heatmaps=False, min_image_size=224, ): self.cfg = cfg.clone() self.model = build_detection_model(cfg) self.model.eval() self.device = torch.device(cfg.MODEL.DEVICE) self.model.to(self.device) self.min_image_size = min_image_size save_dir = cfg.OUTPUT_DIR checkpointer = DetectronCheckpointer(cfg, self.model, save_dir=save_dir) _ = checkpointer.load(cfg.MODEL.WEIGHT) self.transforms = self.build_transform() mask_threshold = -1 if show_mask_heatmaps else 0.5 self.masker = Masker(threshold=mask_threshold, padding=1) self.cpu_device = torch.device("cpu") self.confidence_threshold = confidence_threshold
def __init__( self, cfg, confidence_threshold=0.7, show_mask_heatmaps=False, masks_per_dim=2, min_image_size=224, ): self.cfg = cfg.clone() self.model = build_detection_model(cfg) self.model.eval() self.device = torch.device(cfg.MODEL.DEVICE) self.model.to(self.device) self.min_image_size = min_image_size save_dir = cfg.OUTPUT_DIR checkpointer = DetectronCheckpointer(cfg, self.model, save_dir=save_dir) _ = checkpointer.load(cfg.MODEL.WEIGHT) logging.info('model loaded from: {}'.format(cfg.MODEL.WEIGHT)) self.transforms = self.build_transform() mask_threshold = -1 if show_mask_heatmaps else 0.5 self.masker = Masker(threshold=mask_threshold, padding=1) # used to make colors for each class self.palette = torch.tensor([2**25 - 1, 2**15 - 1, 2**21 - 1]) self.cpu_device = torch.device("cpu") self.confidence_threshold = confidence_threshold self.show_mask_heatmaps = show_mask_heatmaps self.masks_per_dim = masks_per_dim
def __init__(self, cfg): super(RPNModule, self).__init__() self.cfg = cfg.clone() anchor_generator = make_anchor_generator(cfg) in_channels = cfg.MODEL.BACKBONE.OUT_CHANNELS rpn_head = registry.RPN_HEADS[cfg.MODEL.RPN.RPN_HEAD] head = rpn_head(cfg, in_channels, anchor_generator.num_anchors_per_location()[0]) rpn_box_coder = BoxCoder(weights=(1.0, 1.0, 1.0, 1.0)) box_selector_train = make_rpn_postprocessor(cfg, rpn_box_coder, is_train=True) box_selector_test = make_rpn_postprocessor(cfg, rpn_box_coder, is_train=False) loss_evaluator = make_rpn_loss_evaluator(cfg, rpn_box_coder) self.anchor_generator = anchor_generator self.head = head self.box_selector_train = box_selector_train self.box_selector_test = box_selector_test self.loss_evaluator = loss_evaluator
def __init__(self, cfg_path, weights_path, input_shape=(608, 608)): cfg.merge_from_file(cfg_path) cfg.merge_from_list(['DTYPE', 'float16']) self._cfg = cfg.clone() self._model = build_detection_model(self._cfg) self._model.eval() self._device = 'cuda' self._model.to(self._device) self.shape = input_shape save_dir = cfg.OUTPUT_DIR checkpoint = torch.load(weights_path, map_location=torch.device("cpu")) load_state_dict(self._model, checkpoint.pop("model")) self._transform = self._build_transform()
def __init__( self, cfg, min_image_size = 224 ): self.cfg = cfg.clone() self.model = build_detection_model(cfg) self.model.eval() self.device = torch.device(cfg.MODEL.DEVICE) self.min_image_size = min_image_size checkpointer = DetectronCheckpointer(cfg,self.model) _ = checkpointer.load(cfg.MODEL.WEIGHT) self.model.to(self.device) self.transforms = self.build_trainsform() self.cpu_device = torch.device("cpu")
def __init__( self, model, CATEGORIES, dataset, confidence_threshold=0.5, show_mask_heatmaps=False, masks_per_dim=2, min_image_size=224, ): if model == 'faster': config_file = "faster-retina/configs/e2e_faster_rcnn_R_50_FPN_1x_{}_test.yaml".format( dataset) if model == 'retinanet': config_file = 'faster-retina/configs/retinanet_R-50-FPN_1x-{}.yaml'.format( dataset) if model == 'maskrcnn': config_file = 'faster-retina/configs/e2e_mask_rcnn_R_50_FPN_1x-{}.yaml'.format( dataset) cfg.merge_from_file(config_file) self.cfg = cfg.clone() self.CATEGORIES = CATEGORIES self.model = build_detection_model(cfg) self.model.eval() self.device = torch.device('cuda') self.model.to(self.device) self.min_image_size = min_image_size self.feat_extractor = FeatureExtractorFromBoxes(self.model) save_dir = cfg.OUTPUT_DIR checkpointer = DetectronCheckpointer(cfg, self.model, save_dir=save_dir) _ = checkpointer.load(cfg.MODEL.WEIGHT) self.transforms = self.build_transform() mask_threshold = -1 if show_mask_heatmaps else 0.5 self.masker = Masker(threshold=mask_threshold, padding=1) # used to make colors for each class self.palette = torch.tensor([2**25 - 1, 2**15 - 1, 2**21 - 1]) self.cpu_device = torch.device("cpu") self.confidence_threshold = confidence_threshold self.show_mask_heatmaps = show_mask_heatmaps self.masks_per_dim = masks_per_dim
def __init__(self, cfg, image_size=(100, 100)): super(OPN, self).__init__() self.cfg = cfg.clone() self.model = build_detection_model(cfg) self.model.eval() self.backbone = self.model.backbone self.rpn = self.model.rpn self.device = torch.device(cfg.MODEL.DEVICE) self.model.to(self.device) self.image_size = image_size # (h, w) checkpointer = DetectronCheckpointer(cfg, self.model) _ = checkpointer.load(cfg.MODEL.WEIGHT) self.transforms = self.build_transform() self.cpu_device = torch.device("cpu")
def __init__(self, cfg): super(RetinaNetModule, self).__init__() self.cfg = cfg.clone() anchor_generator = make_anchor_generator_retinanet(cfg) head = RetinaNetHead(cfg) box_coder = BoxCoder(weights=(10., 10., 5., 5.)) if self.cfg.MODEL.SPARSE_MASK_ON: raise NotImplementedError else: box_selector_test = make_retinanet_postprocessor( cfg, 100, box_coder) self.anchor_generator = anchor_generator self.head = head self.box_selector_test = box_selector_test
def __init__(self, cfg, confidence_threshold=0.7, min_image_size=224, output_polygon=True): self.cfg = cfg.clone() self.model = build_detection_model(cfg) self.model.eval() self.device = torch.device(cfg.MODEL.DEVICE) self.model.to(self.device) self.min_image_size = min_image_size checkpointer = DetectronCheckpointer(cfg, self.model) _ = checkpointer.load(cfg.MODEL.WEIGHT) self.transforms = self.build_transform() self.cpu_device = torch.device("cpu") self.confidence_threshold = confidence_threshold self.output_polygon = output_polygon
def __init__( self, cfg, show_mask_heatmaps=False, min_image_size=224, ): self.cfg = cfg.clone() self.min_image_size = min_image_size self.transforms = self.build_transform() mask_threshold = -1 if show_mask_heatmaps else 0.5 self.masker = Masker(threshold=mask_threshold, padding=1) self.cpu_device = torch.device("cpu") self.visualizer = Visualizer(categories=ro_categories.CATEGORIES, cfg=cfg, confidence_threshold=0.8, show_mask_heatmaps=show_mask_heatmaps)
def __init__( self, config_file, checkpoint_path=None ): cfg.merge_from_file(config_file) cfg.defrost() if checkpoint_path: cfg.MODEL.WEIGHT = checkpoint_path cfg.freeze() self.cfg = cfg.clone() self.model = build_detection_model(cfg) self.device = torch.device(cfg.MODEL.DEVICE) self.model.to(self.device) self.cpu_device = torch.device("cpu") self.checkpointer = DetectronCheckpointer(cfg, self.model) _ = self.checkpointer.load(cfg.MODEL.WEIGHT) if hasattr(self.checkpointer, "classes"): self.classes = self.checkpointer.classes
def __init__(self, cfg: CfgNode, args): self.args = args self.cfg = cfg.clone() self.model: nn.Module = build_detection_model(cfg) self.model.eval() self.device = torch.device(cfg.MODEL.DEVICE) self.model.to(self.device) # Load weight save_dir = cfg.OUTPUT_DIR checkpointer = DetectronCheckpointer(cfg, self.model, save_dir=save_dir) _ = checkpointer.load(cfg.MODEL.WEIGHT) self.cpu_device = torch.device('cpu') # Keep all result self.confidence_threshold = 0. self.datasets = GIFDataset(cfg, args) self.results = defaultdict(list)
def __init__( self, confidence_threshold=0.7, show_mask_heatmaps=False, masks_per_dim=3, min_image_size=224, ): self.res_label_mask_scorse = [] self.res_dir = './res_person/' config_file = "../configs/e2e_fashion_mask_rcnn_R_50_FPN_1x.yaml" cfg.merge_from_file(config_file) # 设置配置文件 cfg.merge_from_list(["MODEL.MASK_ON", True]) # cfg.merge_from_list(["MODEL.DEVICE", "cpu"]) # 指定为CPU cfg.merge_from_list(["MODEL.DEVICE", "cuda"]) # 指定为GPU self.cfg = cfg.clone() self.model = build_detection_model(cfg) self.model.eval() self.device = torch.device(cfg.MODEL.DEVICE) self.model.to(self.device) self.min_image_size = min_image_size save_dir = cfg.OUTPUT_DIR checkpointer = DetectronCheckpointer(cfg, self.model, save_dir=save_dir) _ = checkpointer.load(cfg.MODEL.WEIGHT) self.transforms = self.build_transform() mask_threshold = -1 if show_mask_heatmaps else 0.5 self.masker = Masker(threshold=mask_threshold, padding=1) # used to make colors for each class self.palette = torch.tensor([2**25 - 1, 2**15 - 1, 2**21 - 1]) self.confidence_threshold = confidence_threshold self.show_mask_heatmaps = show_mask_heatmaps self.masks_per_dim = masks_per_dim
def main(args, skip_test=False): cfg = c.clone() num_gpus = int(os.environ["WORLD_SIZE"]) if "WORLD_SIZE" in os.environ else 1 args.distributed = num_gpus > 1 if args.distributed: torch.cuda.set_device(args.local_rank) torch.distributed.init_process_group( backend="nccl", init_method="env://" ) cfg.merge_from_file(args.config_file) cfg.merge_from_list(args.opts) cfg.OUTPUT_DIR = os.path.join("output", cfg.OUTPUT_DIR, "train") cfg.freeze() output_dir = cfg.OUTPUT_DIR if output_dir: mkdir(output_dir) logger = setup_logger("maskrcnn_benchmark", output_dir, get_rank()) logger.info("Using {} GPUs".format(num_gpus)) logger.info(args) logger.info("Collecting env info (might take some time)") logger.info("\n" + collect_env_info()) logger.info("Loaded configuration file {}".format(args.config_file)) with open(args.config_file, "r") as cf: config_str = "\n" + cf.read() logger.info(config_str) logger.info("Running with config:\n{}".format(cfg)) model = train(cfg, args.local_rank, args.distributed) if (not args.skip_test) and (not skip_test): test(cfg, model, args.distributed)
def __init__(self, config_file, confidence_threshold=0.3, mask_threshold=0.5, mask_on=True): cfg.merge_from_file(config_file) cfg.MODEL.MASK_ON = mask_on cfg.freeze() self.cfg = cfg.clone() self.model = build_detection_model(cfg).eval() self.device = torch.device(cfg.MODEL.DEVICE) self.model.to(self.device) checkpointer = DetectronCheckpointer(cfg, self.model, save_dir=cfg.OUTPUT_DIR) _ = checkpointer.load(cfg.MODEL.WEIGHT) self.transforms = self.build_transform() self.masker = Masker(threshold=mask_threshold, padding=1) self.cpu_device = torch.device("cpu") self.confidence_threshold = confidence_threshold
def __init__(self, cfg, conf_thresh=0.7, mask_thresh=0.5, masks_per_dim=2, min_image_size=224, ckpt=None): assert 0.0 <= conf_thresh <= 1.0 and isinstance(conf_thresh, float) assert 0.0 <= mask_thresh <= 1.0 and isinstance(mask_thresh, float) assert masks_per_dim > 0 and isinstance(masks_per_dim, int) assert min_image_size > 0 and isinstance(min_image_size, int) self.cfg = cfg.clone() self.model = build_detection_model(cfg) self.model.eval() self.device = torch.device(cfg.MODEL.DEVICE) self.model.to(self.device) self.min_image_size = min_image_size save_dir = cfg.OUTPUT_DIR checkpointer = DetectronCheckpointer(cfg, self.model, save_dir=save_dir) if ckpt: checkpointer.load_model(checkpointer.load_file(ckpt)) else: _ = checkpointer.load() self.transforms = self.build_transform() self.masker = Masker(threshold=mask_thresh, padding=1) # used to make colors for each class self.palette = torch.tensor([2**25 - 1, 2**15 - 1, 2**21 - 1]) self.cpu_device = torch.device("cpu") self.conf_thresh = conf_thresh self.masks_per_dim = masks_per_dim
def __init__(self, cfg, confidence_threshold=0.7, show_mask_heatmaps=False, masks_per_dim=2, min_image_size=224, cache_loc='cache/unknown_object_detection', image_extension='.jpg', overwrite=False): self.cache_loc = cache_loc if not os.path.exists(self.cache_loc): os.makedirs(self.cache_loc) self.extension_length = len(image_extension) self.overwrite = overwrite self.cfg = cfg.clone() self.model = build_detection_model(cfg) self.model.eval() self.device = torch.device(cfg.MODEL.DEVICE) self.model.to(self.device) self.min_image_size = min_image_size save_dir = cfg.OUTPUT_DIR checkpointer = DetectronCheckpointer(cfg, self.model, save_dir=save_dir) _ = checkpointer.load(cfg.MODEL.WEIGHT) self.transforms = self.build_transform() mask_threshold = -1 if show_mask_heatmaps else 0.5 self.masker = Masker(threshold=mask_threshold, padding=1) # used to make colors for each class self.palette = torch.tensor([2 ** 25 - 1, 2 ** 15 - 1, 2 ** 21 - 1]) self.cpu_device = torch.device("cpu") self.confidence_threshold = confidence_threshold self.show_mask_heatmaps = show_mask_heatmaps self.masks_per_dim = masks_per_dim
def __init__(self, config, ): self.confidence_threshold = 0.5 self.min_image_size = 800 self.mask_shrink = 4 # set cpu/gpu self.cpu_only = True visible_device_list = [] for handle in config.devices: if int(handle.type) == DeviceType.GPU.value: visible_device_list.append(handle.id) self.cpu_only = False if not self.cpu_only: print('Using GPU: {}'.format(visible_device_list[0])) torch.cuda.set_device(visible_device_list[0]) else: print('Using CPU') # set maskrcnn config cfg.merge_from_file(CONFIG_FILE) if self.cpu_only: cfg.merge_from_list(["MODEL.DEVICE", "cpu"]) else: cfg.merge_from_list(["MODEL.DEVICE", "cuda"]) cfg.merge_from_list(["INPUT.TO_BGR255", False]) self.cfg = cfg.clone() # build model and transform self.model = build_detection_model(cfg) self.model.eval() self.device = torch.device(cfg.MODEL.DEVICE) self.model.to(self.device) self.transforms = self.build_transform() self.masker = Masker(threshold=0.5, padding=1) self.cpu_device = torch.device("cpu")
def __init__(self, cfg, confidence_threshold=0.5): self.cfg = cfg.clone() self.model = build_detection_model(cfg) self.model.eval() self.device = torch.device(cfg.MODEL.DEVICE) self.model.to(self.device) checkpointer = DetectronCheckpointer(cfg, self.model) _ = checkpointer.load(cfg.MODEL.WEIGHT) self.transforms = self.build_transform() self.cpu_device = torch.device("cpu") self.confidence_threshold = confidence_threshold show_mask_heatmaps = True mask_threshold = -1 if show_mask_heatmaps else 0.5 self.masker = Masker(threshold=mask_threshold, padding=1) self.show_mask_heatmaps = show_mask_heatmaps self.palette = torch.tensor([2**25 - 1, 2**15 - 1, 2**21 - 1]) masks_per_dim = 2 self.cpu_device = torch.device("cpu") self.confidence_threshold = confidence_threshold self.show_mask_heatmaps = show_mask_heatmaps self.masks_per_dim = masks_per_dim
def main(): parser = argparse.ArgumentParser( description="PyTorch Object Detection Training") parser.add_argument( "--config-file", default="", metavar="FILE", help="path to config file", type=str, ) parser.add_argument("--local_rank", type=int, default=0) parser.add_argument( "--skip-test", dest="skip_test", help="Do not test the final model", action="store_true", ) parser.add_argument( "opts", help="Modify config options using the command-line", default=None, nargs=argparse.REMAINDER, ) args = parser.parse_args() num_gpus = int( os.environ["WORLD_SIZE"]) if "WORLD_SIZE" in os.environ else 1 args.distributed = num_gpus > 1 if args.distributed: torch.cuda.set_device(args.local_rank) torch.distributed.init_process_group(backend="nccl", init_method="env://") synchronize() cfg.merge_from_file(args.config_file) cfg.merge_from_list(args.opts) cfg.MODEL.ROI_BOX_HEAD.NUM_CLASSES = 31 cfg.SOLVER.BASE_LR = 0.01 cfg_origial = cfg.clone() cfg.freeze() output_dir = cfg.OUTPUT_DIR if output_dir: mkdir(output_dir) logger = setup_logger("maskrcnn_benchmark", output_dir, get_rank()) logger.info("Using {} GPUs".format(num_gpus)) logger.info(args) logger.info("Collecting env info (might take some time)") logger.info("\n" + collect_env_info()) logger.info("Loaded configuration file {}".format(args.config_file)) with open(args.config_file, "r") as cf: config_str = "\n" + cf.read() logger.info(config_str) logger.info("Running with config:\n{}".format(cfg)) model = train(cfg, cfg_origial, args.local_rank, args.distributed) if not args.skip_test: run_test(cfg, model, args.distributed)
def load_cfg(conf): config = cfg.clone() config.merge_from_file(conf.config_file) config.freeze() return config
def main(): parser = argparse.ArgumentParser( description="PyTorch Object Detection Inference") parser.add_argument( "--config", default="maskrcnn_benchmark/data/datasets/predictor.yaml", metavar="FILE", help="path to config file", ) parser.add_argument("--local_rank", type=int, default=0) parser.add_argument( "opts", help="Modify config options using the command-line", default=None, nargs=argparse.REMAINDER, ) args = parser.parse_args() num_gpus = int( os.environ["WORLD_SIZE"]) if "WORLD_SIZE" in os.environ else 1 distributed = num_gpus > 1 # if distributed: # torch.cuda.set_device(args.local_rank) # torch.distributed.init_process_group( # backend="nccl", init_method="env://" # ) # synchronize() cfg.merge_from_file(args.config) # cfg.merge_from_list(args.opts) # cfg.freeze() mycfg = cfg.clone() save_dir = "" logger = setup_logger("maskrcnn_benchmark", save_dir, get_rank()) logger.info("Using {} GPUs".format(num_gpus)) logger.info(cfg) logger.info("Collecting env info (might take some time)") logger.info("\n" + collect_env_info()) model = build_detection_model(mycfg) model.eval() model.to(mycfg.MODEL.DEVICE) output_dir = cfg.OUTPUT_DIR # checkpointer = DetectronCheckpointer(cfg, model, save_dir=output_dir) # _ = checkpointer.load(cfg.MODEL.WEIGHT) iou_types = ("bbox", ) if cfg.MODEL.MASK_ON: iou_types = iou_types + ("segm", ) if cfg.MODEL.KEYPOINT_ON: iou_types = iou_types + ("keypoints", ) output_folders = [None] * len(cfg.DATASETS.TEST) dataset_names = cfg.DATASETS.TEST if cfg.OUTPUT_DIR: for idx, dataset_name in enumerate(dataset_names): output_folder = os.path.join(cfg.OUTPUT_DIR, "inference", dataset_name) mkdir(output_folder) output_folders[idx] = output_folder data_loaders_val = make_data_loader(cfg, is_train=False, is_distributed=distributed) for output_folder, dataset_name, data_loader_val in zip( output_folders, dataset_names, data_loaders_val): inference( model, data_loader_val, dataset_name=dataset_name, iou_types=iou_types, box_only=False if cfg.MODEL.RETINANET_ON else cfg.MODEL.RPN_ONLY, device=cfg.MODEL.DEVICE, expected_results=cfg.TEST.EXPECTED_RESULTS, expected_results_sigma_tol=cfg.TEST.EXPECTED_RESULTS_SIGMA_TOL, output_folder=output_folder, ) synchronize()
def __init__(self, cfg): self.cfg = cfg.clone() self.transforms = self.build_transform()
def train_with_validation(cfg, local_rank, distributed, test_weights=None): arguments = {} arguments["iteration"] = 0 if test_weights: cfg.MODEL.WEIGHT = test_weights cfg.SOLVER.MAX_ITER = 0 ignore_labels = (cfg.MODEL.ROI_BOX_HEAD.NUM_CLASSES == 0) # prepare training data root_path = os.path.expanduser(os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", "..", "data"))) data_loader, class_ids = make_data_loader( root_path, cfg, is_train=True, is_distributed=distributed, start_iter=arguments["iteration"], ignore_labels=ignore_labels, ) # overwrite the number of classes by considering the training set if cfg.MODEL.ROI_BOX_HEAD.NUM_CLASSES <= 0: # if we have binary classification or unknown number of classes cfg.MODEL.ROI_BOX_HEAD.NUM_CLASSES = len(class_ids) # prepare model model = build_detection_model(cfg) device = torch.device(cfg.MODEL.DEVICE) model.to(device) # prepare optimizer optimizer = make_optimizer(cfg, model) scheduler = make_lr_scheduler(cfg, optimizer) if distributed: model = torch.nn.parallel.DistributedDataParallel( model, device_ids=[local_rank], output_device=local_rank, # this should be removed if we update BatchNorm stats broadcast_buffers=False, ) # prepare validation run_validation_for_model = partial(run_validation, root_path=root_path, cfg=cfg.clone(), class_ids=class_ids, ignore_labels=ignore_labels, distributed=distributed) # setup checkpointer output_dir = cfg.OUTPUT_DIR save_to_disk = get_rank() == 0 checkpointer = DetectronCheckpointer( cfg, model, optimizer, scheduler, output_dir, save_to_disk ) extra_checkpoint_data = checkpointer.load(cfg.MODEL.WEIGHT, use_latest=False if test_weights else True) arguments.update(extra_checkpoint_data) checkpoint_period = cfg.SOLVER.CHECKPOINT_PERIOD validation_period = cfg.SOLVER.VALIDATION_PERIOD # start training do_train( model, data_loader, optimizer, scheduler, checkpointer, device, validation_period, checkpoint_period, arguments, run_validation_for_model) return model