arthurdouillard
commented
3 years ago I used to have a very dirty script, that I haven't cleaned yet:
import argparse
import os
import json
import tasks
from skimage import color
from argparser import get_argparser, modify_command_options
from segmentation_module import make_model
from tasks import get_per_task_classes
from dataset import transform
import torch
from PIL import Image
import numpy as np
import glob
import os
from torchvision.transforms import functional as Fv
import matplotlib.pyplot as plt
from torch import distributed
from apex.parallel import DistributedDataParallel
def color_map_viz():
labels = ['background', 'aeroplane', 'bicycle', 'bird', 'boat', 'bottle', 'bus', 'car', 'cat', 'chair', 'cow', 'diningtable', 'dog', 'horse', 'motorbike', 'person', 'pottedplant', 'sheep', 'sofa', 'train', 'tvmonitor', 'void']
nclasses = 21
row_size = 40
col_size = 20
cmap = color_map()
array = np.empty((row_size*(nclasses+1), col_size, cmap.shape[1]), dtype=cmap.dtype)
for i in range(nclasses):
array[i*row_size:i*row_size+row_size, :] = cmap[i]
array[nclasses*row_size:nclasses*row_size+row_size, :] = cmap[-1]
plt.imshow(array)
plt.yticks([row_size*i+row_size/2 for i in range(nclasses+1)], labels)
plt.xticks([])
if __name__ == "__main__":
parser = get_argparser()
folder = "/local/douillard/continual_segmentation/visualization/"
opts = parser.parse_args()
opts = modify_command_options(opts)
distributed.init_process_group(backend='nccl', init_method='env://')
device_id, device = opts.local_rank, torch.device(opts.local_rank)
rank, world_size = distributed.get_rank(), distributed.get_world_size()
torch.cuda.set_device(device_id)
classes = get_per_task_classes(opts.dataset, opts.task, opts.step[0])
model = make_model(opts, classes=classes)
model = DistributedDataParallel(model.cuda())
checkpoint = torch.load(opts.ckpt, map_location="cpu")
model.load_state_dict(checkpoint["model_state"], strict=True)
old_model.module.in_eval = True
model.eval()
subset = "val"
step = opts.step
task = opts.task
print(f'task {task} on step {step}')
base_path = f"/home/douillard/continual_segmentation/data/voc/15-5s-ov/{subset}-*.npy"
paths = sorted(glob.glob(base_path))
print(paths)
indexes = set()
for step_index in range(int(step[0]) + 1):
indexes |= set(np.load(paths[step_index]))
indexes = sorted(list(indexes))
if subset == "train":
split_f = f"/local/douillard/pascal_voc_2012/list/{subset}_aug.txt"
else:
split_f = f"/local/douillard/pascal_voc_2012/list/{subset}.txt"
with open(os.path.join(split_f), "r") as f:
file_names = [x[:-1].split(' ') for x in f.readlines()]
voc_root = "/local/douillard/pascal_voc_2012/"
targets_paths = [
(
os.path.join(voc_root, "VOCdevkit/VOC2012", x[0][1:]),
os.path.join(voc_root, x[1][1:])
)
for x in file_names
]
targets_paths = np.array(targets_paths)
targets_paths = targets_paths[indexes]
print(len(targets_paths))
os.makedirs("images_viz", exist_ok=True)
class myResize:
def __call__(self, x):
return Fv.resize(x, opts.crop_size, Image.BILINEAR)
viz_transform = transform.Compose(
[
myResize(),
transform.CenterCrop(size=opts.crop_size),
]
)
val_transform = transform.Compose(
[
transform.Resize(size=opts.crop_size),
transform.CenterCrop(size=opts.crop_size),
transform.ToTensor(),
transform.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
]
)
def color_map(N=256, normalized=False):
def bitget(byteval, idx):
return ((byteval & (1 << idx)) != 0)
dtype = 'float32' if normalized else 'uint8'
cmap = np.zeros((N, 3), dtype=dtype)
for i in range(N):
r = g = b = 0
c = i
for j in range(8):
r = r | (bitget(c, 0) << 7-j)
g = g | (bitget(c, 1) << 7-j)
b = b | (bitget(c, 2) << 7-j)
c = c >> 3
cmap[i] = np.array([r, g, b])
cmap = cmap/255 if normalized else cmap
return cmap
cmap = color_map()[:, np.newaxis, :]
for i, (image_path, target_path) in enumerate(targets_paths):
transform.Resize(size=opts.crop_size),
transform.CenterCrop(size=opts.crop_size)
image_name = image_path.split("/")[-1].split(".")[0]
# Image
image, target = val_transform(Image.open(image_path).convert('RGB'), Image.open(target_path))
raw_image = viz_transform(Image.open(image_path).convert('RGB'))
raw_image.save(os.path.join(folder, f"{image_name}_pixels.png"))
# Preds
outputs, _ = model(image[None].cuda())
preds = outputs.argmax(dim=1)[0].cpu().numpy().astype("uint8")[:, :, np.newaxis]
np.save("preds.npy", preds)
raw_image.save("image.png")
np.save("targets.npy", target.numpy())
new_im = np.dot(preds == 0, cmap[0])
for i in range(1, cmap.shape[0]):
new_im += np.dot(preds == i, cmap[i])
new_im = Image.fromarray(new_im.astype(np.uint8))
blend_preds = Image.blend(raw_image, new_im, alpha=0.8)
blend_preds.save(os.path.join(folder, f"{image_name}_{step[0]}_{opts.name}.png"))
# Target
shp = target.shape
target = target.reshape(-1)
min_idx = {0: -1, 1: 16, 2: 17, 3: 18, 4: 19, 5: 20}[step[0]]
if step[0] == 0:
mask = target < 16
else:
mask = np.logical_and(target >= min_idx, target < min_idx + 1)
target[~mask] = 0
target = target.reshape(*shp)
target = target[:, :, np.newaxis]
new_im = np.dot(target == 0, cmap[0])
for i in range(1, cmap.shape[0]):
new_im += np.dot(target == i, cmap[i])
new_im = Image.fromarray(new_im.astype(np.uint8))
blend_image = Image.blend(raw_image, new_im, alpha=0.8)
blend_image.save(os.path.join(folder, f"{image_name}_{step[0]}_gt.png"))
if i % 100 == 0:
print(f"{step}: {i}/{len(targets_paths)}")It should work I guess. You can use a command like that (replace the value in {...}): CUDA_VISIBLE_DEVICES=${GPU} python3 -m torch.distributed.launch --master_port ${PORT} --nproc_per_node=${NB_GPU} visualize.py --dataset ${DATASET} --name ${NAME} --task ${TASK} --method ${METHOD} --overlap --step 1 --ckpt ${SECONDMODEL} --opt_level O1
AI-student-wl
Hello, I have run your code model on my server. I know your code has completed the training and prediction, but I can't find the prediction picture in the file. Only the predicted score, my deep learning ability is limited, so I ask you for help on how to get the segmented image in the paper. Thank you. My email: wlmail2021@163.com