def main(args): cfg = setup(args) if args.eval_only: model = Trainer.build_model(cfg) DetectionCheckpointer(model, save_dir=cfg.OUTPUT_DIR).resume_or_load( cfg.MODEL.WEIGHTS, resume=args.resume) # res = Trainer.test_with_TTA(cfg, model) Trainer.inference_with_TTA(cfg, model, args) return else: raise Exception("Only evaluation supported")
def main(args): cfg = setup(args) if args.eval_only: model = TrainerNoMeta.build_model(cfg) DetectionCheckpointer(model, save_dir=cfg.OUTPUT_DIR).resume_or_load( cfg.MODEL.WEIGHTS, resume=args.resume) res = TrainerNoMeta.test(cfg, model) return res trainer = TrainerNoMeta(cfg) trainer.resume_or_load(resume=args.resume) return trainer.train()
def __init__(self, cfg): """ Args: cfg (CfgNode): """ super().__init__() logger = logging.getLogger("detectron2") if not logger.isEnabledFor( logging.INFO): # setup_logger is not called for d2 setup_logger() cfg = DefaultTrainer.auto_scale_workers(cfg, comm.get_world_size()) # Assume these objects must be constructed in this order. model = self.build_model(cfg) optimizer = self.build_optimizer(cfg, model) data_loader = self.build_train_loader(cfg) # For training, wrap with DDP. But don't need this for inference. if comm.get_world_size() > 1: model = DistributedDataParallel(model, device_ids=[comm.get_local_rank()], broadcast_buffers=False) self._trainer = (AMPTrainer if cfg.SOLVER.AMP.ENABLED else SimpleTrainer)(model, data_loader, optimizer) self.scheduler = self.build_lr_scheduler(cfg, optimizer) # Assume no other objects need to be checkpointed. # We can later make it checkpoint the stateful hooks self.checkpointer = DetectionCheckpointer( # Assume you want to save checkpoints together with logs/statistics model, cfg.OUTPUT_DIR, optimizer=optimizer, scheduler=self.scheduler, ) self.start_iter = 0 self.max_iter = cfg.SOLVER.MAX_ITER self.cfg = cfg self.register_hooks(self.build_hooks())
def evaluate(args, mode, _appcfg): name = "hmd" #uncomment if using trainer.model # for subset in ["train", "val"]: # metadata = load_and_register_dataset(name, subset, _appcfg) subset = "test" # subset = "val" metadata = load_and_register_dataset(name, subset, _appcfg) dataset_name = get_dataset_name(name, subset) dataset_dicts = DatasetCatalog.get(dataset_name) cfg = config.get_cfg() cfg.merge_from_file( "/codehub/external/detectron2/configs/COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml" ) cfg.DATASETS.TRAIN = ("hmd_train", "hmd_val") cfg.DATASETS.TEST = (dataset_name) cfg.OUTPUT_DIR = "/codehub/apps/detectron2/release" cfg.MODEL.WEIGHTS = os.path.join(cfg.OUTPUT_DIR, "model_final.pth") cfg.DATALOADER.NUM_WORKERS = 2 cfg.SOLVER.IMS_PER_BATCH = 2 cfg.SOLVER.BASE_LR = 0.00025 cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 128 # faster, and good enough for this toy dataset cfg.MODEL.ROI_HEADS.NUM_CLASSES = 3 cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.7 # mapper = DatasetMapper(cfg, False) _loader = build_detection_test_loader(cfg, dataset_name) evaluator = COCOEvaluator(dataset_name, cfg, False, output_dir=cfg.OUTPUT_DIR) # trainer = DefaultTrainer(cfg) # trainer.resume_or_load(resume=True) # model = trainer.model # predictor = DefaultPredictor(cfg) # model = predictor.model file_path = cfg.MODEL.WEIGHTS model = build_model(cfg) DetectionCheckpointer(model).load(file_path) inference_on_dataset(model, _loader, evaluator)
def load_model(config_file, model_weights, model_device): print('Loading cfg') args = [ '--config-file', config_file, 'MODEL.WEIGHTS', model_weights, 'MODEL.DEVICE', model_device, ] args = default_argument_parser().parse_args(args) cfg = get_cfg() cfg.merge_from_file(args.config_file) cfg.merge_from_list(args.opts) cfg.freeze() default_setup(cfg, args) print('Loading model') model = build_model(cfg) checkpointer = DetectionCheckpointer(model, cfg.OUTPUT_DIR) checkpointer.resume_or_load(cfg.MODEL.WEIGHTS, resume=True) return model
class Tester: def __init__(self, cfg): self.cfg = cfg self.model = Trainer.build_model(cfg) self.check_pointer = DetectionCheckpointer( self.model, save_dir=cfg.OUTPUT_DIR ) self.best_res = None self.best_file = None self.all_res = {} def test(self, ckpt): self.check_pointer._load_model(self.check_pointer._load_file(ckpt)) print("evaluating checkpoint {}".format(ckpt)) res = Trainer.test(self.cfg, self.model) if comm.is_main_process(): verify_results(self.cfg, res) print(res) if (self.best_res is None) or ( self.best_res is not None and self.best_res["bbox"]["AP"] < res["bbox"]["AP"] ): self.best_res = res self.best_file = ckpt print("best results from checkpoint {}".format(self.best_file)) print(self.best_res) self.all_res["best_file"] = self.best_file self.all_res["best_res"] = self.best_res self.all_res[ckpt] = res os.makedirs( os.path.join(self.cfg.OUTPUT_DIR, "inference"), exist_ok=True ) with open( os.path.join(self.cfg.OUTPUT_DIR, "inference", "all_res.json"), "w", ) as fp: json.dump(self.all_res, fp)
def build_train_loader(cls, cfg: CfgNode): data_loader = build_detection_train_loader(cfg, mapper=DatasetMapper(cfg, True)) if not has_inference_based_loaders(cfg): return data_loader model = cls.build_model(cfg) model.to(cfg.BOOTSTRAP_MODEL.DEVICE) DetectionCheckpointer(model).resume_or_load(cfg.BOOTSTRAP_MODEL.WEIGHTS, resume=False) inference_based_loaders, ratios = build_inference_based_loaders(cfg, model) loaders = [data_loader] + inference_based_loaders ratios = [1.0] + ratios combined_data_loader = build_combined_loader(cfg, loaders, ratios) sample_counting_loader = SampleCountingLoader(combined_data_loader) return sample_counting_loader
def __init__(self, num_classes, num_votes, fix_bn=True, freeze_at=2, norm='GN', fpn_norm='', conv_dims=128, **kwargs): # FPN_C, backbone='resnet50'): super().__init__() d2_cfg = detectron2.config.get_cfg() d2_cfg.merge_from_file( '/home-nfs/whc/glab/detectron2/configs/COCO-PanopticSegmentation/panoptic_fpn_R_50_1x.yaml' ) if fix_bn: d2_cfg.MODEL.RESNETS.NORM = "FrozenBN" else: d2_cfg.MODEL.RESNETS.NORM = "BN" d2_cfg.MODEL.BACKBONE.FREEZE_AT = freeze_at d2_cfg.MODEL.FPN.NORM = 'BN' self.backbone = build_backbone(d2_cfg) d2_cfg.MODEL.SEM_SEG_HEAD.NORM = norm d2_cfg.MODEL.SEM_SEG_HEAD.CONVS_DIM = conv_dims d2_cfg.MODEL.SEM_SEG_HEAD.NUM_CLASSES = num_classes self.sem_seg_head = build_sem_seg_head(d2_cfg, self.backbone.output_shape()) self.sem_classifier = self.sem_seg_head.predictor del self.sem_seg_head.predictor d2_cfg.MODEL.SEM_SEG_HEAD.NUM_CLASSES = num_votes tmp = build_sem_seg_head(d2_cfg, self.backbone.output_shape()) self.vote_classifier = tmp.predictor assert cfg.data.dataset.params.caffe_mode == True checkpointer = DetectionCheckpointer(self) checkpointer.load(d2_cfg.MODEL.WEIGHTS)
def run_train(): torch.multiprocessing.freeze_support() cfg = get_cfg() cfg.merge_from_file(model_zoo.get_config_file('COCO-Detection/faster_rcnn_R_101_FPN_3x.yaml')) # cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.75 # Threshold cfg.MODEL.WEIGHTS = "detectron2://COCO-Detection/faster_rcnn_R_101_FPN_3x/137851257/model_final_f6e8b1.pkl" cfg.MODEL.DEVICE = "cpu" # cpu or cuda register_datasets() cfg.DATASETS.TRAIN = ('grini_nc_merged_bbox_only_train',) cfg.DATASETS.TEST = ('grini_nc_merged_bbox_only_val',) # cfg.MODEL.WEIGHTS = get_checkpoint_url('COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_1x.yaml') cfg.MODEL.DEVICE = "cpu" # cpu or cuda # todo find out how rescale images and annotations first... # Parameters fixed cfg.SOLVER.IMS_PER_BATCH = 4 cfg.SOLVER.BASE_LR = 0.001 cfg.SOLVER.WARMUP_ITERS = 1000 cfg.SOLVER.MAX_ITER = 1500 # adjust up if val mAP is still rising, adjust down if overfit cfg.SOLVER.STEPS = (1000, 1500) cfg.SOLVER.GAMMA = 0.05 cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 12 cfg.MODEL.ROI_HEADS.NUM_CLASSES = 3 cfg.TEST.EVAL_PERIOD = 500 # makedirs(cfg.OUTPUT_DIR, exist_ok=True) setup_logger(output=cfg.OUTPUT_DIR, distributed_rank=comm.get_rank(), name="my_model") # DetectionCheckpointer(cfg).load(file_path_or_url) # load a file, usually from cfg.MODEL.WEIGHTS checkpointer = DetectionCheckpointer(build_model(cfg), save_dir=cfg.OUTPUT_DIR) checkpointer.save("model_faster_rcnn_unscaled") # save to output/model_999.pth trainer = CocoTrainer(cfg) trainer.resume_or_load(resume=False) trainer.train()
def do_activation(cfg): if isinstance(cfg, CfgNode): data_loader = build_detection_test_loader(cfg, cfg.DATASETS.TEST[0]) model = build_model(cfg) DetectionCheckpointer(model).load(cfg.MODEL.WEIGHTS) else: data_loader = instantiate(cfg.dataloader.test) model = instantiate(cfg.model) model.to(cfg.train.device) DetectionCheckpointer(model).load(cfg.train.init_checkpoint) model.eval() counts = Counter() total_activations = [] for idx, data in zip(tqdm.trange(args.num_inputs), data_loader): # noqa count = activation_count_operators(model, data) counts += count total_activations.append(sum(count.values())) logger.info("(Million) Activations for Each Type of Operators:\n" + str([(k, v / idx) for k, v in counts.items()])) logger.info("Total (Million) Activations: {}±{}".format( np.mean(total_activations), np.std(total_activations)))
def do_activation(cfg): data_loader = build_detection_test_loader(cfg, cfg.DATASETS.TEST[0]) model = build_model(cfg) DetectionCheckpointer(model).load(cfg.MODEL.WEIGHTS) model.eval() counts = Counter() for idx, data in zip(tqdm.trange(args.num_inputs), data_loader): # noqa counts += activation_count_operators(model, data) logger.info( "(Million) Activations for Each Type of Operators:\n" + str([(k, v / idx) for k, v in counts.items()]) )
def main(args): cfg = setup(args) model = build_model(cfg) logger.info("Model:\n{}".format(model)) if args.eval_only: DetectionCheckpointer(model, save_dir=cfg.OUTPUT_DIR).resume_or_load( cfg.MODEL.WEIGHTS, resume=args.resume ) return do_test(cfg, model) do_train(cfg, model, resume=args.resume) return do_test(cfg, model)
def __init__(self, cfg_file, model_path, classes, confidence_thresh=0.5): self.cpu_device = torch.device("cpu") self.cfg = get_cfg() self.cfg.merge_from_file(cfg_file) self.cfg.MODEL.WEIGHTS = os.path.join(model_path, "model_final.pth") self.cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = confidence_thresh self.model = build_model(self.cfg) self.model.eval() MetadataCatalog.get(self.cfg.DATASETS.TEST[0]).thing_classes = classes self.metadata = MetadataCatalog.get(self.cfg.DATASETS.TEST[0]) checkpointer = DetectionCheckpointer(self.model) checkpointer.load(self.cfg.MODEL.WEIGHTS) self.transform_gen = T.ResizeShortestEdge( [self.cfg.INPUT.MIN_SIZE_TEST, self.cfg.INPUT.MIN_SIZE_TEST], self.cfg.INPUT.MAX_SIZE_TEST) self.input_format = self.cfg.INPUT.FORMAT assert self.input_format in ["RGB", "BGR"], self.input_format
def main(args): cfg = LazyConfig.load(args.config_file) cfg = LazyConfig.apply_overrides(cfg, args.opts) default_setup(cfg, args) if args.eval_only: model = instantiate(cfg.model) model.to(cfg.train.device) model = create_ddp_model(model) DetectionCheckpointer(model).load(cfg.train.init_checkpoint) print(do_test(cfg, model)) else: do_train(args, cfg)
def load_model(config_file, model_path): cfg = get_cfg() add_centernet_config(cfg) cfg.merge_from_file(config_file) forward = {'centerX': centerX_forward} # model model = build_model(cfg) model.forward = MethodType(forward['centerX'], model) DetectionCheckpointer(model).load(model_path) model.eval() model.cuda() return model
def benchmark_eval(args): cfg = setup(args) if args.config_file.endswith(".yaml"): model = build_model(cfg) DetectionCheckpointer(model).load(cfg.MODEL.WEIGHTS) cfg.defrost() cfg.DATALOADER.NUM_WORKERS = 0 data_loader = build_detection_test_loader(cfg, cfg.DATASETS.TEST[0]) else: model = instantiate(cfg.model) model.to(cfg.train.device) DetectionCheckpointer(model).load(cfg.train.init_checkpoint) cfg.dataloader.num_workers = 0 data_loader = instantiate(cfg.dataloader.test) model.eval() logger.info("Model:\n{}".format(model)) dummy_data = DatasetFromList(list(itertools.islice(data_loader, 100)), copy=False) def f(): while True: yield from dummy_data for k in range(5): # warmup model(dummy_data[k]) max_iter = 300 timer = Timer() with tqdm.tqdm(total=max_iter) as pbar: for idx, d in enumerate(f()): if idx == max_iter: break model(d) pbar.update() logger.info("{} iters in {} seconds.".format(max_iter, timer.seconds()))
def __init__(self, cfg): # Create common attributes. self.cfg = cfg.clone() # cfg can be modified by model self.model = build_model(self.cfg) self.model_list = [] # Parse config self.inference_mode = self.cfg.PROBABILISTIC_INFERENCE.INFERENCE_MODE self.mc_dropout_enabled = self.cfg.PROBABILISTIC_INFERENCE.MC_DROPOUT.ENABLE self.num_mc_dropout_runs = self.cfg.PROBABILISTIC_INFERENCE.MC_DROPOUT.NUM_RUNS # Set model to train for MC-Dropout runs if self.mc_dropout_enabled: self.model.train() else: self.model.eval() # Create ensemble if applicable. if self.inference_mode == 'ensembles': ensemble_random_seeds = self.cfg.PROBABILISTIC_INFERENCE.ENSEMBLES.RANDOM_SEED_NUMS for i, random_seed in enumerate(ensemble_random_seeds): model = build_model(self.cfg) model.eval() checkpoint_dir = os.path.join( os.path.split(self.cfg.OUTPUT_DIR)[0], 'random_seed_' + str(random_seed)) # Load last checkpoint. DetectionCheckpointer(model, save_dir=checkpoint_dir).resume_or_load( cfg.MODEL.WEIGHTS, resume=True) self.model_list.append(model) else: # Or Load single model last checkpoint. DetectionCheckpointer(self.model, save_dir=cfg.OUTPUT_DIR).resume_or_load( cfg.MODEL.WEIGHTS, resume=True)
def __init__(self, cf, cf_add_d2, dataset, out): num_classes = len(dataset.action_names) TEST_DATASET_NAME = 'daly_objaction_test' # / Define d2 conf d2_output_dir = str(small.mkdir(out / 'd2_output')) d_cfg = set_detectron_cfg_base(d2_output_dir, num_classes, cf['seed']) d_cfg = set_detectron_cfg_test(d_cfg, TEST_DATASET_NAME, cf['d2_rcnn.model'], cf['d2_rcnn.conf_thresh'], cf_add_d2, freeze=False) d_cfg.MODEL.PROPOSAL_GENERATOR.NAME = "PrecomputedProposals" d_cfg.freeze() # / Start d2 simple_d2_setup(d_cfg) # Predictor without proposal generator model = build_model(d_cfg) model.eval() checkpointer = DetectionCheckpointer(model) checkpointer.load(d_cfg.MODEL.WEIGHTS) MIN_SIZE_TEST = d_cfg.INPUT.MIN_SIZE_TEST MAX_SIZE_TEST = d_cfg.INPUT.MAX_SIZE_TEST transform_gen = d2_transforms.ResizeShortestEdge( [MIN_SIZE_TEST, MIN_SIZE_TEST], MAX_SIZE_TEST) # Instance monkeypatching # https://stackoverflow.com/questions/50599045/python-replacing-a-function-within-a-class-of-a-module/50600307#50600307 model.forward = MethodType(genrcnn_rcnn_roiscores_forward, model) self.d_cfg = d_cfg self.rcnn_roiscores_model = model self.cpu_device = torch.device("cpu") self.transform_gen = transform_gen
def __init__(self, cfg): """ Args: cfg (CfgNode): """ # Assume these objects must be constructed in this order. model = self.build_model(cfg) optimizer = self.build_optimizer(cfg, model) data_loader = self.build_train_loader(cfg) # For training, wrap with DDP. But don't need this for inference. if comm.get_world_size() > 1: model = DistributedDataParallel(model, device_ids=[comm.get_local_rank()], broadcast_buffers=False, find_unused_parameters=True) super().__init__(model, data_loader, optimizer) self.scheduler = self.build_lr_scheduler(cfg, optimizer) # Assume no other objects need to be checkpointed. # We can later make it checkpoint the stateful hooks checkpoint_dir = osp.join(cfg.OUTPUT_DIR, "checkpoints") if not osp.exists(checkpoint_dir): PathManager.mkdirs(checkpoint_dir) self.checkpointer = DetectionCheckpointer( # Assume you want to save checkpoits together with logs/statistics model, checkpoint_dir, optimizer=optimizer, scheduler=self.scheduler, ) self.start_iter = 0 self.max_iter = cfg.SOLVER.MAX_ITER self.cfg = cfg self.register_hooks(self.build_hooks())
def main(args): cfg = setup(args) register_lofar_datasets(cfg) if args.eval_only: model = LOFARTrainer.build_model(cfg) DetectionCheckpointer(model, save_dir=cfg.OUTPUT_DIR).resume_or_load( cfg.MODEL.WEIGHTS, resume=args.resume) res = LOFARTrainer.test(cfg, model) return res trainer = LOFARTrainer(cfg) trainer.resume_or_load(resume=args.resume) return trainer.train()
def __init__(self, context: PyTorchTrialContext): self.context = context self.cfg = self.setup_cfg() model = build_model(self.cfg) checkpointer = DetectionCheckpointer( model, self.cfg.OUTPUT_DIR ) checkpointer.resume_or_load(self.cfg.MODEL.WEIGHTS, resume=False) self.model = self.context.wrap_model(checkpointer.model) optimizer = build_optimizer(self.cfg, self.model) self.optimizer = self.context.wrap_optimizer(optimizer) self.scheduler = build_lr_scheduler(self.cfg, self.optimizer) self.scheduler = self.context.wrap_lr_scheduler(self.scheduler,LRScheduler.StepMode.STEP_EVERY_BATCH) self.dataset_name = self.cfg.DATASETS.TEST[0] self.evaluators = get_evaluator(self.cfg, self.dataset_name, self.context.get_hparam("output_dir"), self.context.get_hparam('fake_data')) self.val_reducer = self.context.wrap_reducer(EvaluatorReducer(self.evaluators), for_training=False) self.context.experimental.disable_dataset_reproducibility_checks()
def main(args): cfg = setup(args) model = build_model(cfg) DetectionCheckpointer(model, save_dir=cfg.OUTPUT_DIR).resume_or_load( cfg.MODEL.WEIGHTS, resume=False) dataset_path = '' with open(dataset_path, 'r') as f: dataset = json.load(f) image_path = './images/' image_list = list(map(str, list(vqamb.keys()))) do_feature_extraction(cfg, model, image_path, image_list[args.lower:args.higher])
def main(args): cfg = setup(args) if args.eval_only: model = Trainer.build_model(cfg) DetectionCheckpointer(model, save_dir=cfg.OUTPUT_DIR).resume_or_load( cfg.MODEL.WEIGHTS, resume=args.resume) res = Trainer.test(cfg, model) if comm.is_main_process(): verify_results(cfg, res) return res trainer = Trainer(cfg) trainer.resume_or_load(resume=args.resume) return trainer.train()
def create_model(): from detectron2 import model_zoo from detectron2.config import get_cfg from detectron2.checkpoint import DetectionCheckpointer from detectron2.modeling import build_model # Create a config cfg = get_cfg() cfg.merge_from_file( model_zoo.get_config_file( "COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml")) cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.7 # set threshold for this model cfg.TEST.DETECTIONS_PER_IMAGE = 5 cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url( "COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml") model = build_model(cfg) model.eval() checkpointer = DetectionCheckpointer(model) checkpointer.load(cfg.MODEL.WEIGHTS) return model, cfg
def __init__(self, cfg): """ Args: cfg (CfgNode): """ # Assume these objects must be constructed in this order. model = self.build_model(cfg) if cfg.FREEZE_ALL: # Freeze model parameters except for viewpoint for name,param in model.named_parameters(): if param.requires_grad and 'viewpoint' not in name: param.requires_grad = False optimizer = self.build_optimizer(cfg, model) data_loader = self.build_train_loader(cfg) # For training, wrap with DDP. But don't need this for inference. if comm.get_world_size() > 1: model = DistributedDataParallel( model, device_ids=[comm.get_local_rank()], broadcast_buffers=False ) super().__init__(model, data_loader, optimizer) self.scheduler = self.build_lr_scheduler(cfg, optimizer) # Assume no other objects need to be checkpointed. # We can later make it checkpoint the stateful hooks self.checkpointer = DetectionCheckpointer( # Assume you want to save checkpoints together with logs/statistics model, cfg.OUTPUT_DIR, optimizer=optimizer, scheduler=self.scheduler, ) self.start_iter = 0 self.max_iter = cfg.SOLVER.MAX_ITER self.cfg = cfg self.register_hooks(self.build_hooks())
def get_teacher(args, cfg, logger): class ModelWrapper: def __init__(self, model): self.model = model @torch.no_grad() def __call__(self, img: torch.Tensor) -> dict: img = img.permute(2, 0, 1).type(torch.float32) input = {'image': img, 'height': 720, 'width': 1280} return self.model([input])[0] class ModelTimerWrapper: def __init__(self, model): self.model = model self.inference_timer = TimeAverageMeter('MRCNN inference', 'SERVER', True) @torch.no_grad() def __call__(self, img: torch.Tensor) -> dict: self.inference_timer.tic() img = img.permute(2, 0, 1).type(torch.float32) input = {'image': img, 'height': 720, 'width': 1280} output = self.model([input])[0] self.inference_timer.toc() return output model = build_model(cfg) checkpointer = DetectionCheckpointer(model) checkpointer.load(cfg.MODEL.WEIGHTS) model.eval() if args.time: model = ModelTimerWrapper(model) else: model = ModelWrapper(model) logger.critical('Loaded teacher from model zoo') return model
def main(args): # register_coco_instances('asparagus_train', {'_background_': 0, 'clump': 1, 'stalk': 2, 'spear': 3, 'bar': 4} , "./datasets/coco/annotations/train/annotations.json", "./datasets/coco/annotations/train") # register_coco_instances('asparagus_val', {'_background_': 0, 'clump': 1, 'stalk': 2, 'spear': 3, 'bar': 4} , "./datasets/coco/annotations/test_458/annotations.json", "./datasets/coco/annotations/test_458") register_coco_instances( 'asparagus_train', { '_background_': 0, 'clump': 1, 'stalk': 2, 'spear': 3, 'bar': 4, 'straw': 5 }, "./datasets/coco/annotations/straw/train/annotations.json", "./datasets/coco/annotations/straw/train") register_coco_instances( 'asparagus_val', { '_background_': 0, 'clump': 1, 'stalk': 2, 'spear': 3, 'bar': 4, 'straw': 5 }, "./datasets/coco/annotations/val_straw/val/annotations.json", "./datasets/coco/annotations/val_straw/val") cfg = setup(args) if args.eval_only: model = Trainer.build_model(cfg) DetectionCheckpointer(model, save_dir=cfg.OUTPUT_DIR).resume_or_load( cfg.MODEL.WEIGHTS, resume=args.resume) res = Trainer.test(cfg, model) if cfg.TEST.AUG.ENABLED: res.update(Trainer.test_with_TTA(cfg, model)) if comm.is_main_process(): verify_results(cfg, res) return res """ If you'd like to do anything fancier than the standard training logic, consider writing your own training loop (see plain_train_net.py) or subclassing the trainer. """ trainer = Trainer(cfg) trainer.resume_or_load(resume=args.resume) if cfg.TEST.AUG.ENABLED: trainer.register_hooks([ hooks.EvalHook(0, lambda: trainer.test_with_TTA(cfg, trainer.model)) ]) return trainer.train()
def get_detector(config_file, nms_thresh, data_type, weight_dir=None, dataset_name=None): """ for each image you make an prediction and find the indexes of correctly classified rois returns: - model: object detector(nn.module) - cfg: config file for the module params: - config_file: - nms_thresh: it determins the minimum confidence threshold for rois to survive for NMS - data_type: - weight_file: """ cfg = get_cfg() # add project-specific config (e.g., TensorMask) here if you're not running a model in detectron2's core library cfg.merge_from_file(model_zoo.get_config_file(config_file)) # set threshold for this model cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = nms_thresh # Find a model from detectron2's model zoo. You can use the https://dl.fbaipublicfiles... url as well if weight_dir is None: cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url( config_file) # use the pretrained weights else: cfg.MODEL.WEIGHTS = os.path.join( weight_dir, 'model_final.pth') # use the custom trained weights cfg.MODEL.ROI_HEADS.NUM_CLASSES = len( MetadataCatalog.get( dataset_name).thing_classes) # number of classes if data_type == 'model': model = build_model(cfg) # returns a torch.nn.Module DetectionCheckpointer(model).load(cfg.MODEL.WEIGHTS) device = torch.device( 'cuda') if torch.cuda.is_available() else torch.device('cpu') model = model.to(device) model.eval() if data_type == 'predictor': model = DefaultPredictor(cfg) return model, cfg
def __init__(self, sess, model_saved_path, labels_path, fix_ratio, config_path): with open(labels_path) as f: self.labels = f.read().splitlines() self.image_shape = [600, 600, 3] # TODO self.cfg = cfg = get_cfg() cfg.merge_from_file(config_path) cfg.freeze() self.model = build_model(cfg) self.model.eval() DetectionCheckpointer(self.model).load(model_saved_path) self.fix_ratio = fix_ratio if not fix_ratio: self.transform_gen = T.ResizeShortestEdge( [cfg.INPUT.MIN_SIZE_TEST, cfg.INPUT.MIN_SIZE_TEST], cfg.INPUT.MAX_SIZE_TEST)
def _get_model(args): cfg = get_cfg() cfg.merge_from_file(args.config) cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.7 # set threshold for this model model = build_model(cfg) # returns a torch.nn.Module weights = args.weights if args.weights != None else cfg.MODEL.WEIGHTS DetectionCheckpointer(model).load( weights ) # must load weights this way, can't use cfg.MODEL.WEIGHTS = "..." model.train(False) # inference mode return model