n_blocks = [3, 4, 6, 3] feat_size = 60 mid_stride = False param_fn = 'weights/pspnet101_ADE20K_473_reference.chainer' base_size = 512 crop_size = 473 dataset = TransformDataset(dataset, preprocess) print('dataset:', len(dataset)) chainer.config.train = False model = PSPNet(n_class, n_blocks, feat_size, mid_stride=mid_stride) serializers.load_npz(param_fn, model) if args.gpu >= 0: chainer.cuda.get_device_from_id(args.gpu) model.to_gpu(args.gpu) for i in tqdm(range(args.start_i, args.end_i + 1)): img = dataset[i] out_fn = os.path.join(args.out_dir, os.path.basename(dataset._dataset.img_fns[i])) pred = inference(model, n_class, base_size, crop_size, img, args.scales) assert pred.ndim == 2 if args.model == 'Cityscapes': if args.color_out_dir is not None: color_out = np.zeros((pred.shape[0], pred.shape[1], 3), dtype=np.uint8) label_out = np.zeros_like(pred) for label in cityscapes_labels:
class ModelPSPNet: def __init__(self, model): chainer.config.stride_rate = float(2.0 / 3.0) chainer.config.save_test_image = False chainer.config.train = False self.scales = None self.gpu = 0 self.model = model self.img_fn = 'test.jpg' self.n_class = 19 self.n_blocks = [3, 4, 23, 3] self.feat_size = 90 self.mid_stride = True self.param_fn = 'weights/pspnet101_cityscapes_713_reference.chainer' self.base_size = 2048 self.crop_size = 713 self.labels = cityscapes_label_names self.colors = cityscapes_label_colors if self.model == 'VOC': self.n_class = 21 self.n_blocks = [3, 4, 23, 3] self.feat_size = 60 self.mid_stride = True self.param_fn = 'weights/pspnet101_VOC2012_473_reference.chainer' self.base_size = 512 self.crop_size = 473 self.labels = voc_semantic_segmentation_label_names self.colors = voc_semantic_segmentation_label_colors elif self.model == 'Cityscapes': self.n_class = 19 self.n_blocks = [3, 4, 23, 3] self.feat_size = 90 self.mid_stride = True self.param_fn = 'weights/pspnet101_cityscapes_713_reference.chainer' self.base_size = 2048 self.crop_size = 713 self.labels = cityscapes_label_names self.colors = cityscapes_label_colors elif self.model == 'ADE20K': self.n_class = 150 self.n_blocks = [3, 4, 6, 3] self.feat_size = 60 self.mid_stride = False self.param_fn = 'weights/pspnet101_ADE20K_473_reference.chainer' self.base_size = 512 self.crop_size = 473 self.net = PSPNet(self.n_class, self.n_blocks, self.feat_size, mid_stride=self.mid_stride) serializers.load_npz(self.param_fn, self.net) if self.gpu >= 0: chainer.cuda.get_device_from_id(self.gpu).use() self.net.to_gpu(self.gpu) @staticmethod def load_image(image_data): f = StringIO.BytesIO(image_data) image = Image.open(f) if image.mode not in ('L', 'RGB'): image = image.convert('RGB') image = np.asarray(image, dtype=np.float32) if image.ndim == 2: # reshape (H, W) -> (1, H, W) return image[np.newaxis] else: # transpose (H, W, C) -> (C, H, W) and return image.transpose(2, 0, 1) def do(self, image_data): img = preprocess(self.load_image(image_data=image_data)) # Inference pred = inference(self.net, self.n_class, self.base_size, self.crop_size, img, self.scales) # Save the result image # ax = vis_image(img) # _, legend_handles = vis_label(pred, self.labels, self.colors, alpha=1.0, ax=ax) # ax.legend(handles=legend_handles, bbox_to_anchor=(1.05, 1), loc=2, # borderaxespad=0.) # base = os.path.splitext(os.path.basename(self.img_fn))[0] # plot.savefig('predict_{}.png'.format(base), bbox_inches='tight', dpi=400) # if self.model == 'Cityscapes': # label_out = np.zeros((img.shape[1], img.shape[2], 3), dtype=np.uint8) # for label in cityscapes_labels: # label_out[np.where(pred == label.trainId)] = label.color # # io.imsave( # 'predict_{}_color({}).png'.format("ss", self.scales), label_out) label_out = np.zeros((img.shape[1], img.shape[2], 3), dtype=np.uint8) for label in cityscapes_labels: label_out[np.where(pred == label.trainId)] = label.color pred_data = StringIO.BytesIO() pred_img = Image.fromarray(label_out) pred_img.save(pred_data, format="PNG") pred_data = pred_data.getvalue() print('finish') return pred_data