def prepare_train_img(self, idx): img_info = self.img_infos[idx] # load image img = mmcv.imread(osp.join(self.img_prefix, img_info['filename'])) ann = self.get_ann_info(idx) gt_bboxes = ann['bboxes'] gt_labels = ann['labels'] if self.with_ignore: gt_bboxes_ignore = ann['bboxes_ignore'] # skip the image if there is no valid gt bbox if len(gt_bboxes) == 0: return None # aug 1: apply extra augmentation if self.extra_aug is not None: img, gt_bboxes, gt_labels = \ self.extra_aug(img, gt_bboxes, gt_labels) # aug 2: apply ordinary augmentations: flipping flip = True if np.random.rand() < self.flip_ratio else False # aug 3: apply ordinary augmentations: scaling img_scale = random_scale(self.img_scales, self.multiscale_mode) img, img_shape, pad_shape, scale_factor = \ self.img_transform( img=img, scale=img_scale, flip=flip, pad_val=self.pad_values, keep_ratio=self.resize_keep_ratio ) img = img.copy() gt_bboxes = self.bbox_transform(bboxes=gt_bboxes, img_shape=img_shape, scale_factor=scale_factor, flip=flip) if self.with_ignore: gt_bboxes_ignore = self.bbox_transform(bboxes=gt_bboxes_ignore, img_shape=img_shape, scale_factor=scale_factor, flip=flip) img_meta = dict(ori_shape=(img_info['height'], img_info['width'], 3), img_shape=img_shape, pad_shape=pad_shape, scale_factor=scale_factor, flip=flip) data = dict(img=DataContainer(to_tensor(img), stack=True), img_meta=DataContainer(data=img_meta, cpu_only=True), gt_bboxes=DataContainer(data=to_tensor(gt_bboxes))) if self.with_ignore: data['gt_bboxes_ignore'] = DataContainer( data=to_tensor(gt_bboxes_ignore)) if self.with_label: data['gt_labels'] = DataContainer(to_tensor(gt_labels)) return data
def prepare_train_img(self, idx): img_info = self.img_infos[idx] # load image img = mmcv.imread(osp.join(self.img_prefix, img_info['filename'])) ann = self.get_ann_info(idx) gt_labels = np.zeros([len(self.CLASSES)], dtype=np.float32) for index, item in enumerate(self.CLASSES): gt_labels[index] = int(ann[item] == 1) # aug 1: apply extra augmentation if self.extra_aug is not None: img, gt_bboxes, gt_labels = \ self.extra_aug(img, None, gt_labels) # aug 2: apply ordinary augmentations: flipping flip = True if np.random.rand() < self.flip_ratio else False # aug 3: apply ordinary augmentations: scaling img_scale = random_scale(self.img_scales, self.multiscale_mode) img, img_shape, pad_shape, scale_factor = \ self.img_transform( img=img, scale=img_scale, flip=flip, pad_val=self.pad_values, keep_ratio=self.resize_keep_ratio ) img = img.copy() img_meta = dict(ori_shape=(img_info['height'], img_info['width'], 3), img_shape=img_shape, pad_shape=pad_shape, scale_factor=scale_factor, flip=flip) data = dict( img=DataContainer(to_tensor(img), stack=True), img_meta=DataContainer(data=img_meta, cpu_only=True), ) if self.with_label: data['gt_labels'] = DataContainer(to_tensor(gt_labels), stack=False) return data
def prepare_test_img(self, idx): """Prepare an image for testing""" img_info = self.img_infos[idx] # load image img = mmcv.imread(osp.join(self.img_prefix, img_info['filename'])) img, img_shape, pad_shape, scale_factor = \ self.img_transform( img=img, scale=self.img_scales[0], flip=False, pad_val=self.pad_values, keep_ratio=self.resize_keep_ratio ) img_meta = dict( ori_shape=(img_info['height'], img_info['width'], 3), img_shape=img_shape, pad_shape=pad_shape, scale_factor=scale_factor, flip=False ) data = dict( img=DataContainer(to_tensor(img), stack=True), img_meta=DataContainer(img_meta, cpu_only=True) ) return data
def seq_collate(batch, samples_per_gpu=1): """ Modified from mmcv.collate Puts each data field into a tensor/DataContainer with outer dimension Seq Length. Extend default_collate to add support for :type:`~mmcv.parallel.DataContainer`. There are 3 cases. 1. cpu_only = True, e.g., meta data 2. cpu_only = False, stack = True, e.g., images tensors 3. cpu_only = False, stack = False, e.g., gt bboxes """ if not isinstance(batch, collections.Sequence): raise TypeError("{} is not supported.".format(batch.dtype)) if isinstance(batch[0], DataContainer): assert len(batch) % samples_per_gpu == 0 if batch[0].cpu_only: stacked = [sample.data for sample in batch] return DataContainer(stacked, batch[0].stack, batch[0].padding_value, cpu_only=True) elif batch[0].stack: assert isinstance(batch[0].data, torch.Tensor) if batch[0].pad_dims is not None: ndim = batch[0].dim() assert ndim > batch[0].pad_dims max_shape = [0 for _ in range(batch[0].pad_dims)] for dim in range(1, batch[0].pad_dims + 1): max_shape[dim - 1] = batch[0].size(-dim) for sample in batch: for dim in range(0, ndim - batch[0].pad_dims): assert batch[0].size(dim) == sample.size(dim) for dim in range(1, batch[0].pad_dims + 1): max_shape[dim - 1] = max(max_shape[dim - 1], sample.size(-dim)) padded_samples = [] for sample in batch: pad = [0 for _ in range(batch[0].pad_dims * 2)] for dim in range(1, batch[0].pad_dims + 1): pad[2 * dim - 1] = max_shape[dim - 1] - sample.size(-dim) padded_samples.append( F.pad(sample.data, pad, value=sample.padding_value)) stacked = default_collate(padded_samples) elif batch[0].pad_dims is None: stacked = default_collate([sample.data for sample in batch]) else: raise ValueError( 'pad_dims should be either None or integers (1-3)') else: stacked = [sample.data for sample in batch] return DataContainer(stacked, batch[0].stack, batch[0].padding_value) elif isinstance(batch[0], collections.Sequence): transposed = zip(*batch) return [ seq_collate(samples, samples_per_gpu) for samples in transposed ] elif isinstance(batch[0], collections.Mapping): collate_ret = { key: seq_collate([d[key] for d in batch], samples_per_gpu) for key in batch[0] } return collate_ret else: return default_collate(batch)
def main(): # load image img = mmcv.imread(IMG_PATH) img_height = img.shape[0] img_width = img.shape[1] # image pre-processing img_transform = ImageTransform( mean=IMG_TRANSFORM_CONFIG['mean'], std=IMG_TRANSFORM_CONFIG['std'], to_rgb=IMG_TRANSFORM_CONFIG['to_rgb'], ) img, img_shape, pad_shape, scale_factor = \ img_transform( img=img, scale=IMG_SCALE, flip=False, pad_val=IMG_TRANSFORM_CONFIG['pad_values'], keep_ratio=IMG_TRANSFORM_CONFIG['resize_keep_ratio'] ) img_meta = dict(ori_shape=(img_height, img_width, 3), img_shape=img_shape, pad_shape=pad_shape, scale_factor=scale_factor, flip=False) data = dict(img=DataContainer(to_tensor(img), stack=True), img_meta=DataContainer(img_meta, cpu_only=True)) # define the model model = SSDDetector( # basic input_size=IMG_SCALE, num_classes=NUM_CLASS, in_channels=(512, 1024, 512, 256, 256), use_dropout=False, dropout_rate=None, # anchor generate anchor_ratios=([1 / 2.0, 1.0, 2.0], [1 / 3.0, 1 / 2.0, 1.0, 2.0, 3.0], [1 / 3.0, 1 / 2.0, 1.0, 2.0, 3.0], [1 / 3.0, 1 / 2.0, 1.0, 2.0, 3.0], [1 / 2.0, 1.0, 2.0]), anchor_strides=((16, 16), (16, 16), (30, 30), (60, 60), (100, 100)), basesizes=((12, 12), (16, 16), (24, 24), (30, 30), (36, 36)), allowed_border=-1, # regression target_means=(.0, .0, .0, .0), target_stds=(0.1, 0.1, 0.2, 0.2), # box assign pos_iou_thr=0.5, neg_iou_thr=0.5, min_pos_iou=0., gt_max_assign_all=False, # sampling sampling=False, # balancing the loss neg_pos_ratio=3, # loss smoothl1_beta=1., # inference nms nms_pre=-1, score_thr=0.02, min_size=100.0, max_scale_ratio=10.0, nms_cfg=['nms', 0.45, None], max_per_img=200, # device device='cpu', ) # load checkpoint _ = load_checkpoint(model=model, filename=CHECKPOINT_FILE, map_location='cpu', strict=True, logger=None) # parallelize model model.eval() # results and progress bar # inference the data with torch.no_grad(): result = model(is_test=True, rescale=True, img=data['img'].data.unsqueeze(0), img_meta=(data['img_meta'].data, )) show_one_image(result[0], IMG_PATH, SAVE_PATH)
def collate(batch, samples_per_gpu=1): """Puts each data field into a tensor/DataContainer with outer dimension batch size. Extend default_collate to add support for :type:`~mmcv.parallel.DataContainer`. There are 3 cases. 1. cpu_only = True, e.g., meta data 2. cpu_only = False, stack = True, e.g., images tensors 3. cpu_only = False, stack = False, e.g., gt bboxes """ if not isinstance(batch, collections.Sequence): raise TypeError("{} is not supported.".format(batch.dtype)) if isinstance(batch[0], DataContainer): assert len(batch) % samples_per_gpu == 0 stacked = [] if batch[0].cpu_only: for i in range(0, len(batch), samples_per_gpu): stacked.append( [sample.data for sample in batch[i:i + samples_per_gpu]]) return DataContainer(stacked, batch[0].stack, batch[0].padding_value, cpu_only=True) elif batch[0].stack: for i in range(0, len(batch), samples_per_gpu): assert isinstance(batch[i].data, torch.Tensor) if batch[i].dim() == 3: # TODO: handle tensors other than 3d assert batch[i].dim() == 3 #选择一个batch里面最大的 h w c, h, w = batch[0].size() for sample in batch[i:i + samples_per_gpu]: assert c == sample.size(0) h = max(h, sample.size(1)) w = max(w, sample.size(2)) padded_samples = [ F.pad( sample.data, (0, w - sample.size(2), 0, h - sample.size(1)), #在右侧和下册pad value=sample.padding_value) for sample in batch[i:i + samples_per_gpu] ] stacked.append(default_collate(padded_samples)) elif batch[i].dim() == 1: padded_samples = [ sample.data for sample in batch[i:i + samples_per_gpu] ] stacked.append(default_collate(padded_samples)) else: for i in range(0, len(batch), samples_per_gpu): stacked.append( [sample.data for sample in batch[i:i + samples_per_gpu]]) return DataContainer(stacked, batch[0].stack, batch[0].padding_value) elif isinstance(batch[0], collections.Sequence): transposed = zip(*batch) return [collate(samples, samples_per_gpu) for samples in transposed] elif isinstance(batch[0], collections.Mapping): return { key: collate([d[key] for d in batch], samples_per_gpu) for key in batch[0] } else: return default_collate(batch)
def prepare_train_img(self, idx): img_info = self.img_infos[idx] # load image img = mmcv.imread(osp.join(self.img_prefix, img_info['filename'])) ann = self.get_ann_info(idx) gt_bboxes = ann['bboxes'] gt_labels = ann['labels'] # skip the image if there is no valid gt bbox if len(gt_bboxes) == 0: return None # aug 1: apply extra augmentation if self.extra_aug is not None: img, gt_bboxes, gt_labels = \ self.extra_aug(img, gt_bboxes, gt_labels) # visualize augmented data # import string # import random # import os # temp_data = np.copy(img) # letters = string.ascii_lowercase # name = ''.join(random.choice(letters) for i in range(10)) # dir_path = os.path.dirname(os.getcwd()) + '/MobileDentist' # work_dir = dir_path + '/work_dirs/dental_1009_w_pretrained_wt_fix_w_phonaugment/' # mmcv.imwrite(temp_data, work_dir + '{}.jpg'.format(name)) # aug 2: apply ordinary augmentations: flipping flip = True if np.random.rand() < self.flip_ratio else False # aug 3: apply ordinary augmentations: scaling img_scale = random_scale(self.img_scales, self.multiscale_mode) img, img_shape, pad_shape, scale_factor = \ self.img_transform( img=img, scale=img_scale, flip=flip, pad_val=self.pad_values, keep_ratio=self.resize_keep_ratio ) img = img.copy() gt_bboxes = self.bbox_transform( bboxes=gt_bboxes, img_shape=img_shape, scale_factor=scale_factor, flip=flip ) img_meta = dict( ori_shape=(img_info['height'], img_info['width'], 3), img_shape=img_shape, pad_shape=pad_shape, scale_factor=scale_factor, flip=flip ) data = dict( img=DataContainer(to_tensor(img), stack=True), img_meta=DataContainer(data=img_meta, cpu_only=True), gt_bboxes=DataContainer(data=to_tensor(gt_bboxes)), gt_labels = DataContainer(to_tensor(gt_labels.astype(np.long))) ) return data