Mask weights are added!

guided_diffusion/dist_util.py

@@ -23,9 +23,10 @@ def setup_dist():
     Setup a distributed process group.
     """
     if dist.is_initialized():
+        print("Distributed training already initialized")
         return
     # os.environ["CUDA_VISIBLE_DEVICES"] = f"{MPI.COMM_WORLD.Get_rank() % GPUS_PER_NODE}"
-    os.environ["CUDA_VISIBLE_DEVICES"] = "4"
+    os.environ["CUDA_VISIBLE_DEVICES"] = "0"
 
     comm = MPI.COMM_WORLD
     backend = "gloo" if not th.cuda.is_available() else "nccl"

guided_diffusion/gaussian_diffusion.py

@@ -818,7 +818,8 @@ class GaussianDiffusion:
                 ModelMeanType.EPSILON: noise,
             }[self.model_mean_type]
             assert model_output.shape == target.shape == x_start.shape
-            terms["mse"] = mean_flat((target - model_output) ** 2)
+            semantic_weights = model_kwargs['semantic_weights']
+            terms["mse"] = mean_flat(((target - model_output) ** 2) * semantic_weights)
             if "vb" in terms:
                 terms["loss"] = terms["mse"] + terms["vb"]
             else:

guided_diffusion/image_datasets.py

@@ -58,6 +58,10 @@ def load_data(
         all_files = _list_image_files_recursively(os.path.join(data_dir, 'train' if is_train else 'test', 'images'))
         classes = _list_image_files_recursively(os.path.join(data_dir, 'train' if is_train else 'test', 'labels'))
         instances = _list_image_files_recursively(os.path.join(data_dir, 'train' if is_train else 'test', 'labels'))
+    elif dataset_mode == 'miniFrance':
+        all_files = _list_image_files_recursively(os.path.join(data_dir, 'train/images' if is_train else 'val/images'))
+        classes = _list_image_files_recursively(os.path.join(data_dir, 'train/label_indices' if is_train else 'val/label_indices'))
+        instances = None
     else:
         raise NotImplementedError('{} not implemented'.format(dataset_mode))
 
@@ -150,7 +154,7 @@ class ImageDataset(Dataset):
         else:
             pil_instance = None
 
-        if self.dataset_mode == 'cityscapes':
+        if self.dataset_mode == 'cityscapes' or self.dataset_mode == 'miniFrance':
             arr_image, arr_class, arr_instance = resize_arr([pil_image, pil_class, pil_instance], self.resolution)
         else:
             if self.is_train:

guided_diffusion/mask.py

@@ -0,0 +1,338 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+import math
+
+import cv2
+import numpy as np
+from PIL import Image, ImageDraw
+
+
+def random_cropping_bbox(img_shape=(256,256), mask_mode='onedirection'):
+    h, w = img_shape
+    if mask_mode == 'onedirection':
+        _type = np.random.randint(0, 4)
+        if _type == 0:
+            top, left, height, width = 0, 0, h, w//2
+        elif _type == 1:
+            top, left, height, width = 0, 0, h//2, w
+        elif _type == 2:
+            top, left, height, width = h//2, 0, h//2, w
+        elif _type == 3:
+            top, left, height, width = 0, w//2, h, w//2
+    else:
+        target_area = (h*w)//2
+        width = np.random.randint(target_area//h, w)
+        height = target_area//width
+        if h==height:
+            top = 0
+        else:
+            top = np.random.randint(0, h-height)
+        if w==width:
+            left = 0
+        else:
+            left = np.random.randint(0, w-width)
+    return (top, left, height, width)
+
+def random_bbox(img_shape=(256,256), max_bbox_shape=(128, 128), max_bbox_delta=40, min_margin=20):
+    """Generate a random bbox for the mask on a given image.
+
+    In our implementation, the max value cannot be obtained since we use
+    `np.random.randint`. And this may be different with other standard scripts
+    in the community.
+
+    Args:
+        img_shape (tuple[int]): The size of a image, in the form of (h, w).
+        max_bbox_shape (int | tuple[int]): Maximum shape of the mask box,
+            in the form of (h, w). If it is an integer, the mask box will be
+            square.
+        max_bbox_delta (int | tuple[int]): Maximum delta of the mask box,
+            in the form of (delta_h, delta_w). If it is an integer, delta_h
+            and delta_w will be the same. Mask shape will be randomly sampled
+            from the range of `max_bbox_shape - max_bbox_delta` and
+            `max_bbox_shape`. Default: (40, 40).
+        min_margin (int | tuple[int]): The minimum margin size from the
+            edges of mask box to the image boarder, in the form of
+            (margin_h, margin_w). If it is an integer, margin_h and margin_w
+            will be the same. Default: (20, 20).
+
+    Returns:
+        tuple[int]: The generated box, (top, left, h, w).
+    """
+    if not isinstance(max_bbox_shape, tuple):
+        max_bbox_shape = (max_bbox_shape, max_bbox_shape)
+    if not isinstance(max_bbox_delta, tuple):
+        max_bbox_delta = (max_bbox_delta, max_bbox_delta)
+    if not isinstance(min_margin, tuple):
+        min_margin = (min_margin, min_margin)
+        
+    img_h, img_w = img_shape[:2]
+    max_mask_h, max_mask_w = max_bbox_shape
+    max_delta_h, max_delta_w = max_bbox_delta
+    margin_h, margin_w = min_margin
+
+    if max_mask_h > img_h or max_mask_w > img_w:
+        raise ValueError(f'mask shape {max_bbox_shape} should be smaller than '
+                         f'image shape {img_shape}')
+    if (max_delta_h // 2 * 2 >= max_mask_h
+            or max_delta_w // 2 * 2 >= max_mask_w):
+        raise ValueError(f'mask delta {max_bbox_delta} should be smaller than'
+                         f'mask shape {max_bbox_shape}')
+    if img_h - max_mask_h < 2 * margin_h or img_w - max_mask_w < 2 * margin_w:
+        raise ValueError(f'Margin {min_margin} cannot be satisfied for img'
+                         f'shape {img_shape} and mask shape {max_bbox_shape}')
+
+    # get the max value of (top, left)
+    max_top = img_h - margin_h - max_mask_h
+    max_left = img_w - margin_w - max_mask_w
+    # randomly select a (top, left)
+    top = np.random.randint(margin_h, max_top)
+    left = np.random.randint(margin_w, max_left)
+    # randomly shrink the shape of mask box according to `max_bbox_delta`
+    # the center of box is fixed
+    delta_top = np.random.randint(0, max_delta_h // 2 + 1)
+    delta_left = np.random.randint(0, max_delta_w // 2 + 1)
+    top = top + delta_top
+    left = left + delta_left
+    h = max_mask_h - delta_top
+    w = max_mask_w - delta_left
+    return (top, left, h, w)
+
+
+def bbox2mask(img_shape, bbox, dtype='uint8'):
+    """Generate mask in ndarray from bbox.
+
+    The returned mask has the shape of (h, w, 1). '1' indicates the
+    hole and '0' indicates the valid regions.
+
+    We prefer to use `uint8` as the data type of masks, which may be different
+    from other codes in the community.
+
+    Args:
+        img_shape (tuple[int]): The size of the image.
+        bbox (tuple[int]): Configuration tuple, (top, left, height, width)
+        dtype (str): Indicate the data type of returned masks. Default: 'uint8'
+
+    Return:
+        numpy.ndarray: Mask in the shape of (h, w, 1).
+    """
+
+    height, width = img_shape[:2]
+
+    mask = np.zeros((height, width, 1), dtype=dtype)
+    mask[bbox[0]:bbox[0] + bbox[2], bbox[1]:bbox[1] + bbox[3], :] = 1
+
+    return mask
+
+
+def brush_stroke_mask(img_shape,
+                      num_vertices=(4, 12),
+                      mean_angle=2 * math.pi / 5,
+                      angle_range=2 * math.pi / 15,
+                      brush_width=(12, 40),
+                      max_loops=4,
+                      dtype='uint8'):
+    """Generate free-form mask.
+
+    The method of generating free-form mask is in the following paper:
+    Free-Form Image Inpainting with Gated Convolution.
+
+    When you set the config of this type of mask. You may note the usage of
+    `np.random.randint` and the range of `np.random.randint` is [left, right).
+
+    We prefer to use `uint8` as the data type of masks, which may be different
+    from other codes in the community.
+
+    TODO: Rewrite the implementation of this function.
+
+    Args:
+        img_shape (tuple[int]): Size of the image.
+        num_vertices (int | tuple[int]): Min and max number of vertices. If
+            only give an integer, we will fix the number of vertices.
+            Default: (4, 12).
+        mean_angle (float): Mean value of the angle in each vertex. The angle
+            is measured in radians. Default: 2 * math.pi / 5.
+        angle_range (float): Range of the random angle.
+            Default: 2 * math.pi / 15.
+        brush_width (int | tuple[int]): (min_width, max_width). If only give
+            an integer, we will fix the width of brush. Default: (12, 40).
+        max_loops (int): The max number of for loops of drawing strokes.
+        dtype (str): Indicate the data type of returned masks.
+            Default: 'uint8'.
+
+    Returns:
+        numpy.ndarray: Mask in the shape of (h, w, 1).
+    """
+
+    img_h, img_w = img_shape[:2]
+    if isinstance(num_vertices, int):
+        min_num_vertices, max_num_vertices = num_vertices, num_vertices + 1
+    elif isinstance(num_vertices, tuple):
+        min_num_vertices, max_num_vertices = num_vertices
+    else:
+        raise TypeError('The type of num_vertices should be int'
+                        f'or tuple[int], but got type: {num_vertices}')
+
+    if isinstance(brush_width, tuple):
+        min_width, max_width = brush_width
+    elif isinstance(brush_width, int):
+        min_width, max_width = brush_width, brush_width + 1
+    else:
+        raise TypeError('The type of brush_width should be int'
+                        f'or tuple[int], but got type: {brush_width}')
+
+    average_radius = math.sqrt(img_h * img_h + img_w * img_w) / 8
+    mask = Image.new('L', (img_w, img_h), 0)
+
+    loop_num = np.random.randint(1, max_loops)
+    num_vertex_list = np.random.randint(
+        min_num_vertices, max_num_vertices, size=loop_num)
+    angle_min_list = np.random.uniform(0, angle_range, size=loop_num)
+    angle_max_list = np.random.uniform(0, angle_range, size=loop_num)
+
+    for loop_n in range(loop_num):
+        num_vertex = num_vertex_list[loop_n]
+        angle_min = mean_angle - angle_min_list[loop_n]
+        angle_max = mean_angle + angle_max_list[loop_n]
+        angles = []
+        vertex = []
+
+        # set random angle on each vertex
+        angles = np.random.uniform(angle_min, angle_max, size=num_vertex)
+        reverse_mask = (np.arange(num_vertex, dtype=np.float32) % 2) == 0
+        angles[reverse_mask] = 2 * math.pi - angles[reverse_mask]
+
+        h, w = mask.size
+
+        # set random vertices
+        vertex.append((np.random.randint(0, w), np.random.randint(0, h)))
+        r_list = np.random.normal(
+            loc=average_radius, scale=average_radius // 2, size=num_vertex)
+        for i in range(num_vertex):
+            r = np.clip(r_list[i], 0, 2 * average_radius)
+            new_x = np.clip(vertex[-1][0] + r * math.cos(angles[i]), 0, w)
+            new_y = np.clip(vertex[-1][1] + r * math.sin(angles[i]), 0, h)
+            vertex.append((int(new_x), int(new_y)))
+        # draw brush strokes according to the vertex and angle list
+        draw = ImageDraw.Draw(mask)
+        width = np.random.randint(min_width, max_width)
+        draw.line(vertex, fill=1, width=width)
+        for v in vertex:
+            draw.ellipse((v[0] - width // 2, v[1] - width // 2,
+                          v[0] + width // 2, v[1] + width // 2),
+                         fill=1)
+    # randomly flip the mask
+    if np.random.normal() > 0:
+        mask.transpose(Image.FLIP_LEFT_RIGHT)
+    if np.random.normal() > 0:
+        mask.transpose(Image.FLIP_TOP_BOTTOM)
+    mask = np.array(mask).astype(dtype=getattr(np, dtype))
+    mask = mask[:, :, None]
+    return mask
+
+
+def random_irregular_mask(img_shape,
+                          num_vertices=(4, 8),
+                          max_angle=4,
+                          length_range=(10, 100),
+                          brush_width=(10, 40),
+                          dtype='uint8'):
+    """Generate random irregular masks.
+
+    This is a modified version of free-form mask implemented in
+    'brush_stroke_mask'.
+
+    We prefer to use `uint8` as the data type of masks, which may be different
+    from other codes in the community.
+
+    TODO: Rewrite the implementation of this function.
+
+    Args:
+        img_shape (tuple[int]): Size of the image.
+        num_vertices (int | tuple[int]): Min and max number of vertices. If
+            only give an integer, we will fix the number of vertices.
+            Default: (4, 8).
+        max_angle (float): Max value of angle at each vertex. Default 4.0.
+        length_range (int | tuple[int]): (min_length, max_length). If only give
+            an integer, we will fix the length of brush. Default: (10, 100).
+        brush_width (int | tuple[int]): (min_width, max_width). If only give
+            an integer, we will fix the width of brush. Default: (10, 40).
+        dtype (str): Indicate the data type of returned masks. Default: 'uint8'
+
+    Returns:
+        numpy.ndarray: Mask in the shape of (h, w, 1).
+    """
+
+    h, w = img_shape[:2]
+
+    mask = np.zeros((h, w), dtype=dtype)
+    if isinstance(length_range, int):
+        min_length, max_length = length_range, length_range + 1
+    elif isinstance(length_range, tuple):
+        min_length, max_length = length_range
+    else:
+        raise TypeError('The type of length_range should be int'
+                        f'or tuple[int], but got type: {length_range}')
+    if isinstance(num_vertices, int):
+        min_num_vertices, max_num_vertices = num_vertices, num_vertices + 1
+    elif isinstance(num_vertices, tuple):
+        min_num_vertices, max_num_vertices = num_vertices
+    else:
+        raise TypeError('The type of num_vertices should be int'
+                        f'or tuple[int], but got type: {num_vertices}')
+
+    if isinstance(brush_width, int):
+        min_brush_width, max_brush_width = brush_width, brush_width + 1
+    elif isinstance(brush_width, tuple):
+        min_brush_width, max_brush_width = brush_width
+    else:
+        raise TypeError('The type of brush_width should be int'
+                        f'or tuple[int], but got type: {brush_width}')
+
+    num_v = np.random.randint(min_num_vertices, max_num_vertices)
+
+    for i in range(num_v):
+        start_x = np.random.randint(w)
+        start_y = np.random.randint(h)
+        # from the start point, randomly setlect n \in [1, 6] directions.
+        direction_num = np.random.randint(1, 6)
+        angle_list = np.random.randint(0, max_angle, size=direction_num)
+        length_list = np.random.randint(
+            min_length, max_length, size=direction_num)
+        brush_width_list = np.random.randint(
+            min_brush_width, max_brush_width, size=direction_num)
+        for direct_n in range(direction_num):
+            angle = 0.01 + angle_list[direct_n]
+            if i % 2 == 0:
+                angle = 2 * math.pi - angle
+            length = length_list[direct_n]
+            brush_w = brush_width_list[direct_n]
+            # compute end point according to the random angle
+            end_x = (start_x + length * np.sin(angle)).astype(np.int32)
+            end_y = (start_y + length * np.cos(angle)).astype(np.int32)
+
+            cv2.line(mask, (start_y, start_x), (end_y, end_x), 1, brush_w)
+            start_x, start_y = end_x, end_y
+    mask = np.expand_dims(mask, axis=2)
+
+    return mask
+
+
+def get_irregular_mask(img_shape, area_ratio_range=(0.15, 0.5), **kwargs):
+    """Get irregular mask with the constraints in mask ratio
+
+    Args:
+        img_shape (tuple[int]): Size of the image.
+        area_ratio_range (tuple(float)): Contain the minimum and maximum area
+        ratio. Default: (0.15, 0.5).
+
+    Returns:
+        numpy.ndarray: Mask in the shape of (h, w, 1).
+    """
+
+    mask = random_irregular_mask(img_shape, **kwargs)
+    min_ratio, max_ratio = area_ratio_range
+
+    while not min_ratio < (np.sum(mask) /
+                           (img_shape[0] * img_shape[1])) < max_ratio:
+        mask = random_irregular_mask(img_shape, **kwargs)
+
+    return mask

guided_diffusion/train_util.py

@@ -1,7 +1,8 @@
 import copy
 import functools
 import os
-
+import random
+from .mask import (bbox2mask, brush_stroke_mask, get_irregular_mask, random_bbox, random_cropping_bbox)
 import blobfile as bf
 import torch as th
 import torch.distributed as dist
@@ -40,6 +41,7 @@ class TrainLoop:
         schedule_sampler=None,
         weight_decay=0.0,
         lr_anneal_steps=0,
+        inpainting=False,
     ):
         self.model = model
         self.diffusion = diffusion
@@ -62,6 +64,7 @@ class TrainLoop:
         self.schedule_sampler = schedule_sampler or UniformSampler(diffusion)
         self.weight_decay = weight_decay
         self.lr_anneal_steps = lr_anneal_steps
+        self.inpainting = inpainting
 
         self.step = 0
         self.resume_step = 0
@@ -118,10 +121,15 @@ class TrainLoop:
             self.resume_step = parse_resume_step_from_filename(resume_checkpoint)
             if dist.get_rank() == 0:
                 logger.log(f"loading model from checkpoint: {resume_checkpoint}...")
+                checkpoint = th.load(resume_checkpoint, map_location=dist_util.dev())
+                model_dict = self.model.state_dict()
+
+                # Filter out mismatched params
+                filtered_checkpoint = {k: v for k, v in checkpoint.items() if k in model_dict and v.size() == model_dict[k].size()}
+
                 self.model.load_state_dict(
-                    th.load(
-                        resume_checkpoint, map_location=dist_util.dev()
-                    )
+                    filtered_checkpoint,
+                    strict=False
                 )
 
         dist_util.sync_params(self.model.parameters())
@@ -163,7 +171,7 @@ class TrainLoop:
             or self.step + self.resume_step < self.lr_anneal_steps
         ):
             batch, cond = next(self.data)
-            cond = self.preprocess_input(cond)
+            cond = self.preprocess_input(batch, cond)
             self.run_step(batch, cond)
             if self.step % self.log_interval == 0:
                 logger.dumpkvs()
@@ -262,17 +270,33 @@ class TrainLoop:
 
         dist.barrier()
 
-    def preprocess_input(self, data):
+    def preprocess_input(self, batch, data):
         # move to GPU and change data types
         data['label'] = data['label'].long()
 
         # create one-hot label map
         label_map = data['label']
-        bs, _, h, w = label_map.size()
+        bs, c, h, w = label_map.size()
+        # print(label_map.shape, batch.shape)
         nc = self.num_classes
         input_label = th.FloatTensor(bs, nc, h, w).zero_()
+        if self.inpainting:
+            unique_per_batch = [th.unique(x) for x in label_map]
+            bypass_label_indices = [th.randperm(len(x))[:random.randint(1, len(x))] for x in unique_per_batch]
+            bypass_labels = [x[bypass_label_indices[idx]] for idx, x in enumerate(unique_per_batch)]
+            for idx, bypass_label in enumerate(bypass_labels):
+                if random.random() < 0.3:
+                    bypassed_label = th.where(th.isin(label_map[idx], bypass_label), 0, label_map[idx])
+                    bypassed_image = th.where(th.isin(label_map[idx], bypass_label), batch[idx], 0)
+                else:
+                    mask = self.get_mask(image_size=(h, w))
+                    bypassed_label = th.where(mask == 1, label_map[idx], 0)
+                    bypassed_image = th.where(mask == 1, 0, batch[idx])
+                
+                label_map[idx, :, :, :] = bypassed_label
+                input_label[idx, -3:, :, :] = bypassed_image
         input_semantics = input_label.scatter_(1, label_map, 1.0)
-
+        semantic_weights = th.where(input_semantics[:, :1] == 1, 0.1, 2.0)        
         # concatenate instance map if it exists
         if 'instance' in data:
             inst_map = data['instance']
@@ -284,10 +308,21 @@ class TrainLoop:
             input_semantics = input_semantics * mask
 
         cond = {key: value for key, value in data.items() if key not in ['label', 'instance', 'path', 'label_ori']}
+        cond['semantic_weights'] = semantic_weights
         cond['y'] = input_semantics
 
         return cond
 
+    def get_mask(self, image_size):    
+        if random.random() < 0.5:
+            regular_mask = bbox2mask(image_size, random_bbox())
+            irregular_mask = brush_stroke_mask(image_size, )
+            mask = regular_mask | irregular_mask
+        else:
+            irregular_mask = brush_stroke_mask(image_size, max_loops=10)
+            mask = irregular_mask
+        return th.from_numpy(mask).permute(2,0,1)
+
     def get_edges(self, t):
         edge = th.ByteTensor(t.size()).zero_()
         edge[:, :, :, 1:] = edge[:, :, :, 1:] | (t[:, :, :, 1:] != t[:, :, :, :-1])

image_sample.py

@@ -5,10 +5,11 @@ numpy array. This can be used to produce samples for FID evaluation.
 
 import argparse
 import os
-
+import random
 import torch as th
 import torch.distributed as dist
 import torchvision as tv
+from guided_diffusion.mask import (bbox2mask, brush_stroke_mask, get_irregular_mask, random_bbox, random_cropping_bbox)
 
 from guided_diffusion.image_datasets import load_data
 
@@ -20,10 +21,30 @@ from guided_diffusion.script_util import (
     args_to_dict,
 )
 
+label_color_mapping = {
+    0: (0, 0, 0),           # No information
+    1: (128, 64, 128),      # Urban fabric
+    2: (244, 35, 232),      # Industrial, commercial, public, military, private and transport units
+    3: (70, 70, 70),        # Mine, dump and construction sites
+    4: (102, 102, 156),     # Artificial non-agricultural vegetated areas
+    5: (190, 153, 153),     # Arable land (annual crops)
+    6: (153, 153, 153),     # Permanent crops
+    7: (250, 170, 30),      # Pastures
+    8: (220, 220, 0),       # Complex and mixed cultivation patterns
+    9: (107, 142, 35),      # Orchards at the fringe of urban classes
+    10: (152, 251, 152),    # Forests
+    11: (70, 130, 180),     # Herbaceous vegetation associations
+    12: (220, 20, 60),      # Open spaces with little or no vegetation
+    13: (255, 0, 0),        # Wetlands
+    14: (0, 0, 142),        # Water
+    15: (220, 220, 220)     # Clouds and shadows
+}
+label_color_mapping_ts = th.tensor(list(label_color_mapping.values()))
+print(label_color_mapping_ts)
 
 def main():
     args = create_argparser().parse_args()
-
+    args.num_classes = args.num_classes + 3 if args.inpainting else args.num_classes
     dist_util.setup_dist()
     logger.configure()
 
@@ -59,13 +80,16 @@ def main():
     os.makedirs(label_path, exist_ok=True)
     sample_path = os.path.join(args.results_path, 'samples')
     os.makedirs(sample_path, exist_ok=True)
-
+    debug_path = os.path.join(args.results_path, 'debug')
+    os.makedirs(debug_path, exist_ok=True)
     logger.log("sampling...")
     all_samples = []
     for i, (batch, cond) in enumerate(data):
         image = ((batch + 1.0) / 2.0).cuda()
         label = (cond['label_ori'].float() / 255.0).cuda()
-        model_kwargs = preprocess_input(cond, num_classes=args.num_classes)
+        model_kwargs = preprocess_input(batch, cond, num_classes=args.num_classes, inpainting=args.inpainting)
+        debug_img = model_kwargs['debug_img']                
+        print(cond['path'])
 
         # set hyperparameter
         model_kwargs['s'] = args.s
@@ -85,14 +109,17 @@ def main():
         gathered_samples = [th.zeros_like(sample) for _ in range(dist.get_world_size())]
         dist.all_gather(gathered_samples, sample)  # gather not supported with NCCL
         all_samples.extend([sample.cpu().numpy() for sample in gathered_samples])
-
         for j in range(sample.shape[0]):
-            tv.utils.save_image(image[j], os.path.join(image_path, cond['path'][j].split('/')[-1].split('.')[0] + '.png'))
-            tv.utils.save_image(sample[j], os.path.join(sample_path, cond['path'][j].split('/')[-1].split('.')[0] + '.png'))
-            tv.utils.save_image(label[j], os.path.join(label_path, cond['path'][j].split('/')[-1].split('.')[0] + '.png'))
-
+            print(os.path.join(image_path, cond['path'][j].split('/')[-1]))
+            print(os.path.join(sample_path, cond['path'][j].split('/')[-1]))
+            print(os.path.join(label_path, cond['path'][j].split('/')[-1]))
+            print(os.path.join(debug_path, cond['path'][j].split('/')[-1]))            
+            tv.utils.save_image(image[j], os.path.join(image_path, cond['path'][j].split('/')[-1]))
+            tv.utils.save_image(sample[j], os.path.join(sample_path, cond['path'][j].split('/')[-1]))
+            tv.utils.save_image(label[j], os.path.join(label_path, cond['path'][j].split('/')[-1]))
+            tv.utils.save_image(debug_img[j], os.path.join(debug_path, cond['path'][j].split('/')[-1]))            
         logger.log(f"created {len(all_samples) * args.batch_size} samples")
-
+        # break
         if len(all_samples) * args.batch_size > args.num_samples:
             break
 
@@ -100,7 +127,7 @@ def main():
     logger.log("sampling complete")
 
 
-def preprocess_input(data, num_classes):
+def preprocess_input(batch, data, num_classes, inpainting=False):
     # move to GPU and change data types
     data['label'] = data['label'].long()
 
@@ -108,6 +135,31 @@ def preprocess_input(data, num_classes):
     label_map = data['label']
     bs, _, h, w = label_map.size()
     input_label = th.FloatTensor(bs, num_classes, h, w).zero_()
+    debug_img = th.FloatTensor(bs, h, w, 3).zero_()    
+    if inpainting:
+        unique_per_batch = [th.unique(x) for x in label_map]
+        bypass_label_indices = [th.randperm(len(x))[:random.randint(0, len(x))] for x in unique_per_batch]
+        bypass_labels = [x[bypass_label_indices[idx]] for idx, x in enumerate(unique_per_batch)]
+        print(unique_per_batch, bypass_labels)
+        for idx, bypass_label in enumerate(bypass_labels):
+            debug_norm_img = (batch[idx, :, :, :].permute(1, 2, 0) + 1.0) / 2.0
+            if random.random() < 1.0:
+                bypassed_label = th.where(th.isin(label_map[idx], bypass_label), 0, label_map[idx])
+                bypassed_image = th.where(th.isin(label_map[idx], bypass_label), batch[idx], 0) 
+                debug_img[idx, :, :, :] = th.where(th.isin(label_map[idx], bypass_label).permute(1,2,0),
+                    debug_norm_img, label_color_mapping_ts[label_map[idx].squeeze()]/255.0)
+            else:
+                mask = get_mask(image_size=(h, w))
+                bypassed_label = th.where(mask == 1, label_map[idx], 0)
+                bypassed_image = th.where(mask == 1, 0, batch[idx])
+                debug_img[idx, :, :, :] = th.where(mask.permute(1,2,0) == 1,                                                   
+                                                   label_color_mapping_ts[label_map[idx].squeeze()]/255.0,
+                                                   debug_norm_img)
+                                                   
+            
+            label_map[idx, :, :, :] = bypassed_label
+            input_label[idx, -3:, :, :] = bypassed_image
+
     input_semantics = input_label.scatter_(1, label_map, 1.0)
 
     # concatenate instance map if it exists
@@ -116,7 +168,18 @@ def preprocess_input(data, num_classes):
         instance_edge_map = get_edges(inst_map)
         input_semantics = th.cat((input_semantics, instance_edge_map), dim=1)
 
-    return {'y': input_semantics}
+    return {'y': input_semantics, 'debug_img': debug_img.permute(0, 3, 1, 2)}
+
+def get_mask(image_size):
+    if random.random() < 0.5:
+        regular_mask = bbox2mask(image_size, random_bbox())
+        irregular_mask = brush_stroke_mask(image_size, )
+        mask = regular_mask | irregular_mask
+    else:
+        irregular_mask = brush_stroke_mask(image_size, max_loops=10)
+        mask = irregular_mask
+
+    return th.from_numpy(mask).permute(2,0,1)
 
 
 def get_edges(t):
@@ -139,7 +202,8 @@ def create_argparser():
         model_path="",
         results_path="",
         is_train=False,
-        s=1.0
+        s=1.0,
+        inpainting=False
     )
     defaults.update(model_and_diffusion_defaults())
     parser = argparse.ArgumentParser()

image_train.py

@@ -19,7 +19,7 @@ from guided_diffusion.train_util import TrainLoop
 
 def main():
     args = create_argparser().parse_args()
-
+    args.num_classes = args.num_classes + 3 if args.inpainting else args.num_classes
     dist_util.setup_dist()
     logger.configure()
 
@@ -27,6 +27,7 @@ def main():
     model, diffusion = create_model_and_diffusion(
         **args_to_dict(args, model_and_diffusion_defaults().keys())
     )
+    print(dist_util.dev())
     model.to(dist_util.dev())
     schedule_sampler = create_named_schedule_sampler(args.schedule_sampler, diffusion)
 
@@ -59,6 +60,7 @@ def main():
         schedule_sampler=schedule_sampler,
         weight_decay=args.weight_decay,
         lr_anneal_steps=args.lr_anneal_steps,
+        inpainting=args.inpainting
     ).run_loop()
 
 
@@ -79,7 +81,8 @@ def create_argparser():
         resume_checkpoint="",
         use_fp16=False,
         fp16_scale_growth=1e-3,
-        is_train=True
+        is_train=True,
+        inpainting=False
     )
     defaults.update(model_and_diffusion_defaults())
     parser = argparse.ArgumentParser()

scripts/miniFrance.sh

@@ -0,0 +1,11 @@
+#!/bin/bash
+
+# Training
+export OPENAI_LOGDIR='OUTPUT/MiniFrance-SDM-256CH'
+python image_train.py --data_dir ../output2/balanced_patches2 --dataset_mode miniFrance --lr 1e-4 \
+    --batch_size 4 \
+    --attention_resolutions 32,16,8 --diffusion_steps 1000 --image_size 128 --learn_sigma True \
+    --noise_schedule linear --num_channels 256 --num_head_channels 64 --num_res_blocks 2 \
+    --resblock_updown True --use_fp16 True --use_scale_shift_norm True \
+    --use_checkpoint True --num_classes 16 \
+    --class_cond True --no_instance True

scripts/miniFrance_inpaint.sh

@@ -0,0 +1,12 @@
+#!/bin/bash
+
+# Inpainting Finetune
+export OPENAI_LOGDIR='OUTPUT/MiniFrance-SDM-256CH-INPAINT'
+python image_train.py --data_dir ../output2/balanced_patches2 --dataset_mode miniFrance --lr 2e-5 \
+    --batch_size 1 --attention_resolutions 32,16,8 --diffusion_steps 1000 --image_size 256 \
+    --learn_sigma True --noise_schedule linear --num_channels 256 --num_head_channels 64 \
+    --num_res_blocks 2 --resblock_updown True --use_fp16 True --use_scale_shift_norm True \
+    --use_checkpoint True --num_classes 16 \
+    --class_cond True --no_instance True \
+    --resume_checkpoint OUTPUT/MiniFrance-SDM-256CH/model_best.pt \
+    --inpainting True

scripts/miniFrance_test.sh

@@ -0,0 +1,12 @@
+#!/bin/bash
+
+# Testing
+export OPENAI_LOGDIR='OUTPUT/MiniFrance-SDM-256CH-TEST-INPAINT'
+python image_sample.py --data_dir ../output2/balanced_patches2 --dataset_mode miniFrance \
+    --attention_resolutions 32,16,8 --diffusion_steps 1000 --image_size 256 --learn_sigma True \
+    --noise_schedule linear --num_channels 256 --num_head_channels 64 --num_res_blocks 2 \
+    --resblock_updown True --use_fp16 True --use_scale_shift_norm True --num_classes 16 \
+    --class_cond True --no_instance True --batch_size 1 --num_samples 6 \
+    --model_path OUTPUT/MiniFrance-SDM-256CH-INPAINT/model016000.pt \
+    --results_path RESULTS/MiniFrance-SDM-256CH-INPAINT --s 1.5 \
+    --inpainting True