def test_mask_rcnn(self): model = mask_rcnn.maskrcnn_resnet50_fpn(pretrained=False, pretrained_backbone=True, min_size=200, max_size=300) images, test_images = _get_test_images() self.run_test(model, (images, ), rtol=1e-3, atol=1e-5) self.run_test( model, (images, ), input_names=["images_tensors"], output_names=["boxes", "labels", "scores", "masks"], dynamic_axes={ "images_tensors": [0, 1, 2], "boxes": [0, 1], "labels": [0], "scores": [0], "masks": [0, 1, 2], }, rtol=1e-3, atol=1e-5, ) dummy_image = [torch.ones(3, 100, 100) * 0.3] self.run_test( model, (images, ), additional_test_inputs=[(images, ), (test_images, ), (dummy_image, )], input_names=["images_tensors"], output_names=["boxes", "labels", "scores", "masks"], dynamic_axes={ "images_tensors": [0, 1, 2], "boxes": [0, 1], "labels": [0], "scores": [0], "masks": [0, 1, 2], }, rtol=1e-3, atol=1e-5, ) self.run_test( model, (dummy_image, ), additional_test_inputs=[(dummy_image, ), (images, )], input_names=["images_tensors"], output_names=["boxes", "labels", "scores", "masks"], dynamic_axes={ "images_tensors": [0, 1, 2], "boxes": [0, 1], "labels": [0], "scores": [0], "masks": [0, 1, 2], }, rtol=1e-3, atol=1e-5, )
def multitask_maskrcnn_resnet50_fpn( **kwargs ): task_heads = [ ClassificationTaskHead(name, num_classes) for name, num_classes in [("roof_style", 4), ("roof_material", 4)] ] task_heads = [head.cuda() for head in task_heads] model = maskrcnn_resnet50_fpn(num_classes=NUM_CLASSES, **kwargs) [model.add_module(head.name, head) for head in task_heads] model.roi_heads.forward = partial(forward, model.roi_heads, task_heads=task_heads) return model
def get_model(num_classes, pretrained=True): model = maskrcnn_resnet50_fpn(pretrained=pretrained) in_features = model.roi_heads.box_predictor.cls_score.in_features model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes) in_features_mask = model.roi_heads.mask_predictor.conv5_mask.in_channels print("IFM", in_features_mask) hidden_layer = 256 model.roi_heads.mask_predictor = MaskRCNNPredictor(in_features_mask, hidden_layer, num_classes) return model
def model( num_classes: int, backbone: Optional[nn.Module] = None, remove_internal_transforms: bool = True, pretrained: bool = True, **mask_rcnn_kwargs ) -> nn.Module: """MaskRCNN model implemented by torchvision. # Arguments num_classes: Number of classes. backbone: Backbone model to use. Defaults to a resnet50_fpn model. remove_internal_transforms: The torchvision model internally applies transforms like resizing and normalization, but we already do this at the `Dataset` level, so it's safe to remove those internal transforms. pretrained: Argument passed to `maskrcnn_resnet50_fpn` if `backbone is None`. By default it is set to True: this is generally used when training a new model (transfer learning). `pretrained = False` is used during inference (prediction) for cases where the users have their own pretrained weights. **mask_rcnn_kwargs: Keyword arguments that internally are going to be passed to `torchvision.models.detection.mask_rcnn.MaskRCNN`. # Return A Pytorch `nn.Module`. """ if backbone is None: model = maskrcnn_resnet50_fpn( pretrained=pretrained, pretrained_backbone=pretrained, **mask_rcnn_kwargs ) in_features_box = model.roi_heads.box_predictor.cls_score.in_features model.roi_heads.box_predictor = FastRCNNPredictor(in_features_box, num_classes) in_features_mask = model.roi_heads.mask_predictor.conv5_mask.in_channels model.roi_heads.mask_predictor = MaskRCNNPredictor( in_channels=in_features_mask, dim_reduced=256, num_classes=num_classes ) backbone_param_groups = resnet_fpn.param_groups(model.backbone) else: model = MaskRCNN(backbone, num_classes=num_classes, **mask_rcnn_kwargs) backbone_param_groups = backbone.param_groups() patch_param_groups(model=model, backbone_param_groups=backbone_param_groups) if remove_internal_transforms: remove_internal_model_transforms(model) return model
def model(num_classes: int, backbone: Optional[TorchvisionBackboneConfig] = None, remove_internal_transforms: bool = True, **mask_rcnn_kwargs) -> nn.Module: """MaskRCNN model implemented by torchvision. # Arguments num_classes: Number of classes. backbone: Backbone model to use. Defaults to a resnet50_fpn model. remove_internal_transforms: The torchvision model internally applies transforms like resizing and normalization, but we already do this at the `Dataset` level, so it's safe to remove those internal transforms. **mask_rcnn_kwargs: Keyword arguments that internally are going to be passed to `torchvision.models.detection.mask_rcnn.MaskRCNN`. # Return A Pytorch `nn.Module`. """ if backbone is None: model = maskrcnn_resnet50_fpn(pretrained=True, pretrained_backbone=True, **mask_rcnn_kwargs) in_features_box = model.roi_heads.box_predictor.cls_score.in_features model.roi_heads.box_predictor = FastRCNNPredictor( in_features_box, num_classes) in_features_mask = model.roi_heads.mask_predictor.conv5_mask.in_channels model.roi_heads.mask_predictor = MaskRCNNPredictor( in_channels=in_features_mask, dim_reduced=256, num_classes=num_classes) resnet_fpn.patch_param_groups(model.backbone) else: model = MaskRCNN(backbone.backbone, num_classes=num_classes, **mask_rcnn_kwargs) patch_rcnn_param_groups(model=model) if remove_internal_transforms: remove_internal_model_transforms(model) return model
def model(num_classes: int, backbone: Optional[nn.Module] = None, remove_internal_transforms: bool = True, **faster_rcnn_kwargs) -> nn.Module: """ FasterRCNN model given by torchvision Args: num_classes (int): Number of classes. backbone (nn.Module): Backbone model to use. Defaults to a resnet50_fpn model. Return: nn.Module """ if backbone is None: model = maskrcnn_resnet50_fpn(pretrained=True, **faster_rcnn_kwargs) in_features_box = model.roi_heads.box_predictor.cls_score.in_features model.roi_heads.box_predictor = FastRCNNPredictor( in_features_box, num_classes) in_features_mask = model.roi_heads.mask_predictor.conv5_mask.in_channels model.roi_heads.mask_predictor = MaskRCNNPredictor( in_channels=in_features_mask, dim_reduced=256, num_classes=num_classes) backbone_param_groups = resnet_fpn.param_groups(model.backbone) else: model = MaskRCNN(backbone, num_classes=num_classes, **faster_rcnn_kwargs) backbone_param_groups = backbone.param_groups() patch_param_groups(model=model, backbone_param_groups=backbone_param_groups) if remove_internal_transforms: remove_internal_model_transforms(model) return model
def get_segmentation_learner(dls, number_classes, segmentation_type, architecture_name, backbone_name, loss_func=None, opt_func=Adam, lr=defaults.lr, splitter=trainable_params, cbs=None, pretrained=True, normalize=True, image_size=None, metrics=None, path=None, model_dir='models', wd=None, wd_bn_bias=False, train_bn=True, moms=(0.95, 0.85, 0.95)): """This function return a learner for the provided architecture and backbone Parameters: dls (DataLoader): the dataloader to use with the learner number_classes (int): the number of clases in the project. It should be >=2 segmentation_type (str): just `Semantic Segmentation` accepted for now architecture_name (str): name of the architecture. The following ones are supported: `unet`, `deeplabv3+`, `hrnet`, `maskrcnn` and `u2^net` backbone_name (str): name of the backbone loss_func (): loss function. opt_func (): opt function. lr (): learning rates splitter (): splitter function for freazing the learner cbs (List[cb]): list of callbacks pretrained (bool): it defines if a trained backbone is needed normalize (bool): image_size (int): REQUIRED for MaskRCNN. It indicates the desired size of the image. metrics (List[metric]): list of metrics path (): path parameter model_dir (str): the path in which save models wd (float): wieght decay wd_bn_bias (bool): train_bn (bool): moms (Tuple(float)): tuple of different momentuns Returns: learner: value containing the learner object """ number_classes_name = "" if number_classes == 2: number_classes_name = "binary" elif number_classes > 2: number_classes_name = "multiple" else: raise Exception("The number of classes must be >=2") check_architecture_configuration(number_classes=number_classes_name, segmentation_type=segmentation_type, architecture_name=architecture_name, backbone_name=backbone_name) learner = None if architecture_name == "unet": # TODO -> Revisar arch learner = unet_learner(dls=dls, arch=unet_backbone_name[backbone_name], metrics=metrics, wd=wd, loss_func=loss_func, opt_func=opt_func, lr=lr, splitter=splitter, cbs=cbs, path=path, model_dir=model_dir, wd_bn_bias=wd_bn_bias, train_bn=train_bn, pretrained=pretrained, normalize=normalize, moms=moms) elif architecture_name == "deeplabv3+": model = DeepLabV3Plus(backbone_name=backbone_name, nclass=number_classes, pretrained=pretrained) learner = Learner(dls=dls, model=model, loss_func=loss_func, opt_func=opt_func, lr=lr, splitter=splitter, cbs=cbs, metrics=metrics, path=path, model_dir=model_dir, wd=wd, wd_bn_bias=wd_bn_bias, train_bn=train_bn) elif architecture_name == "hrnet": model = HRNet(nclass=number_classes, backbone_name=backbone_name, pretrained=pretrained) learner = Learner(dls=dls, model=model, loss_func=loss_func, opt_func=opt_func, lr=lr, splitter=splitter, cbs=cbs, metrics=metrics, path=path, model_dir=model_dir, wd=wd, wd_bn_bias=wd_bn_bias, train_bn=train_bn) elif architecture_name == "maskrcnn": if image_size is None: raise Exception( "MaskRCNN need to define image_size. This values are for reescaling the image" ) model = maskrcnn_resnet50_fpn(num_classes=number_classes, min_size=image_size, max_size=image_size) learner = mask_rcnn.MaskRCNNLearner(dls=dls, model=model, loss_func=loss_func, opt_func=opt_func, lr=lr, splitter=splitter, cbs=cbs, metrics=metrics, path=path, model_dir=model_dir, wd=wd, wd_bn_bias=wd_bn_bias, train_bn=train_bn) elif architecture_name == "u2^net": model = None if backbone_name == "small": model = u2net.U2NETP(3, 1) elif backbone_name == "normal": model = u2net.U2NET(3, 1) learner = u2net.USquaredNetLearner(dls=dls, model=model, opt_func=opt_func, lr=lr, splitter=splitter, cbs=cbs, metrics=metrics, path=path, model_dir=model_dir, wd=wd, wd_bn_bias=wd_bn_bias, train_bn=train_bn) return learner
box_head = TwoMLPHead(in_channels=7 * 7 * 256, representation_size=128) box_predictor = FastRCNNPredictor(in_channels=128, num_classes=3) mask_roi_pool = torchvision.ops.MultiScaleRoIAlign(featmap_names=[0, 1, 2, 3], output_size=14, sampling_ratio=2) mask_predictor = MaskRCNNPredictor(in_channels=256, dim_reduced=256, num_classes=3) inference_args['box_head'] = box_head inference_args['rpn_anchor_generator'] = anchor_generator inference_args['mask_roi_pool'] = mask_roi_pool inference_args['mask_predictor'] = mask_predictor inference_args['box_predictor'] = box_predictor maskrcnn_model = maskrcnn_resnet50_fpn(pretrained=False, **inference_args) maskrcnn_model.load_state_dict(covid_detector_weights['model_weights']) maskrcnn_model.eval() maskrcnn_model = maskrcnn_model.to(device) thresholds = torch.arange(0.5, 1, 0.05).to(device) mean_aps_all_th = torch.zeros(thresholds.size()[0]).to(device) def compute_map(iou_th): mean_aps_this_th = torch.zeros(len(dataloader_eval_covid), dtype=torch.float) for id, b in enumerate(dataloader_eval_covid): X, y = b if device == torch.device('cuda'): X, y['labels'], y['boxes'], y['masks'] = X.to( device), y['labels'].to(device), y['boxes'].to(
'79': 'oven', '80': 'toaster', '81': 'sink', '82': 'refrigerator', '84': 'book', '85': 'clock', '86': 'vase', '87': 'scissors', '88': 'teddybear', '89': 'hair drier', '90': 'toothbrush' } SCORE_THRESHOLD = 0.5 model = maskrcnn_resnet50_fpn(pretrained=True) model.eval() img1 = skimage.io.imread("img2.jpeg") img2 = skimage.io.imread("img1.jpeg") x = [ torch.from_numpy(np.transpose(img1 / 255.0, axes=[2, 0, 1])).to(torch.float32), torch.from_numpy(np.transpose(img2 / 255.0, axes=[2, 0, 1])).to(torch.float32) ] out = model(x) filtered_idx = (out[0]['scores'] >= SCORE_THRESHOLD)
def date_for_filename(): tgt = time.localtime() year = str(tgt.tm_year) mon = "{:02}".format(tgt.tm_mon) day = "{:02}".format(tgt.tm_mday) hour = "{:02}".format(tgt.tm_hour) minute = "{:02}".format(tgt.tm_min) datestr = year + mon + day + '_' + hour + minute return datestr if __name__=="__main__": model = maskrcnn_resnet50_fpn(pretrained=False, \ pretrained_backbone=False,\ num_classes=49) device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu') #model.load_state_dict(torch.load('/home/karrington/FantomHD/CDC4-5-6-7-8-9(CleanRedo89Ers)_train-json-20210831_1301.pth')) model = model.to(device) # create test and train dataset dataset = myOwnDataset(root=train_data_dir, annotation=train_coco, transforms=get_transform(train=True) ) dataset_test = myOwnDataset(root=train_data_dir, annotation=test_coco, transforms=get_transform(train=False)
def main(args): utils.init_distributed_mode(args) print(args) device = torch.device(args.device) # Data loading code print("Loading data") dataset, num_classes = get_dataset(args.dataset, "train", get_transform(train=True), args.data_path) dataset_test, _ = get_dataset(args.dataset, "val", get_transform(train=False), args.data_path) print("Creating data loaders") if args.distributed: train_sampler = torch.utils.data.distributed.DistributedSampler( dataset) test_sampler = torch.utils.data.distributed.DistributedSampler( dataset_test) else: train_sampler = torch.utils.data.RandomSampler(dataset) test_sampler = torch.utils.data.SequentialSampler(dataset_test) if args.aspect_ratio_group_factor >= 0: group_ids = create_aspect_ratio_groups( dataset, k=args.aspect_ratio_group_factor) train_batch_sampler = GroupedBatchSampler(train_sampler, group_ids, args.batch_size) else: train_batch_sampler = torch.utils.data.BatchSampler(train_sampler, args.batch_size, drop_last=True) data_loader = torch.utils.data.DataLoader( dataset, batch_sampler=train_batch_sampler, num_workers=args.workers, collate_fn=utils.collate_fn) data_loader_test = torch.utils.data.DataLoader(dataset_test, batch_size=1, sampler=test_sampler, num_workers=args.workers, collate_fn=utils.collate_fn) print("Creating model") kwargs = { "trainable_backbone_layers": args.trainable_backbone_layers, } if (args.tfidf): tfidf = pd.read_csv(f'../{args.dataset}_files/idf_{num_classes}.csv')[ args.tfidf] tfidf = torch.tensor(tfidf, device='cuda').unsqueeze(0) else: tfidf = torch.ones(num_classes, device='cuda', dtype=torch.float).unsqueeze(0) if "rcnn" in args.model: if args.rpn_score_thresh is not None: kwargs["rpn_score_thresh"] = args.rpn_score_thresh if args.model == 'fasterrcnn_resnet50_fpn': model = frcnn.fasterrcnn_resnet50_fpn(pretrained=args.pretrained, num_classes=num_classes, tfidf=tfidf) elif args.model == 'retinanet_resnet50_fpn': model = retinanet.retinanet_resnet50_fpn(pretrained=args.pretrained, num_classes=num_classes, tfidf=tfidf) elif args.model == 'maskrcnn_resnet50_fpn': model = mask_rcnn.maskrcnn_resnet50_fpn(pretrained=args.pretrained, num_classes=num_classes, tfidf=tfidf) # model = torchvision.models.detection.__dict__[args.model](num_classes=num_classes, pretrained=args.pretrained, # **kwargs) model.to(device) model_without_ddp = model if args.distributed: model = torch.nn.parallel.DistributedDataParallel( model, device_ids=[args.gpu]) model_without_ddp = model.module params = [p for p in model.parameters() if p.requires_grad] optimizer = torch.optim.SGD(params, lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay) # lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=args.lr_step_size, gamma=args.lr_gamma) lr_scheduler = torch.optim.lr_scheduler.MultiStepLR( optimizer, milestones=args.lr_steps, gamma=args.lr_gamma) if args.resume: checkpoint = torch.load(args.resume, map_location='cpu') model_without_ddp.load_state_dict(checkpoint['model']) optimizer.load_state_dict(checkpoint['optimizer']) lr_scheduler.load_state_dict(checkpoint['lr_scheduler']) args.start_epoch = checkpoint['epoch'] + 1 if args.test_only: evaluate(model, data_loader_test, device=device) return print("Start training") start_time = time.time() for epoch in range(args.start_epoch, args.epochs): if args.distributed: train_sampler.set_epoch(epoch) train_one_epoch(model, optimizer, data_loader, device, epoch, args.print_freq) lr_scheduler.step() if args.output_dir: utils.save_on_master( { 'model': model_without_ddp.state_dict(), 'optimizer': optimizer.state_dict(), 'lr_scheduler': lr_scheduler.state_dict(), 'args': args, 'epoch': epoch }, os.path.join(args.output_dir, 'model_{}.pth'.format(epoch))) # evaluate after every epoch evaluate(model, data_loader_test, device=device) total_time = time.time() - start_time total_time_str = str(datetime.timedelta(seconds=int(total_time))) print('Training time {}'.format(total_time_str))