MAD/generate_translation.py at main · basiclab/MAD · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
import argparse
import os
from typing import List, Tuple

import torch
from joblib import Parallel, delayed
from loguru import logger
from PIL import Image
from tqdm import tqdm

from config import create_cfg, merge_possible_with_base, show_config
from modeling import build_model
from modeling.translation import TranslationDiffusion


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--config", default=None, type=str)
    parser.add_argument("--save-folder", default="batch_images", type=str)
    parser.add_argument("--source-root", required=True, type=str)
    parser.add_argument("--source-list", required=True, type=str)
    parser.add_argument("--source-label", required=True, type=int)
    parser.add_argument("--target-label", required=True, type=int)
    parser.add_argument("--num-process", default=1, type=int)
    parser.add_argument("--num-of-step", default=700, type=int)
    parser.add_argument("--opts", nargs=argparse.REMAINDER, default=None, type=str)
    return parser.parse_args()


def generate_image(
    cfg,
    save_folder: str,
    source_list: List[Tuple[str, str]],
    source_label: int,
    target_label: int,
    offset: int,
    device: str,
    num_of_step: int,
):
    torch.cuda.set_device(device)
    model = build_model(cfg).to(device)
    model.eval()

    diffuser = TranslationDiffusion(cfg, device)
    os.makedirs(args.save_folder, exist_ok=True)

    count_error = 0
    with tqdm(
        total=len(source_list), position=int(device.split(":")[-1])
    ) as progress_bar:
        for idx, (source_image, source_mask) in enumerate(source_list):
            save_image_name = os.path.join(save_folder, f"pred_{idx + offset}.png")
            if os.path.exists(save_image_name):
                progress_bar.update(1)
                continue
            if source_mask.endswith("jpg"):
                source_mask = source_mask.replace("jpg", "png")
            try:
                transfer_result = diffuser.domain_translation(
                    source_model=model,
                    target_model=model,
                    source_image=source_image,
                    source_class_label=source_label,
                    target_class_label=target_label,
                    parsing_mask=source_mask,
                    start_from_step=num_of_step,
                )
            except Exception as e:
                logger.error(str(e))
                count_error += 1
                continue
            save_image = Image.fromarray(
                (transfer_result[0].permute(1, 2, 0).cpu().numpy() * 255).astype(
                    "uint8"
                )
            )
            save_image.save(save_image_name)

    if count_error != 0:
        print(f"Error in {device}: {count_error}")


if __name__ == "__main__":
    args = parse_args()
    cfg = create_cfg()
    if args.config:
        merge_possible_with_base(cfg, args.config)
    if args.opts:
        cfg.merge_from_list(args.opts)
    show_config(cfg)

    source_list = []
    with open(args.source_list, "r") as f:
        for line in f.readlines():
            source_line = line.strip()
            image_path = os.path.join(args.source_root, source_line)
            mask_path = image_path.replace("images", "parsing")
            source_list.append((image_path, mask_path))

    task_per_process = len(source_list) // args.num_process
    Parallel(n_jobs=args.num_process)(
        delayed(generate_image)(
            cfg,
            args.save_folder,
            source_list=source_list[
                gpu_idx * task_per_process : (
                    ((gpu_idx + 1) * task_per_process)
                    if gpu_idx < args.num_process - 1
                    else len(source_list)
                )
            ],
            source_label=args.source_label,
            target_label=args.target_label,
            offset=gpu_idx * task_per_process,
            device=f"cuda:{gpu_idx}",
            num_of_step=args.num_of_step,
        )
        for gpu_idx in range(args.num_process)
    )