def main(config, model_path, cuda, crf, camera_id): # Configuration CONFIG = Dict(yaml.load(open(config))) cuda = cuda and torch.cuda.is_available() if cuda: current_device = torch.cuda.current_device() print('Running on', torch.cuda.get_device_name(current_device)) # Label list with open(CONFIG.LABELS) as f: classes = {} for label in f: label = label.rstrip().split('\t') classes[int(label[0])] = label[1].split(',')[0] # Load a model state_dict = torch.load(model_path) # Model model = DeepLabV2_ResNet101_MSC(n_classes=CONFIG.N_CLASSES) model.load_state_dict(state_dict) model.eval() if cuda: model.cuda() image_size = (CONFIG.IMAGE.SIZE.TEST, ) * 2 cap = cv2.VideoCapture(camera_id) cap.set(cv2.CAP_PROP_FOURCC, cv2.VideoWriter_fourcc(*'YUYV')) while True: # Image preprocessing ret, frame = cap.read() image = cv2.resize(frame.astype(float), image_size) raw_image = image.astype(np.uint8) image -= np.array([ float(CONFIG.IMAGE.MEAN.B), float(CONFIG.IMAGE.MEAN.G), float(CONFIG.IMAGE.MEAN.R), ]) image = torch.from_numpy(image.transpose(2, 0, 1)).float().unsqueeze(0) image = image.cuda() if cuda else image # Inference output = model(Variable(image, volatile=True)) output = F.upsample(output, size=image_size, mode='bilinear') output = F.softmax(output, dim=1) output = output.data.cpu().numpy()[0] if crf: output = dense_crf(raw_image, output) labelmap = np.argmax(output.transpose(1, 2, 0), axis=2) labelmap = labelmap.astype(float) / CONFIG.N_CLASSES labelmap = cm.jet_r(labelmap)[..., :-1] * 255.0 cv2.addWeighted(np.uint8(labelmap), 0.5, raw_image, 0.5, 0.0, raw_image) cv2.imshow('DeepLabV2', raw_image) cv2.waitKey(50)
def makePrediction(config_obj, image_path, model, cuda=True, crf=False): torch.set_grad_enabled(False) image_id = extractIdFromPath(image_path) # Image preprocessing image = cv2.imread(image_path, cv2.IMREAD_COLOR).astype(float) image = preProcessImage(image, config_obj, cuda) # Inference output = model(image) output = F.interpolate(output, size=image.shape[2:], mode="bilinear", align_corners=True) output = F.softmax(output, dim=1) output = output.data.cpu().numpy()[0] if crf: output = dense_crf(image_original, output) labelmap = np.argmax(output, axis=0) cocoResFormat = cocostuff.segmentationToCocoResult(labelmap, image_id) return cocoResFormat
def process_image(model, config, cuda, image_size, image_path): image = cv2.imread(image_path, cv2.IMREAD_COLOR).astype(float) image = cv2.resize(image, image_size) image_original = image.astype(np.uint8) image = image[..., ::-1] - np.array( [config.IMAGE.MEAN.R, config.IMAGE.MEAN.G, config.IMAGE.MEAN.B]) image = torch.from_numpy(image.transpose(2, 0, 1)).float().unsqueeze(0) image = image.cuda() if cuda else image # Inference output = model(Variable(image, volatile=True)) output = F.upsample(output, size=image_size, mode="bilinear") output = F.softmax(output, dim=1) output = output[0].cpu().data.numpy() output = dense_crf(image_original, output) # At this point output is label - x - y prob_map = output.transpose(1, 2, 0) label_map = np.argmax(prob_map, axis=2) return prob_map, label_map
def main(config, model_path, cuda, crf): # Configuration CONFIG = Dict(yaml.load(open(config))) cuda = cuda and torch.cuda.is_available() if cuda: current_device = torch.cuda.current_device() print('Running on', torch.cuda.get_device_name(current_device)) image_size = ( CONFIG.IMAGE.SIZE.TEST, CONFIG.IMAGE.SIZE.TEST, ) # Dataset dataset = CocoStuff10k( root=CONFIG.ROOT, split='test', image_size=image_size, scale=False, flip=False, preload=False, ) # DataLoader loader = torch.utils.data.DataLoader( dataset=dataset, batch_size=CONFIG.BATCH_SIZE, num_workers=CONFIG.NUM_WORKERS, shuffle=False, ) state_dict = torch.load(model_path, map_location=lambda storage, loc: storage) # Model model = DeepLabV2_ResNet101_MSC(n_classes=CONFIG.N_CLASSES) model.load_state_dict(state_dict) model = nn.DataParallel(model) model.eval() if cuda: model.cuda() targets, outputs = [], [] for data, target in tqdm( loader, total=len(loader), leave=False, dynamic_ncols=True, ): # Image data = data.cuda() if cuda else data data = Variable(data, volatile=True) # Forward propagation output = model(data) output = F.upsample(output, size=image_size, mode='bilinear') output = F.softmax(output, dim=1) output = output.data.cpu().numpy() # Postprocessing if crf: crf_output = np.zeros(output.shape) images = data.data.cpu().numpy().astype(np.uint8) for i, (image, prob_map) in enumerate(zip(images, output)): image = image.transpose(1, 2, 0) crf_output[i] = dense_crf(image, prob_map) output = crf_output output = np.argmax(output, axis=1) target = target.numpy() for o, t in zip(output, target): outputs.append(o) targets.append(t) score, class_iou = scores(targets, outputs, n_class=CONFIG.N_CLASSES) for k, v in score.items(): print(k, v) score['Class IoU'] = {} for i in range(CONFIG.N_CLASSES): score['Class IoU'][i] = class_iou[i] with open(model_path.replace('.pth', '.json'), 'w') as f: json.dump(score, f, indent=4, sort_keys=True)
def main(config, image_path, cuda, crf): CONFIG = Dict(yaml.load(open(config))) cuda = cuda and torch.cuda.is_available() # Label list with open(CONFIG.LABELS) as f: classes = {} for label in f: label = label.rstrip().split("\t") classes[int(label[0])] = label[1].split(",")[0] # Load a model state_dict = torch.load(CONFIG.PYTORCH_MODEL) # Model model = PSPNet(n_classes=CONFIG.N_CLASSES, n_blocks=CONFIG.N_BLOCKS, pyramids=CONFIG.PYRAMIDS) model.load_state_dict(state_dict) model.eval() if cuda: model.cuda() image_size = (CONFIG.IMAGE.SIZE.TEST, ) * 2 # Image preprocessing image = cv2.imread(image_path, cv2.IMREAD_COLOR).astype(float) image = cv2.resize(image, image_size) image_original = image.astype(np.uint8) image = image[..., ::-1] - np.array( [CONFIG.IMAGE.MEAN.R, CONFIG.IMAGE.MEAN.G, CONFIG.IMAGE.MEAN.B]) image = torch.from_numpy(image.transpose(2, 0, 1)).float().unsqueeze(0) image = image.cuda() if cuda else image # Inference with torch.no_grad(): output = model(Variable(image, volatile=True)) output = F.upsample(output, size=image_size, mode="bilinear") output = F.softmax(output, dim=1) output = output[0].cpu().data.numpy() if crf: output = dense_crf(image_original, output) labelmap = np.argmax(output.transpose(1, 2, 0), axis=2) labels = np.unique(labelmap) rows = np.floor(np.sqrt(len(labels) + 1)) cols = np.ceil((len(labels) + 1) / rows) plt.figure(figsize=(10, 10)) ax = plt.subplot(1, 2, 1) ax.set_title("Input image") ax.imshow(image_original[:, :, ::-1]) ax.set_xticks([]) ax.set_yticks([]) for i, label in enumerate(labels): if classes[label] == 'sky': ax = plt.subplot(1, 2, 2) ax.set_title(classes[label]) print("{0:3d}: {1}".format(label, classes[label])) mask = labelmap != label mask = mask.astype(np.uint8) * 255 mask_temp = np.zeros(mask.shape, mask.dtype) mask = np.expand_dims(mask, axis=2) mask_temp = np.expand_dims(mask_temp, axis=2) mask = np.concatenate((mask, mask_temp, mask_temp), axis=-1) mask = cv2.cvtColor(mask, cv2.COLOR_RGB2GRAY) origin = image_original[:, :, ::-1] right_image = cv2.add(origin, np.zeros(np.shape(origin), dtype=np.uint8), mask=mask) ax.imshow(right_image) cv2.imwrite('demo_sky_masked.jpg', cv2.cvtColor(right_image, cv2.COLOR_BGR2RGB)) # ax.imshow(mask.astype(np.float32), alpha=0.5, cmap="viridis") ax.set_xticks([]) ax.set_yticks([]) plt.tight_layout() plt.show()
def main(config, image_path, cuda, crf): CONFIG = Dict(yaml.load(open(config))) cuda = cuda and torch.cuda.is_available() # Label list with open(CONFIG.LABELS) as f: classes = {} for label in f: label = label.rstrip().split('\t') classes[int(label[0])] = label[1].split(',')[0] # Load a model state_dict = torch.load(CONFIG.PYTORCH_MODEL) # Model model = PSPNet(n_classes=CONFIG.N_CLASSES, n_blocks=CONFIG.N_BLOCKS, pyramids=CONFIG.PYRAMIDS) model.load_state_dict(state_dict) model.eval() if cuda: model.cuda() image_size = (CONFIG.IMAGE.SIZE.TEST, ) * 2 # Image preprocessing image = cv2.imread(image_path, cv2.IMREAD_COLOR).astype(float) image = cv2.resize(image, image_size) image_original = image.astype(np.uint8) image = image[..., ::-1] - np.array([ CONFIG.IMAGE.MEAN.R, CONFIG.IMAGE.MEAN.G, CONFIG.IMAGE.MEAN.B, ]) image = torch.from_numpy(image.transpose(2, 0, 1)).float().unsqueeze(0) image = image.cuda() if cuda else image # Inference output = model(Variable(image, volatile=True)) output = F.upsample(output, size=image_size, mode='bilinear') output = F.softmax(output, dim=1) output = output[0].cpu().data.numpy() if crf: output = dense_crf(image_original, output) labelmap = np.argmax(output.transpose(1, 2, 0), axis=2) labels = np.unique(labelmap) rows = np.floor(np.sqrt(len(labels) + 1)) cols = np.ceil((len(labels) + 1) / rows) plt.figure(figsize=(10, 10)) ax = plt.subplot(rows, cols, 1) ax.set_title('Input image') ax.imshow(image_original[:, :, ::-1]) ax.set_xticks([]) ax.set_yticks([]) for i, label in enumerate(labels): print '{0:3d}: {1}'.format(label, classes[label]) mask = labelmap == label ax = plt.subplot(rows, cols, i + 2) ax.set_title(classes[label]) ax.imshow(image_original[:, :, ::-1]) ax.imshow(mask.astype(np.float32), alpha=0.5, cmap='viridis') ax.set_xticks([]) ax.set_yticks([]) plt.tight_layout() plt.show()
def semseg(rootpath, config, image_path, cuda, crf, out_class_figure): CONFIG = Dict(yaml.load(open(config))) cuda = cuda and torch.cuda.is_available() # Label list with open(os.path.join(rootpath, CONFIG.LABELS)) as f: classes = {} for label in f: label = label.rstrip().split("\t") classes[int(label[0])] = label[1].split(",")[0] # Load a model state_dict = torch.load(os.path.join(rootpath, CONFIG.PYTORCH_MODEL)) # Model model = PSPNet(n_classes=CONFIG.N_CLASSES, n_blocks=CONFIG.N_BLOCKS, pyramids=CONFIG.PYRAMIDS) model.load_state_dict(state_dict) model.eval() if cuda: model.cuda() image_size = (CONFIG.IMAGE.SIZE.TEST, ) * 2 # image_size = (512, 288) # Image preprocessing # MES change for grayscale # for color image input, ignoring alpha channel: use cv2.IMREAD_COLOR # for color image input, keeping alpha channel: use cv2.IMREAD_UNCHANGED # for grayscale image input, use cv2.IMREAD_GRAYSCALE img = cv2.imread(image_path, cv2.IMREAD_COLOR) # img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) image = cv2.resize(img, image_size).astype(float) image_original = image.astype(np.uint8) image = image[..., ::-1] - np.array( [CONFIG.IMAGE.MEAN.R, CONFIG.IMAGE.MEAN.G, CONFIG.IMAGE.MEAN.B]) image = torch.from_numpy(image.transpose(2, 0, 1)).float().unsqueeze(0) image = image.cuda() if cuda else image # Inference # MES change to silence 'volatile' deprecation message with torch.no_grad(): # output = model(Variable(image, volatile=True)) output = model(Variable(image)) # MES change to silence 'upsample' deprecation message # MES change to silence 'align_corners' default message # output = F.upsample(output, size=image_size, mode="bilinear") output = F.interpolate(output, size=image_size, mode="bilinear", align_corners=False) output = F.softmax(output, dim=1) output = output[0].cpu().data.numpy() if crf: output = dense_crf(image_original, output) # # labelmap = np.argmax(output.transpose(1, 2, 0), axis=2) labels = np.unique(labelmap) # rows = np.floor(np.sqrt(len(labels) + 1)) cols = np.ceil((len(labels) + 1) / rows) # plt.figure(figsize=(10, 10)) ax = plt.subplot(rows, cols, 1) ax.set_title("Input image") ax.imshow(image_original[:, :, ::-1]) ax.set_xticks([]) ax.set_yticks([]) for i, label in enumerate(labels): print("{0:3d}: {1}".format(label, classes[label])) mask = labelmap == label ax = plt.subplot(rows, cols, i + 2) ax.set_title(classes[label]) ax.imshow(image_original[:, :, ::-1]) ax.imshow(mask.astype(np.float32), alpha=0.5, cmap="viridis") ax.set_xticks([]) ax.set_yticks([]) # MES change to save the sky class image as a separate image if classes[label] == 'sky': mask_invert = labelmap != label # preserve non-sky pixels masked_image = cv2.bitwise_and(image_original, image_original, mask=mask_invert.astype(np.uint8)) # plt.imsave(out_masked_sky_image, masked_image) plt.tight_layout() # MES changes to save output # plt.show() plt.savefig(out_class_figure) # # return the masked sky image return masked_image
def main(config, model_path, cuda, crf, camera_id): cuda = cuda and torch.cuda.is_available() device = torch.device("cuda" if cuda else "cpu") if cuda: current_device = torch.cuda.current_device() print("Running on", torch.cuda.get_device_name(current_device)) else: print("Running on CPU") # Configuration CONFIG = Dict(yaml.load(open(config))) # Label list with open(CONFIG.LABELS) as f: classes = {} for label in f: label = label.rstrip().split("\t") classes[int(label[0])] = label[1].split(",")[0] torch.set_grad_enabled(False) # Model model = DeepLabV2_ResNet101_MSC(n_classes=CONFIG.N_CLASSES) model.load_state_dict(torch.load(model_path)) model.eval() model.to(device) cap = cv2.VideoCapture(camera_id) cap.set(cv2.CAP_PROP_FOURCC, cv2.VideoWriter_fourcc(*"YUYV")) while True: # Image preprocessing ret, frame = cap.read() h, w, c = frame.shape image = frame.astype(np.float32) image -= np.array([ float(CONFIG.IMAGE.MEAN.B), float(CONFIG.IMAGE.MEAN.G), float(CONFIG.IMAGE.MEAN.R), ]) image = torch.from_numpy(image.transpose(2, 0, 1)).float().unsqueeze(0) image = image.to(device) # Inference output = model(image) output = F.upsample(output, size=(h, w), mode="bilinear", align_corners=False) output = F.softmax(output, dim=1) output = output.data.cpu().numpy()[0] if crf: output = dense_crf(raw_image, output) labelmap = np.argmax(output.transpose(1, 2, 0), axis=2) labelmap = labelmap.astype(float) / CONFIG.N_CLASSES labelmap = cm.jet_r(labelmap)[..., :-1] * 255.0 cv2.addWeighted(np.uint8(labelmap), 0.5, frame, 0.5, 0.0, frame) cv2.imshow("DeepLabV2", frame) cv2.waitKey(10)
def main(config, model_path, cuda, crf): cuda = cuda and torch.cuda.is_available() device = torch.device("cuda" if cuda else "cpu") if cuda: current_device = torch.cuda.current_device() print("Running on", torch.cuda.get_device_name(current_device)) else: print("Running on CPU") # Configuration CONFIG = Dict(yaml.load(open(config))) # Dataset 10k or 164k dataset = get_dataset(CONFIG.DATASET)( root=CONFIG.ROOT, split=CONFIG.SPLIT.VAL, base_size=CONFIG.IMAGE.SIZE.TEST, mean=(CONFIG.IMAGE.MEAN.B, CONFIG.IMAGE.MEAN.G, CONFIG.IMAGE.MEAN.R), warp=CONFIG.WARP_IMAGE, scale=None, flip=False, ) # DataLoader loader = torch.utils.data.DataLoader( dataset=dataset, batch_size=CONFIG.BATCH_SIZE.TEST, num_workers=CONFIG.NUM_WORKERS, shuffle=False, ) torch.set_grad_enabled(False) # Model model = DeepLabV2_ResNet101_MSC(n_classes=CONFIG.N_CLASSES) state_dict = torch.load(model_path, map_location=lambda storage, loc: storage) model.load_state_dict(state_dict) model = nn.DataParallel(model) model.eval() model.to(device) targets, outputs = [], [] for data, target in tqdm(loader, total=len(loader), leave=False, dynamic_ncols=True): # Image data = data.to(device) # Forward propagation output = model(data) output = F.interpolate(output, size=data.shape[2:], mode="bilinear") output = F.softmax(output, dim=1) output = output.data.cpu().numpy() # Postprocessing if crf: crf_output = np.zeros(output.shape) images = data.data.cpu().numpy().astype(np.uint8) for i, (image, prob_map) in enumerate(zip(images, output)): image = image.transpose(1, 2, 0) crf_output[i] = dense_crf(image, prob_map) output = crf_output output = np.argmax(output, axis=1) target = target.numpy() for o, t in zip(output, target): outputs.append(o) targets.append(t) score, class_iou = scores(targets, outputs, n_class=CONFIG.N_CLASSES) for k, v in score.items(): print(k, v) score["Class IoU"] = {} for i in range(CONFIG.N_CLASSES): score["Class IoU"][i] = class_iou[i] with open(model_path.replace(".pth", ".json"), "w") as f: json.dump(score, f, indent=4, sort_keys=True)
def main(config, image_path, model_path, cuda, crf): # Configuration CONFIG = Dict(yaml.load(open(config))) cuda = cuda and torch.cuda.is_available() if cuda: current_device = torch.cuda.current_device() print('Running on', torch.cuda.get_device_name(current_device)) # Label list with open(CONFIG.LABELS) as f: classes = {} for label in f: label = label.rstrip().split('\t') classes[int(label[0])] = label[1].split(',')[0] # Load a model state_dict = torch.load(model_path) # Model model = DeepLabV2_ResNet101_MSC(n_classes=CONFIG.N_CLASSES) model.load_state_dict(state_dict) model.eval() if cuda: model.cuda() image_size = (CONFIG.IMAGE.SIZE.TEST, ) * 2 # Image preprocessing image = cv2.imread(image_path, cv2.IMREAD_COLOR).astype(float) image = cv2.resize(image, image_size) image_original = image.astype(np.uint8) image -= np.array([ float(CONFIG.IMAGE.MEAN.B), float(CONFIG.IMAGE.MEAN.G), float(CONFIG.IMAGE.MEAN.R), ]) image = torch.from_numpy(image.transpose(2, 0, 1)).float().unsqueeze(0) image = image.cuda() if cuda else image # Inference output = model(Variable(image, volatile=True)) output = F.upsample(output, size=image_size, mode='bilinear') output = F.softmax(output, dim=1) output = output.data.cpu().numpy()[0] if crf: output = dense_crf(image_original, output) labelmap = np.argmax(output.transpose(1, 2, 0), axis=2) labels = np.unique(labelmap) # Show results rows = np.floor(np.sqrt(len(labels) + 1)) cols = np.ceil((len(labels) + 1) / rows) plt.figure(figsize=(10, 10)) ax = plt.subplot(rows, cols, 1) ax.set_title('Input image') ax.imshow(image_original[:, :, ::-1]) ax.set_xticks([]) ax.set_yticks([]) for i, label in enumerate(labels): print('{0:3d}: {1}'.format(label, classes[label])) mask = labelmap == label ax = plt.subplot(rows, cols, i + 2) ax.set_title(classes[label]) ax.imshow(image_original[:, :, ::-1]) ax.imshow(mask.astype(np.float32), alpha=0.5, cmap='viridis') ax.set_xticks([]) ax.set_yticks([]) plt.tight_layout() plt.savefig('demo.jpg') plt.show()
def main(config, image_path, model_path, cuda, crf): cuda = cuda and torch.cuda.is_available() device = torch.device("cuda" if cuda else "cpu") if cuda: current_device = torch.cuda.current_device() print("Running on", torch.cuda.get_device_name(current_device)) else: print("Running on CPU") # Configuration CONFIG = Dict(yaml.load(open(config))) # Label list with open(CONFIG.LABELS) as f: classes = {} for label in f: label = label.rstrip().split("\t") classes[int(label[0])] = label[1].split(",")[0] torch.set_grad_enabled(False) # Model model = DeepLabV2_ResNet101_MSC(n_classes=CONFIG.N_CLASSES) state_dict = torch.load(model_path, map_location=lambda storage, loc: storage) model.load_state_dict(state_dict) model.eval() model.to(device) # Image preprocessing image = cv2.imread(image_path, cv2.IMREAD_COLOR).astype(float) scale = CONFIG.IMAGE.SIZE.TEST / max(image.shape[:2]) image = cv2.resize(image, dsize=None, fx=scale, fy=scale) image_original = image.astype(np.uint8) image -= np.array([ float(CONFIG.IMAGE.MEAN.B), float(CONFIG.IMAGE.MEAN.G), float(CONFIG.IMAGE.MEAN.R), ]) image = torch.from_numpy(image.transpose(2, 0, 1)).float().unsqueeze(0) image = image.to(device) # Inference output = model(image) output = F.interpolate(output, size=image.shape[2:], mode="bilinear", align_corners=True) output = F.softmax(output, dim=1) output = output.data.cpu().numpy()[0] if crf: output = dense_crf(image_original, output) labelmap = np.argmax(output, axis=0) labels = np.unique(labelmap) # Show results rows = np.floor(np.sqrt(len(labels) + 1)) cols = np.ceil((len(labels) + 1) / rows) plt.figure(figsize=(10, 10)) ax = plt.subplot(rows, cols, 1) ax.set_title("Input image") ax.imshow(image_original[:, :, ::-1]) ax.set_xticks([]) ax.set_yticks([]) for i, label in enumerate(labels): print("{0:3d}: {1}".format(label, classes[label])) mask = labelmap == label ax = plt.subplot(rows, cols, i + 2) ax.set_title(classes[label]) ax.imshow(image_original[..., ::-1]) ax.imshow(mask.astype(np.float32), alpha=0.5, cmap="viridis") ax.set_xticks([]) ax.set_yticks([]) plt.tight_layout() plt.show()
def main(config, excludeval, embedding, model_path, run, cuda, crf, redo, imagedataset, threshold): pth_extn = '.pth.tar' if osp.isfile(model_path.replace( pth_extn, "_" + run + ".json")) and not threshold and not redo: print("Already Done!") with open(model_path.replace(pth_extn, "_" + run + ".json")) as json_file: data = json.load(json_file) for key, value in data.items(): if not key == "Class IoU": print(key, value) sys.exit() cuda = cuda and torch.cuda.is_available() device = torch.device("cuda" if cuda else "cpu") if cuda: current_device = torch.cuda.current_device() print("Running on", torch.cuda.get_device_name(current_device)) else: print("Running on CPU") # Configuration CONFIG = Dict(yaml.load(open(config))) datadir = os.path.join('data/datasets', imagedataset) print("Split dir: ", datadir) savedir = osp.dirname(model_path) epoch = re.findall("checkpoint_(.*)\." + pth_extn[1:], osp.basename(model_path))[-1] val = None visible_classes = None if run == 'zlss' or run == 'flss': val = np.load(datadir + '/split/test_list.npy') visible_classes = np.load(datadir + '/split/novel_cls.npy') elif run == 'gzlss' or run == 'gflss': val = np.load(datadir + '/split/test_list.npy') if excludeval: vals_cls = np.asarray(np.load(datadir + '/split/seen_cls.npy'), dtype=int) else: vals_cls = np.asarray(np.concatenate([ np.load(datadir + '/split/seen_cls.npy'), np.load(datadir + '/split/val_cls.npy') ]), dtype=int) valu_cls = np.load(datadir + '/split/novel_cls.npy') visible_classes = np.concatenate([vals_cls, valu_cls]) else: print("invalid run ", run) sys.exit() if threshold is not None and run != 'gzlss': print("invalid run for threshold", run) sys.exit() cls_map = np.array([255] * 256) for i, n in enumerate(visible_classes): cls_map[n] = i if threshold is not None: savedir = osp.join(savedir, str(threshold)) if crf is not None: savedir = savedir + '-crf' if run == 'gzlss' or run == 'gflss': novel_cls_map = np.array([255] * 256) for i, n in enumerate(list(valu_cls)): novel_cls_map[cls_map[n]] = i seen_cls_map = np.array([255] * 256) for i, n in enumerate(list(vals_cls)): seen_cls_map[cls_map[n]] = i if threshold is not None: thresholdv = np.asarray(np.zeros((visible_classes.shape[0], 1)), dtype=np.float) thresholdv[np.in1d(visible_classes, vals_cls), 0] = threshold thresholdv = torch.tensor(thresholdv).float().cuda() visible_classesp = np.concatenate([visible_classes, [255]]) all_labels = np.genfromtxt(datadir + '/labels_2.txt', delimiter='\t', usecols=1, dtype='str') print("Visible Classes: ", visible_classes) # Dataset dataset = get_dataset(CONFIG.DATASET)( train=None, test=val, root=CONFIG.ROOT, split=CONFIG.SPLIT.TEST, base_size=CONFIG.IMAGE.SIZE.TEST, mean=(CONFIG.IMAGE.MEAN.B, CONFIG.IMAGE.MEAN.G, CONFIG.IMAGE.MEAN.R), warp=CONFIG.WARP_IMAGE, scale=None, flip=False, ) if embedding == 'word2vec': class_emb = pickle.load( open(datadir + '/word_vectors/word2vec.pkl', "rb")) elif embedding == 'fasttext': class_emb = pickle.load( open(datadir + '/word_vectors/fasttext.pkl', "rb")) elif embedding == 'fastnvec': class_emb = np.concatenate([ pickle.load(open(datadir + '/word_vectors/fasttext.pkl', "rb")), pickle.load(open(datadir + '/word_vectors/word2vec.pkl', "rb")) ], axis=1) else: print("invalid emb ", embedding) sys.exit() class_emb = class_emb[visible_classes] class_emb = F.normalize(torch.tensor(class_emb), p=2, dim=1).cuda() print("Embedding dim: ", class_emb.shape[1]) print("# Visible Classes: ", class_emb.shape[0]) # DataLoader loader = torch.utils.data.DataLoader( dataset=dataset, batch_size=CONFIG.BATCH_SIZE.TEST, num_workers=CONFIG.NUM_WORKERS, shuffle=False, ) torch.set_grad_enabled(False) # Model model = DeepLabV2_ResNet101_MSC(class_emb.shape[1], class_emb) sdir = osp.join(savedir, model_path.replace(pth_extn, ""), str(epoch), run) state_dict = torch.load(model_path, map_location=lambda storage, loc: storage) model = nn.DataParallel(model) model.load_state_dict(state_dict['state_dict']) model.eval() model.to(device) imgfeat = [] targets, outputs = [], [] for data, target, img_id in tqdm(loader, total=len(loader), leave=False, dynamic_ncols=True): # Image data = data.to(device) # Forward propagation output = model(data) output = F.interpolate(output, size=data.shape[2:], mode="bilinear", align_corners=False) output = F.softmax(output, dim=1) if threshold is not None: output = output - thresholdv.view(1, -1, 1, 1) target = cls_map[target.numpy()] # Postprocessing if crf: output = output.data.cpu().numpy() crf_output = np.zeros(output.shape) images = data.data.cpu().numpy().astype(np.uint8) for i, (image, prob_map) in enumerate(zip(images, output)): image = image.transpose(1, 2, 0) crf_output[i] = dense_crf(image, prob_map) output = crf_output output = np.argmax(output, axis=1) else: output = torch.argmax(output, dim=1).cpu().numpy() for o, t in zip(output, target): outputs.append(o) targets.append(t) if run == 'gzlss' or run == 'gflss': score, class_iou = scores_gzsl(targets, outputs, n_class=len(visible_classes), seen_cls=cls_map[vals_cls], unseen_cls=cls_map[valu_cls]) else: score, class_iou = scores(targets, outputs, n_class=len(visible_classes)) for k, v in score.items(): print(k, v) score["Class IoU"] = {} for i in range(len(visible_classes)): score["Class IoU"][all_labels[visible_classes[i]]] = class_iou[i] if threshold is not None: with open( model_path.replace(pth_extn, "_" + run + '_T' + str(threshold) + ".json"), "w") as f: json.dump(score, f, indent=4, sort_keys=True) else: with open(model_path.replace(pth_extn, "_" + run + ".json"), "w") as f: json.dump(score, f, indent=4, sort_keys=True) print(score["Class IoU"])
def main(config, model_path, cuda): # Configuration with open(config) as f: CONFIG = yaml.load(f) cuda = cuda and torch.cuda.is_available() image_size = (CONFIG['IMAGE']['SIZE']['TEST'], CONFIG['IMAGE']['SIZE']['TEST']) n_classes = CONFIG['N_CLASSES'] # Dataset dataset = get_dataset(CONFIG['DATASET'])(root=CONFIG['ROOT'], split='test', image_size=image_size, scale=False, flip=False, preload=False) # DataLoader loader = torch.utils.data.DataLoader(dataset=dataset, batch_size=CONFIG['BATCH_SIZE'], num_workers=CONFIG['NUM_WORKERS'], shuffle=False) state_dict = torch.load(model_path, map_location=lambda storage, loc: storage) # Model model = DeepLabV2_ResNet101_MSC(n_classes=n_classes) model.load_state_dict(state_dict) model.eval() if cuda: model.cuda() targets, outputs = [], [] for data, target in tqdm(loader, total=len(loader), leave=False, dynamic_ncols=True): # Image data = data.cuda() if cuda else data data = Variable(data, volatile=True) # Forward propagation output = model(data) output = F.upsample(output, size=image_size, mode='bilinear') output = F.softmax(output, dim=1) output = output.data.cpu().numpy() crf_output = np.zeros(output.shape) images = data.data.cpu().numpy().astype(np.uint8) for i, (image, prob_map) in enumerate(zip(images, output)): image = image.transpose(1, 2, 0) crf_output[i] = dense_crf(image, prob_map) output = crf_output output = np.argmax(output, axis=1) target = target.numpy() for o, t in zip(output, target): outputs.append(o) targets.append(t) score, class_iou = scores(targets, outputs, n_class=n_classes) for k, v in score.items(): print k, v score['Class IoU'] = {} for i in range(n_classes): score['Class IoU'][i] = class_iou[i] with open('results.json', 'w') as f: json.dump(score, f)
def dense_crf_wrapper(args): return dense_crf(args[0], args[1])
def main(dataset, image_path, model_path, cuda, crf): CONFIG = { 'voc12': { 'path_pytorch_model': 'data/models/deeplab_resnet101/voc12/deeplabv2_resnet101_VOC2012.pth', 'label_list': 'data/datasets/voc12/labels.txt', 'n_classes': 21, 'n_blocks': [3, 4, 23, 3], 'pyramids': [6, 3, 2, 1], 'image': { 'size': { 'train': 473, 'test': 473, }, 'mean': { 'R': 122.675, 'G': 116.669, 'B': 104.008, } }, }, 'cocostuff': { 'path_pytorch_model': 'data/models/deeplab_resnet101/cocostuff_rgb/deeplabv2_resnet101_VOC2012.pth', 'label_list': 'data/datasets/voc12/labels.txt', 'n_classes': 183, 'n_blocks': [3, 4, 23, 3], 'pyramids': [6, 3, 2, 1], 'image': { 'size': { 'train': 321, 'test': 513, }, 'mean': { 'R': 122.675, 'G': 116.669, 'B': 104.008, } }, }, }.get(dataset) cuda = cuda and torch.cuda.is_available() # Label list with open(CONFIG['label_list']) as f: classes = {} for label in f: label = label.rstrip().split('\t') classes[int(label[0])] = label[1].split(',')[0] # Load a model if model_path is None: state_dict = torch.load(CONFIG['path_pytorch_model']) else: state_dict = torch.load(model_path) # Model model = DeepLabV2_ResNet101_MSC(n_classes=CONFIG['n_classes']) model.load_state_dict(state_dict) model.eval() if cuda: model.cuda() image_size = (CONFIG['image']['size']['test'], CONFIG['image']['size']['test']) # Image preprocessing image = cv2.imread(image_path, cv2.IMREAD_COLOR).astype(float) image = cv2.resize(image, image_size) image_original = image.astype(np.uint8) image -= np.array([ float(CONFIG['image']['mean']['B']), float(CONFIG['image']['mean']['G']), float(CONFIG['image']['mean']['R']) ]) image = torch.from_numpy(image.transpose(2, 0, 1)).float().unsqueeze(0) image = image.cuda() if cuda else image # Inference output = model(Variable(image, volatile=True)) output = F.upsample(output, size=image_size, mode='bilinear') output = F.softmax(output, dim=1) output = output.data.cpu().numpy()[0] if crf: output = dense_crf(image_original, output) labelmap = np.argmax(output.transpose(1, 2, 0), axis=2) labels = np.unique(labelmap) # Show results rows = np.floor(np.sqrt(len(labels) + 1)) cols = np.ceil((len(labels) + 1) / rows) plt.figure(figsize=(10, 10)) ax = plt.subplot(rows, cols, 1) ax.set_title('Input image') ax.imshow(image_original[:, :, ::-1]) ax.set_xticks([]) ax.set_yticks([]) for i, label in enumerate(labels): print '{0:3d}: {1}'.format(label, classes[label]) mask = labelmap == label ax = plt.subplot(rows, cols, i + 2) ax.set_title(classes[label]) ax.imshow(np.dstack((mask, ) * 3) * image_original[:, :, ::-1]) ax.set_xticks([]) ax.set_yticks([]) plt.show()