def __init__(self, num_classes, backbone_fn, chip_size=224): super().__init__() if getattr(backbone_fn, '_is_multispectral', False): self.backbone = create_body(backbone_fn, pretrained=True, cut=_get_backbone_meta( backbone_fn.__name__)['cut']) else: self.backbone = create_body(backbone_fn, pretrained=True) backbone_name = backbone_fn.__name__ ## Support for different backbones if "densenet" in backbone_name or "vgg" in backbone_name: hookable_modules = list(self.backbone.children())[0] else: hookable_modules = list(self.backbone.children()) if "vgg" in backbone_name: modify_dilation_index = -5 else: modify_dilation_index = -2 if backbone_name == 'resnet18' or backbone_name == 'resnet34': module_to_check = 'conv' else: module_to_check = 'conv2' ## Hook at the index where we need to get the auxillary logits out self.hook = hook_output(hookable_modules[modify_dilation_index]) custom_idx = 0 for i, module in enumerate(hookable_modules[modify_dilation_index:]): dilation = 2 * (i + 1) padding = 2 * (i + 1) for n, m in module.named_modules(): if module_to_check in n: m.dilation, m.padding, m.stride = (dilation, dilation), ( padding, padding), (1, 1) elif 'downsample.0' in n: m.stride = (1, 1) if "vgg" in backbone_fn.__name__: if isinstance(module, nn.Conv2d): dilation = 2 * (custom_idx + 1) padding = 2 * (custom_idx + 1) module.dilation, module.padding, module.stride = ( dilation, dilation), (padding, padding), (1, 1) custom_idx += 1 ## returns the size of various activations feature_sizes = model_sizes(self.backbone, size=(chip_size, chip_size)) ## Geting the number of channel persent in stored activation inside of the hook num_channels_aux_classifier = self.hook.stored.shape[1] ## Get number of channels in the last layer num_channels_classifier = feature_sizes[-1][1] self.classifier = DeepLabHead(num_channels_classifier, num_classes) self.aux_classifier = FCNHead(num_channels_aux_classifier, num_classes)
def get_model(self, pretrained=True): if self.model_number == 18: body = create_body( resnet18, pretrained=pretrained, n_in=self.input_channels, cut=-2 ) else: body = create_body( resnet34, pretrained=pretrained, n_in=self.input_channels, cut=-2 ) net_G = DynamicUnet( body, self.output_channels, (self.image_size, self.image_size) ) return net_G
def __init__(self, base_arch, no_diseases, dropout=0.5, init=nn.init.kaiming_normal_): super(CNNPretrainedModel, self).__init__() self.body = create_body(base_arch, pretrained=True) nf = num_features_model(nn.Sequential(*self.body.children())) * 2 self.disease_head = create_head(nf, no_diseases, ps=0.5, concat_pool=True, bn_final=False) #self.age_head = create_head(nf, 1, ps=0.5, concat_pool=True, bn_final=False) self.gender_head = create_head(nf, 2, ps=0.5, concat_pool=True, bn_final=False) #self.projection_head = create_head(nf, 3, ps=0.5, concat_pool=True, bn_final=False) self.disease_model = nn.Sequential(self.body, self.disease_head) self.meta = cnn_config(base_arch) self.split(self.meta['split']) self.freeze() apply_init(self.disease_head, init) #apply_init(self.age_head, init) apply_init(self.gender_head, init)
def __init__(self, data, scales=None, ratios=None, backbone=None, pretrained_path=None, *args, **kwargs): # Set default backbone to be 'resnet50' if backbone is None: backbone = models.resnet50 super().__init__(data, backbone) n_bands = len(getattr(self._data, '_extract_bands', [0, 1, 2])) _backbone = self._backbone if hasattr(self, '_orig_backbone'): _backbone = self._orig_backbone # Check if a backbone provided is compatible, use resnet50 as default if not self._check_backbone_support(_backbone): raise Exception( f"Enter only compatible backbones from {', '.join(self.supported_backbones)}" ) self.name = "RetinaNet" self._code = code self.scales = ifnone(scales, [1, 2**(-1 / 3), 2**(-2 / 3)]) self.ratios = ifnone(ratios, [1 / 2, 1, 2]) self._n_anchors = len(self.scales) * len(self.ratios) self._data = data self._chip_size = (data.chip_size, data.chip_size) # Cut-off the backbone before the penultimate layer self._encoder = create_body(self._backbone, -2) # Initialize the model, loss function and the Learner object self._model = RetinaNetModel(self._encoder, n_classes=data.c - 1, final_bias=-4, chip_size=self._chip_size, n_anchors=self._n_anchors, n_bands=n_bands) self._loss_f = RetinaNetFocalLoss(sizes=self._model.sizes, scales=self.scales, ratios=self.ratios) self.learn = Learner(data, self._model, loss_func=self._loss_f) self.learn.split([self._model.encoder[6], self._model.c5top5]) self.learn.freeze() if pretrained_path is not None: self.load(str(pretrained_path)) self._arcgis_init_callback() # make first conv weights learnable
def __init__(self): super().__init__() concat_pool = True self.encoder = create_body(models.resnet18, True) nf = num_features_model(nn.Sequential( *(self.encoder).children())) * (2 if concat_pool else 1) self.domain_head_1, self.domain_head_2 = create_domain_head( nf, 2, lin_ftrs=[256]) self.class_head_1, self.class_head_2 = create_classification_head( nf, 2, lin_ftrs=[256]) self.alpha = 0
def build_res_unet(n_input=1, n_output=2, size=256): ''' Creates a generator by building a Unet model with a Resnet18 backbone Args: n_input (int): Corresponds to the L feature dimension in L*a*b* color space n_output (int): Corresponds to the ab feature dimensions in L*a*b* color space size (int): height/width of image Returns: Generator model ''' device = torch.device("cuda" if torch.cuda.is_available() else "cpu") body = create_body(resnet18, pretrained=True, n_in=n_input, cut=-2) net_G = DynamicUnet(body, n_output, (size, size)).to(device) return net_G
def unet_learner_wide( data: DataBunch, arch: Callable, pretrained: bool = True, blur_final: bool = True, norm_type: Optional[NormType] = NormType, split_on: Optional[SplitFuncOrIdxList] = None, blur: bool = False, self_attention: bool = False, y_range: Optional[Tuple[float, float]] = None, last_cross: bool = True, bottle: bool = False, nf_factor: int = 1, **kwargs: Any ) -> Learner: "Build Unet learner from `data` and `arch`." meta = cnn_config(arch) body = create_body(arch, pretrained) model = to_device( DynamicUnetWide( body, n_classes=data.c, blur=blur, blur_final=blur_final, self_attention=self_attention, y_range=y_range, norm_type=norm_type, last_cross=last_cross, bottle=bottle, nf_factor=nf_factor, ), data.device, ) learn = Learner(data, model, **kwargs) learn.split(ifnone(split_on, meta['split'])) if pretrained: learn.freeze() apply_init(model[2], nn.init.kaiming_normal_) return learn
def unet4mri_learner(data: DataBunch, arch: Callable, pretrained: bool = True, blur_final: bool = True, norm_type: Optional[NormType] = None, split_on: Optional[SplitFuncOrIdxList] = None, blur: bool = False, self_attention: bool = False, y_range: Optional[Tuple[float, float]] = None, last_cross: bool = True, bottle: bool = False, cut: Union[int, Callable] = None, **learn_kwargs: Any) -> Learner: "Build Unet learner from `data` and `arch`." meta = ln.cnn_config(arch) body = ln.create_body(arch, pretrained, cut) try: size = data.train_ds[0][0].size except: size = next(iter(data.train_dl))[0].shape[-2:] model = to_device( MyDynamicUnet(body, bs=data.batch_size, n_classes=data.c, img_size=size, blur=blur, blur_final=blur_final, self_attention=self_attention, y_range=y_range, norm_type=norm_type, last_cross=last_cross, bottle=bottle), data.device) learn = Learner(data, model, **learn_kwargs) learn.split(ifnone(split_on, meta['split'])) if pretrained: learn.freeze() apply_init(model[2], nn.init.kaiming_normal_) return learn
import argparse if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument( '-i','--image', type=str, default='data/images/Test/xkObNU.jpg' ) parser.add_argument( '-m', '--model', default='model/color_gen.pt' ) args = parser.parse_args() # Load model body = create_body(resnet18, pretrained=True, n_in=1, cut=-2) net_G = DynamicUnet(body, 2, (256, 256)).to('cuda') net_G.load_state_dict(torch.load(args.model)) net_G.to('cuda') net_G.eval() img = cv2.imread(args.image) img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) lab_img = rgb2lab(img).astype(np.float32) lab_img = T.ToTensor()(lab_img) L = lab_img[[0], ...] / 50. - 1. # Between -1 and 1 with torch.no_grad(): pred_ab = net_G(L.unsqueeze(0).to('cuda'))
def __init__(self, data, backbone=None, pretrained_path=None): super().__init__(data, backbone) if self._is_multispectral: self._backbone_ms = self._backbone self._backbone = self._orig_backbone scaled_mean_values = data._scaled_mean_values[ data._extract_bands].tolist() scaled_std_values = data._scaled_std_values[ data._extract_bands].tolist() if backbone is None: self._backbone = models.resnet50 elif type(backbone) is str: self._backbone = getattr(models, backbone) else: self._backbone = backbone if not self._check_backbone_support(self._backbone): raise Exception( f"Enter only compatible backbones from {', '.join(self.supported_backbones)}" ) self._code = instance_detector_prf if self._backbone.__name__ is 'resnet50': model = models.detection.maskrcnn_resnet50_fpn( pretrained=True, min_size=1.5 * data.chip_size, max_size=2 * data.chip_size) if self._is_multispectral: model.backbone = _change_tail(model.backbone, data) model.transform.image_mean = scaled_mean_values model.transform.image_std = scaled_std_values elif self._backbone.__name__ in ['resnet18', 'resnet34']: if self._is_multispectral: backbone_small = create_body(self._backbone_ms, cut=_get_backbone_meta( backbone_fn.__name__)['cut']) backbone_small.out_channels = 512 model = models.detection.MaskRCNN( backbone_small, 91, min_size=1.5 * data.chip_size, max_size=2 * data.chip_size, image_mean=scaled_mean_values, image_std=scaled_std_values) else: backbone_small = create_body(self._backbone) backbone_small.out_channels = 512 model = models.detection.MaskRCNN(backbone_small, 91, min_size=1.5 * data.chip_size, max_size=2 * data.chip_size) else: backbone_fpn = resnet_fpn_backbone(self._backbone.__name__, True) if self._is_multispectral: backbone_fpn = _change_tail(backbone_fpn, data) model = models.detection.MaskRCNN( backbone_fpn, 91, min_size=1.5 * data.chip_size, max_size=2 * data.chip_size, image_mean=scaled_mean_values, image_std=scaled_std_values) else: model = models.detection.MaskRCNN(backbone_fpn, 91, min_size=1.5 * data.chip_size, max_size=2 * data.chip_size) in_features = model.roi_heads.box_predictor.cls_score.in_features model.roi_heads.box_predictor = FastRCNNPredictor(in_features, data.c) in_features_mask = model.roi_heads.mask_predictor.conv5_mask.in_channels hidden_layer = 256 model.roi_heads.mask_predictor = MaskRCNNPredictor( in_features_mask, hidden_layer, data.c) self.learn = Learner(data, model, loss_func=mask_rcnn_loss) self.learn.callbacks.append(train_callback(self.learn)) self.learn.model = self.learn.model.to(self._device) self.learn.c_device = self._device # fixes for zero division error when slice is passed idx = 27 if self._backbone.__name__ in ['resnet18', 'resnet34']: idx = self._freeze() self.learn.layer_groups = split_model_idx(self.learn.model, [idx]) self.learn.create_opt(lr=3e-3) # make first conv weights learnable self._arcgis_init_callback() if pretrained_path is not None: self.load(pretrained_path) if self._is_multispectral: self._orig_backbone = self._backbone self._backbone = self._backbone_ms
def __init__(self, data, grids=[4, 2, 1], zooms=[0.7, 1., 1.3], ratios=[[1., 1.], [1., 0.5], [0.5, 1.]], backbone=None, drop=0.3, bias=-4., focal_loss=False, pretrained_path=None, location_loss_factor=None): super().__init__() self._device = torch.device( 'cuda') if torch.cuda.is_available() else torch.device('cpu') # assert (location_loss_factor is not None) or ((location_loss_factor > 0) and (location_loss_factor < 1)), if location_loss_factor is not None: if not ((location_loss_factor > 0) and (location_loss_factor < 1)): raise Exception( '`location_loss_factor` should be greater than 0 and less than 1' ) self.location_loss_factor = location_loss_factor if not HAS_FASTAI: _raise_fastai_import_error() if backbone is None: self._backbone = models.resnet34 elif type(backbone) is str: self._backbone = getattr(models, backbone) else: self._backbone = backbone self._create_anchors(grids, zooms, ratios) feature_sizes = model_sizes(create_body(self._backbone), size=(data.chip_size, data.chip_size)) num_features = feature_sizes[-1][-1] num_channels = feature_sizes[-1][1] ssd_head = SSDHead(grids, self._anchors_per_cell, data.c, num_features=num_features, drop=drop, bias=bias, num_channels=num_channels) self._data = data self.learn = create_cnn(data=data, arch=self._backbone, custom_head=ssd_head) self.learn.model = self.learn.model.to(self._device) if pretrained_path is not None: self.load(pretrained_path) if focal_loss: self._loss_f = FocalLoss(data.c) else: self._loss_f = BCE_Loss(data.c) self.learn.loss_func = self._ssd_loss
def _pspnet_unet(num_classes, backbone_fn, chip_size=224, pyramid_sizes=(1, 2, 3, 6), pretrained=True): """ Function which returns PPM module attached to backbone which is then used to form the Unet. """ if getattr(backbone_fn, '_is_multispectral', False): backbone = create_body(backbone_fn, pretrained=pretrained, cut=_get_backbone_meta(backbone_fn.__name__)['cut']) else: backbone = create_body(backbone_fn, pretrained=pretrained) backbone_name = backbone_fn.__name__ ## Support for different backbones if "densenet" in backbone_name or "vgg" in backbone_name: hookable_modules = list(backbone.children())[0] else: hookable_modules = list(backbone.children()) if "vgg" in backbone_name: modify_dilation_index = -5 else: modify_dilation_index = -2 if backbone_name == 'resnet18' or backbone_name == 'resnet34': module_to_check = 'conv' else: module_to_check = 'conv2' custom_idx = 0 for i, module in enumerate(hookable_modules[modify_dilation_index:]): dilation = 2 * (i + 1) padding = 2 * (i + 1) # padding = 1 for n, m in module.named_modules(): if module_to_check in n: m.dilation, m.padding, m.stride = (dilation, dilation), (padding, padding), (1, 1) elif 'downsample.0' in n: m.stride = (1, 1) if "vgg" in backbone_fn.__name__: if isinstance(module, nn.Conv2d): dilation = 2 * (custom_idx + 1) padding = 2 * (custom_idx + 1) module.dilation, module.padding, module.stride = (dilation, dilation), (padding, padding), (1, 1) custom_idx += 1 ## returns the size of various activations feature_sizes = model_sizes(backbone, size=(chip_size, chip_size)) ## Get number of channels in the last layer num_channels = feature_sizes[-1][1] penultimate_channels = num_channels / len(pyramid_sizes) ppm = _PyramidPoolingModule(num_channels, int(penultimate_channels), pyramid_sizes) in_final = int(penultimate_channels) * len(pyramid_sizes) + num_channels # Reduce channel size after pyramid pooling module to avoid CUDA OOM error. final_conv = nn.Conv2d(in_channels=in_final, out_channels=512, kernel_size=3, padding=1) ## To make Dynamic Unet work as it expects a backbone which can be indexed. if "densenet" in backbone_name or "vgg" in backbone_name: backbone = backbone[0] layers = [*backbone, ppm, final_conv] return nn.Sequential(*layers)
def __init__(self, num_classes, backbone_fn, chip_size=224, pyramid_sizes=(1, 2, 3, 6), pretrained=True): super(PSPNet, self).__init__() if getattr(backbone_fn, '_is_multispectral', False): self.backbone = create_body(backbone_fn, pretrained=pretrained, cut=_get_backbone_meta(backbone_fn.__name__)['cut']) else: self.backbone = create_body(backbone_fn, pretrained=pretrained) backbone_name = backbone_fn.__name__ ## Support for different backbones if "densenet" in backbone_name or "vgg" in backbone_name: hookable_modules = list(self.backbone.children())[0] else: hookable_modules = list(self.backbone.children()) if "vgg" in backbone_name: modify_dilation_index = -5 else: modify_dilation_index = -2 if backbone_name == 'resnet18' or backbone_name == 'resnet34': module_to_check = 'conv' else: module_to_check = 'conv2' ## Hook at the index where we need to get the auxillary logits out self.hook = hook_output(hookable_modules[modify_dilation_index]) custom_idx = 0 for i, module in enumerate(hookable_modules[modify_dilation_index:]): dilation = 2 * (i + 1) padding = 2 * (i + 1) for n, m in module.named_modules(): if module_to_check in n: m.dilation, m.padding, m.stride = (dilation, dilation), (padding, padding), (1, 1) elif 'downsample.0' in n: m.stride = (1, 1) if "vgg" in backbone_fn.__name__: if isinstance(module, nn.Conv2d): dilation = 2 * (custom_idx + 1) padding = 2 * (custom_idx + 1) module.dilation, module.padding, module.stride = (dilation, dilation), (padding, padding), (1, 1) custom_idx += 1 ## returns the size of various activations feature_sizes = model_sizes(self.backbone, size=(chip_size, chip_size)) ## Geting the stored parameters inside of the hook aux_in_channels = self.hook.stored.shape[1] ## Get number of channels in the last layer num_channels = feature_sizes[-1][1] penultimate_channels = num_channels / len(pyramid_sizes) self.ppm = _PyramidPoolingModule(num_channels, int(penultimate_channels), pyramid_sizes) self.final = nn.Sequential( ## To handle case when the length of pyramid_sizes is odd nn.Conv2d(int(penultimate_channels) * len(pyramid_sizes) + num_channels, math.ceil(penultimate_channels), kernel_size=3, padding=1, bias=False), nn.BatchNorm2d(math.ceil(penultimate_channels)), nn.ReLU(inplace=True), nn.Dropout(0.1), nn.Conv2d(math.ceil(penultimate_channels), num_classes, kernel_size=1) ) self.aux_logits = nn.Conv2d(aux_in_channels, num_classes, kernel_size=1) initialize_weights(self.aux_logits) initialize_weights(self.ppm, self.final)
def __init__(self, data, grids=None, zooms=[1.], ratios=[[1., 1.]], backbone=None, drop=0.3, bias=-4., focal_loss=False, pretrained_path=None, location_loss_factor=None, ssd_version=2, backend='pytorch'): super().__init__(data, backbone) self._backend = backend if self._backend == 'tensorflow': self._intialize_tensorflow(data, grids, zooms, ratios, backbone, drop, bias, pretrained_path, location_loss_factor) else: # assert (location_loss_factor is not None) or ((location_loss_factor > 0) and (location_loss_factor < 1)), if not ssd_version in [1, 2]: raise Exception("ssd_version can be only [1,2]") if location_loss_factor is not None: if not ((location_loss_factor > 0) and (location_loss_factor < 1)): raise Exception( '`location_loss_factor` should be greater than 0 and less than 1' ) self.location_loss_factor = location_loss_factor self._code = code self.ssd_version = ssd_version backbone_cut = None backbone_split = None if hasattr(self, '_orig_backbone'): self._backbone_ms = self._backbone self._backbone = self._orig_backbone _backbone_meta = cnn_config(self._orig_backbone) backbone_cut = _backbone_meta['cut'] backbone_split = _backbone_meta['split'] if backbone is None: self._backbone = models.resnet34 backbone_name = 'res' elif type(backbone) is str: self._backbone = getattr(models, backbone) backbone_name = backbone[:3] else: self._backbone = backbone backbone_name = 'custom' if not self._check_backbone_support(self._backbone): raise Exception( f"Enter only compatible backbones from {', '.join(self.supported_backbones)}" ) if self._backbone == models.mobilenet_v2: backbone_cut = -1 backbone_split = _mobilenet_split if ssd_version == 1: if grids == None: grids = [4, 2, 1] self._create_anchors(grids, zooms, ratios) feature_sizes = model_sizes(create_body(self._backbone, cut=backbone_cut), size=(data.chip_size, data.chip_size)) num_features = feature_sizes[-1][-1] num_channels = feature_sizes[-1][1] ssd_head = SSDHead(grids, self._anchors_per_cell, data.c, num_features=num_features, drop=drop, bias=bias, num_channels=num_channels) elif ssd_version == 2: # find bounding boxes height and width if grids is None: logger.info("Computing optimal grid size...") hw = data.height_width hw = np.array(hw) # find most suitable centroids for dataset centroid = kmeans(hw, 1) avg = avg_iou(hw, centroid) for num_anchor in range(2, 5): new_centroid = kmeans(hw, num_anchor) new_avg = avg_iou(hw, new_centroid) if (new_avg - avg) < 0.05: break avg = new_avg centroid = new_centroid.copy() # find grid size grids = list( map( int, map( round, data.chip_size / np.sort(np.max(centroid, axis=1))))) grids = list(set(grids)) grids.sort(reverse=True) if grids[-1] == 0: grids[-1] = 1 grids = list(set(grids)) self._create_anchors(grids, zooms, ratios) feature_sizes = model_sizes(create_body(self._backbone, cut=backbone_cut), size=(data.chip_size, data.chip_size)) num_features = feature_sizes[-1][-1] num_channels = feature_sizes[-1][1] if grids[0] > 8 and abs(num_features - grids[0] ) > 4 and backbone_name == 'res': num_features = feature_sizes[-2][-1] num_channels = feature_sizes[-2][1] backbone_cut = -3 ssd_head = SSDHeadv2(grids, self._anchors_per_cell, data.c, num_features=num_features, drop=drop, bias=bias, num_channels=num_channels) else: raise Exception('SSDVersion can only be 1 or 2') if hasattr(self, '_backbone_ms'): self._orig_backbone = self._backbone self._backbone = self._backbone_ms self.learn = cnn_learner(data=data, base_arch=self._backbone, cut=backbone_cut, split_on=backbone_split, custom_head=ssd_head) self._arcgis_init_callback() # make first conv weights learnable self.learn.model = self.learn.model.to(self._device) if focal_loss: self._loss_f = FocalLoss(data.c) else: self._loss_f = BCE_Loss(data.c) self.learn.loss_func = self._ssd_loss _set_multigpu_callback(self) if pretrained_path is not None: self.load(pretrained_path)
def build_res_unet(n_input=1, n_output=2, size=224): body = create_body(resnet18, pretrained=False, n_in=n_input, cut=-2) net_G = DynamicUnet(body, n_output, (size, size)) return net_G
def model(): body = create_body(resnet18, pretrained=False) for param in body.parameters(): param.requires_grad = False return DynamicUnet(body, 10)
def build_generator(n_inputs, n_outputs, core): body = create_body(core, pretrained=True, n_in=n_inputs, cut=-2) return DynamicUnet(body, n_outputs, (SIZE, SIZE))
def build_res_unet(n_input=1, n_output=2, size=256): device = torch.device("cuda" if torch.cuda.is_available() else "cpu") body = create_body(resnet18, pretrained=True, n_in=n_input, cut=-2) net_G = DynamicUnet(body, n_output, (size, size)).to(device) return net_G