Paper reproduction

[wanglei]

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: [e-mail]

[Arthur Douillard]

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

[wanglei]

This script is very helpful to me. Thank you very much for your help. Wish everything goes well with your work

[Arthur Douillard]

Thanks! You too!