Spaces:

Kutluhan
/

SemanticInpaint

Build error

A new example is added.

5e56207 over 1 year ago

16.8 kB

import gradio as gr
import numpy as np
import argparse
import os
import random
import torch as th
th.set_float32_matmul_precision('high')
import torchvision as tv

from guided_diffusion import logger
from guided_diffusion.script_util import (
    model_and_diffusion_defaults,
    create_model_and_diffusion,
    recreate_diffusion_with_steps,
    add_dict_to_argparser,
    args_to_dict,
)
import argparse

import warnings
from torchvision import transforms
import torch.nn.functional as F
import huggingface_hub
import functools
# import spaces

# Check if running in Hugging Face Spaces
is_spaces = os.environ.get("SPACE_ID") is not None
spaces_decorator = lambda f: f  # Default no-op decorator

if is_spaces:
    try:
        import spaces
        print("Running in Hugging Face Spaces environment")
        # Create a decorator based on FREE_GPU environment variable
        free_gpu = os.environ.get("FREE_GPU", "false").lower() == "true"
        spaces_decorator = functools.partial(spaces.GPU, duration=480, free_gpu=free_gpu)
    except ImportError:
        print("spaces module not found, continuing without it")
        is_spaces = False

model_path = huggingface_hub.hf_hub_download("Kutluhan/SemanticInpaint", "model008000.pt")


def create_argparser():
    defaults = dict(
        data_dir="",
        dataset_mode="",
        clip_denoised=True,
        num_samples=10000,
        batch_size=1,
        use_ddim=False,
        model_path=model_path,
        results_path="",
        is_train=False,
        s=1.5,
        inpainting=True
    )
    defaults.update(model_and_diffusion_defaults())
    parser = argparse.ArgumentParser()
    add_dict_to_argparser(parser, defaults)
    return parser
    
args = create_argparser().parse_args()
print(args)
def load_model():
    print(args)
    args.num_classes = args.num_classes + 3 if args.inpainting else args.num_classes
    
    logger.configure()

    logger.log("creating model and diffusion...")
    model, diffusion = create_model_and_diffusion(
        **args_to_dict(args, model_and_diffusion_defaults().keys())
    )
    model.load_state_dict(
        th.load(args.model_path, map_location="cpu")
    )
    model.to("cuda")
    return model, diffusion

model, diffusion = load_model()
model.convert_to_fp16()
model.eval()

# Check if torch compile should be enabled via environment variable
use_torch_compile = os.environ.get("USE_TORCH_COMPILE", "0").lower() == "1"
if use_torch_compile:
    try:
        print("Enabling PyTorch compilation mode to improve performance")
        model = th.compile(model, mode="reduce-overhead", fullgraph=True)
    except Exception as e:
        print(f"Failed to enable PyTorch compilation mode: {e}")
        print("Continuing without compilation")

# Label to color mapping
label_color_mapping = {
    0: (0, 0, 0),           # No information
    1: (128, 64, 128),      # Urban fabric
    2: (244, 35, 232),      # Industrial / commercial
    3: (70, 70, 70),        # Mine / dump / construction
    4: (102, 102, 156),     # Artificial non-agricultural
    5: (190, 153, 153),     # Arable land
    6: (153, 153, 153),     # Permanent crops
    7: (250, 170, 30),      # Pastures
    8: (220, 220, 0),       # Complex & mixed cultivation
    9: (107, 142, 35),      # Orchards
    10: (152, 251, 152),    # Forests
    11: (70, 130, 180),     # Herbaceous vegetation
    12: (220, 20, 60),      # Open spaces
    13: (255, 0, 0),        # Wetlands
    14: (0, 0, 142),        # Water
    15: (220, 220, 220),    # Clouds & shadows
}

label_to_name = {
    0: "No information",
    1: "Urban fabric",
    2: "Industrial / commercial",
    3: "Mine / dump / construction",
    4: "Artificial non-agricultural",
    5: "Arable land",
    6: "Permanent crops",
    7: "Pastures",
    8: "Complex & mixed cultivation",
    9: "Orchards",
    10: "Forests",
    11: "Herbaceous vegetation",
    12: "Open spaces",
    13: "Wetlands",
    14: "Water",
    15: "Clouds & shadows"
}

name_to_label = {v: k for k, v in label_to_name.items()}

label_color_mapping_ts = th.tensor(list(label_color_mapping.values()))

def create_drawing_canvas():
    """Create a blank white canvas."""
    canvas = np.ones((256, 256, 3), dtype=np.uint8) * 0
    return canvas

def update_drawing_color(label):
    """
    Update the ImageEditor's brush color based on the 
    chosen label in the dropdown and show helpful hints.
    """
    # Extract label index from dropdown (e.g. "Water [USE ALONE]: (0, 0, 142)" -> 14)
    # Handle labels with tags by splitting on ":" and then extracting the base name
    label_part = label.split(":")[0].strip()
    
    # Extract the base label name without tags
    if "[USE ALONE]" in label_part:
        base_label = label_part.replace("[USE ALONE]", "").strip()
    elif "[DONT USE]" in label_part:
        base_label = label_part.replace("[DONT USE]", "").strip()
    else:
        base_label = label_part
    
    # Get the label index
    label_idx = name_to_label.get(base_label, 0)
    
    # Get RGB color from dictionary, then convert to hex (#RRGGBB)
    color = label_color_mapping[label_idx]
    hex_color = "#{:02x}{:02x}{:02x}".format(*color)
    
    # Define message types and hints for each label
    label_hints = {
        # Multi-Label Combinations That Work Well (info)
        1: ("info", "Urban Fabric: Tends to generate city-like structures and dense urban layouts. Works well with Industrial/Commercial."),
        2: ("info", "Industrial/Commercial: Often produces road networks and industrial complexes. Draw thin lines to simulate roads when using this label."),
        6: ("info", "Permanent Crops: Generates brown or green agricultural fields."),
        10: ("info", "Forests: Creates dense forested regions."),
        14: ("info", "Water: Primarily generates lakes (greenish color) and occasionally seas (bluish color). Combines well with Wetlands."),
        4: ("info", "Artificial Non-Agricultural Vegetated Areas: Complements urban layouts."),
        
        # Labels That Work Well Alone (info)
        7: ("info", "Pastures: Best used independently."),
        5: ("info", "Arable Land: Works well individually."),
        11: ("info", "Herbaceous Vegetation: Sometimes generates white snowy textures. Best used alone."),
        
        # Labels That Do Not Perform Well (warning)
        3: ("warning", "Mine/Dump/Construction: Limited generation quality due to insufficient data."),
        8: ("warning", "Complex & Mixed Cultivation Patterns: Inconsistent results."),
        9: ("warning", "Orchards: Often unreliable."),
        12: ("warning", "Open Spaces with Little or No Vegetation: Sparse or incoherent generation."),
        15: ("warning", "Clouds & Shadows: Limited performance due to insufficient data."),
        
        # Recommended Combinations (info)
        13: ("info", "Wetlands: Works best when combined with Water to generate more realistic transitions."),
        
        # Empty for other labels
        0: ("", ""),
    }
    
    # Get hint for current label
    hint_type, hint_text = label_hints.get(label_idx, ("", ""))
    
    # Return an update for the existing ImageEditor component and the hint
    return gr.ImageEditor(
        label="Semantic Drawing",
        container=True,
        brush=gr.Brush(
            colors=[f'#{c[0]:02x}{c[1]:02x}{c[2]:02x}' for c in label_color_mapping.values()],
            default_color=hex_color,
            color_mode='fixed',
            default_size=20
        ),
        interactive=True
    ), hex_color, hint_type, hint_text

tfs_label = transforms.Compose([
    transforms.ToPILImage(),
    transforms.Resize((256, 256), interpolation=transforms.InterpolationMode.NEAREST),
    transforms.ToTensor(),
    # transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
])
def copy_output_to_input(output_image):
    return output_image
    
# Define the generate_image function with conditional decorator
@spaces_decorator
def generate_image(input_image, semantic_drawing, num_imgs):
    """
    Generate image using the model with adjustable prob_mask parameter.
    """    
    print("num_imgs", num_imgs)
    tfs = transforms.Compose([
        transforms.ToPILImage(),
        transforms.Resize((256, 256)),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
    ])
    diffusion_steps = 1000
    print(diffusion_steps)
    img = tfs(input_image).unsqueeze(0)
    # Convert semantic_drawing to numpy array if it's not already
    # 256x256x3
    raw_label = semantic_drawing["composite"][:,:,:3].copy()
    # Convert RGB colors to semantic indexes
    semantic_label = np.zeros((raw_label.shape[0], raw_label.shape[1]), dtype=np.uint8)
    for label_idx, color in label_color_mapping.items():
        # Find pixels that match this color exactly
        mask = np.all(raw_label == color, axis=2)
        semantic_label[mask] = label_idx
    
    # Check if any pixels couldn't be mapped to a label
    unmapped = np.all(np.logical_not(np.any(
        [np.all(raw_label == color, axis=2) for color in label_color_mapping.values()],
        axis=0
    )), axis=0)
    
    if np.any(unmapped):
        raise ValueError(f"Found colors in the semantic drawing that don't match any label color")
    semantic_label= semantic_label.reshape(1, 1, 256, 256)
    input_label = th.FloatTensor(1, 19, 256, 256).zero_()    
    input_image = th.tensor(np.where(semantic_label == 0, img, 0))
    input_label[0, -3:, :, :] = input_image
    input_semantics = input_label.scatter_(1, th.tensor(semantic_label).long(), 1.0)
    model_kwargs = {'y': input_semantics.tile(num_imgs, 1, 1, 1),}
    model_kwargs['s'] = args.s
    print(args.s)
    
    sample_fn = (
        diffusion.p_sample_loop if not args.use_ddim else diffusion.ddim_sample_loop
    )
    
    output_imgs = []    
    sample = sample_fn(
        model,
        (num_imgs, 3, 256, 256),
        clip_denoised=args.clip_denoised,
        model_kwargs=model_kwargs,
        progress=True,
    )
    sample = (sample + 1) / 2.0
    sample = sample.cpu().permute(0, 2, 3, 1)
    for num in range(num_imgs):
        output_img = sample.cpu().numpy()[num]
        output_imgs.append(output_img)
    
    return output_imgs  # return twice: for output and new input

def put_clicked_into_input(evt: gr.SelectData):
    """
    evt.value  -> value associated with the clicked item
                  (a NumPy array *or* a list, gradio versions differ)
    evt.index  -> zero-based index of the thumbnail that was clicked
    """
    clicked = evt.value
    image = clicked['image']['path']
    if isinstance(clicked, list):          # safety for older gradio versions
        clicked = clicked[evt.index]
    return image                         # single image (NumPy array)

def put_selected_into_input(chosen_img):
    print('chosen_img', chosen_img)
    if chosen_img is None:
        raise gr.Error("Please click an image in the gallery first.")
    return chosen_img
    
with gr.Blocks() as demo:
    gr.Markdown("# Satellite Image Semantic Inpainting with DDPM")
    
    # Add recommendations and best practices
    with gr.Accordion("Usage Guidelines & Best Practices", open=True):
        gr.Markdown("""
        ### General Recommendations
        
        - **Generate Multiple Samples**: Generate 2-4 samples to explore different variations. Not all outputs are visually satisfying, so multiple samples increase your chance of getting high-quality results.
        
        ### Best Practices
        
        - **Use Few Classes at a Time**: Start with 1-2 semantic classes, draw them on the mask, and generate images.
        - **Iterative Editing**: Select the best image from the generated samples as your new base, then add 1-2 more semantic classes and repeat.
        - **Preserve Context**: Using only 1-2 new mask types at a time helps maintain better visual and semantic context.
        - **Avoid Large Masks**: Drawing very large masks disrupts spatial coherence. Prefer smaller, targeted edits for better realism.
        
        ### Label Indications
        - Labels marked with **[USE ALONE]** work best independently.
        - Labels marked with **[DONT USE]** generally produce poor results.
        - Unmarked labels can be combined for interesting effects.
        """)
    
    # Create modified label names with tags
    modified_label_names = {}
    for k, v in label_to_name.items():
        if k in [7, 5, 11]:  # Labels that work well alone
            modified_label_names[k] = f"{v} [USE ALONE]"
        elif k in [3, 8, 9, 12, 15]:  # Labels that don't perform well
            modified_label_names[k] = f"{v} [DONT USE]"
        else:
            modified_label_names[k] = v
    
    with gr.Row():
        with gr.Column(scale=1):
            input_image = gr.Image(
                label="Input Image",
                type="numpy",
                container=True,
                height=350
            )
        with gr.Column(scale=1):
            semantic_drawing = gr.ImageEditor(
                label="Semantic Drawing",
                container=True,
                brush=gr.Brush(
                    colors=[f'#{c[0]:02x}{c[1]:02x}{c[2]:02x}' for c in label_color_mapping.values()],
                    default_color='#000000',
                    color_mode='fixed',
                    default_size=20
                ),
                interactive=True,
                height=350
            )

    with gr.Row(equal_height=True):
        with gr.Column(scale=1):
            label_dropdown = gr.Dropdown(
                choices=[f"{modified_label_names[k]}: {label_color_mapping[k]}" for k in label_color_mapping],
                label="Select Label",
                value=f"{modified_label_names[0]}: {label_color_mapping[0]}",
            )
        with gr.Column(scale=1):
            color_display = gr.ColorPicker(
                label="Selected Color",
                value="#000000",
                interactive=False,
                container=True
            )
    
    with gr.Row():
        num_imgs_slider = gr.Slider(
            minimum=1, maximum=4, value=1, step=1,
            label="Images to generate"
        )
    with gr.Row():
        generate_btn = gr.Button("Generate Image", size="large")
       
    
    with gr.Row():
        output_image = gr.Gallery(
            label="Generated images",
            height="auto", columns=[4],         # 4 columns looks nice up to 8 images
            preview=False,
        )

        selected_img = gr.State(value=None)
    with gr.Row():
        use_as_input_btn = gr.Button("Use Output as New Input")
        
    # Example images
    gr.Examples(
        examples=[
            ["examples/example1.png", create_drawing_canvas()],
            ["examples/example2.png", create_drawing_canvas()],
            ["examples/example3.png", create_drawing_canvas()],
            ["examples/example4.png", "./examples/example4_mask.png"],
            ["examples/example5.png", "./examples/example5_mask.png"],
            ["examples/example6.png", "./examples/example6_mask.png"],
            ["examples/example7.png", "./examples/example7_mask.png"],
            ["examples/example8.png", "./examples/example8_mask.png"],
        ],
        inputs=[input_image, semantic_drawing]
    )

    # Connect the generate button
    generate_btn.click(
        fn=generate_image,
        inputs=[input_image, semantic_drawing, num_imgs_slider],
        outputs=output_image
    )
    
    output_image.select(
        fn=put_clicked_into_input,
        inputs=None,       # the gallery value is the *selected image*, not the list
        outputs=selected_img        # gr.Image expects a single ndarray → OK
    )

    
    use_as_input_btn.click(
        fn=put_selected_into_input,
        inputs=selected_img,     # comes from our hidden State
        outputs=input_image
    )

    
    
    # Add hint components (hidden from UI but used for handling hints)
    hint_type = gr.Textbox(visible=False)
    hint_text = gr.Textbox(visible=False)

    # Function to display appropriate notification based on hint type and text
    def show_notification(hint_type, hint_text):
        if hint_type == "warning":
            return gr.Warning(hint_text)
        elif hint_type == "info":
            return gr.Info(hint_text)
        return None

    # Connect the dropdown to update functions
    label_dropdown.change(
        fn=update_drawing_color,
        inputs=label_dropdown,
        outputs=[semantic_drawing, color_display, hint_type, hint_text]
    ).then(
        fn=show_notification,
        inputs=[hint_type, hint_text],
        outputs=None
    )

if __name__ == "__main__":
    demo.launch()

	import gradio as gr
	import numpy as np
	import argparse
	import os
	import random
	import torch as th
	th.set_float32_matmul_precision('high')
	import torchvision as tv

	from guided_diffusion import logger
	from guided_diffusion.script_util import (
	model_and_diffusion_defaults,
	create_model_and_diffusion,
	recreate_diffusion_with_steps,
	add_dict_to_argparser,
	args_to_dict,
	)
	import argparse

	import warnings
	from torchvision import transforms
	import torch.nn.functional as F
	import huggingface_hub
	import functools
	# import spaces

	# Check if running in Hugging Face Spaces
	is_spaces = os.environ.get("SPACE_ID") is not None
	spaces_decorator = lambda f: f # Default no-op decorator

	if is_spaces:
	try:
	import spaces
	print("Running in Hugging Face Spaces environment")
	# Create a decorator based on FREE_GPU environment variable
	free_gpu = os.environ.get("FREE_GPU", "false").lower() == "true"
	spaces_decorator = functools.partial(spaces.GPU, duration=480, free_gpu=free_gpu)
	except ImportError:
	print("spaces module not found, continuing without it")
	is_spaces = False

	model_path = huggingface_hub.hf_hub_download("Kutluhan/SemanticInpaint", "model008000.pt")


	def create_argparser():
	defaults = dict(
	data_dir="",
	dataset_mode="",
	clip_denoised=True,
	num_samples=10000,
	batch_size=1,
	use_ddim=False,
	model_path=model_path,
	results_path="",
	is_train=False,
	s=1.5,
	inpainting=True
	)
	defaults.update(model_and_diffusion_defaults())
	parser = argparse.ArgumentParser()
	add_dict_to_argparser(parser, defaults)
	return parser

	args = create_argparser().parse_args()
	print(args)
	def load_model():
	print(args)
	args.num_classes = args.num_classes + 3 if args.inpainting else args.num_classes

	logger.configure()

	logger.log("creating model and diffusion...")
	model, diffusion = create_model_and_diffusion(
	**args_to_dict(args, model_and_diffusion_defaults().keys())
	)
	model.load_state_dict(
	th.load(args.model_path, map_location="cpu")
	)
	model.to("cuda")
	return model, diffusion

	model, diffusion = load_model()
	model.convert_to_fp16()
	model.eval()

	# Check if torch compile should be enabled via environment variable
	use_torch_compile = os.environ.get("USE_TORCH_COMPILE", "0").lower() == "1"
	if use_torch_compile:
	try:
	print("Enabling PyTorch compilation mode to improve performance")
	model = th.compile(model, mode="reduce-overhead", fullgraph=True)
	except Exception as e:
	print(f"Failed to enable PyTorch compilation mode: {e}")
	print("Continuing without compilation")

	# Label to color mapping
	label_color_mapping = {
	0: (0, 0, 0), # No information
	1: (128, 64, 128), # Urban fabric
	2: (244, 35, 232), # Industrial / commercial
	3: (70, 70, 70), # Mine / dump / construction
	4: (102, 102, 156), # Artificial non-agricultural
	5: (190, 153, 153), # Arable land
	6: (153, 153, 153), # Permanent crops
	7: (250, 170, 30), # Pastures
	8: (220, 220, 0), # Complex & mixed cultivation
	9: (107, 142, 35), # Orchards
	10: (152, 251, 152), # Forests
	11: (70, 130, 180), # Herbaceous vegetation
	12: (220, 20, 60), # Open spaces
	13: (255, 0, 0), # Wetlands
	14: (0, 0, 142), # Water
	15: (220, 220, 220), # Clouds & shadows
	}

	label_to_name = {
	0: "No information",
	1: "Urban fabric",
	2: "Industrial / commercial",
	3: "Mine / dump / construction",
	4: "Artificial non-agricultural",
	5: "Arable land",
	6: "Permanent crops",
	7: "Pastures",
	8: "Complex & mixed cultivation",
	9: "Orchards",
	10: "Forests",
	11: "Herbaceous vegetation",
	12: "Open spaces",
	13: "Wetlands",
	14: "Water",
	15: "Clouds & shadows"
	}

	name_to_label = {v: k for k, v in label_to_name.items()}

	label_color_mapping_ts = th.tensor(list(label_color_mapping.values()))

	def create_drawing_canvas():
	"""Create a blank white canvas."""
	canvas = np.ones((256, 256, 3), dtype=np.uint8) * 0
	return canvas

	def update_drawing_color(label):
	"""
	Update the ImageEditor's brush color based on the
	chosen label in the dropdown and show helpful hints.
	"""
	# Extract label index from dropdown (e.g. "Water [USE ALONE]: (0, 0, 142)" -> 14)
	# Handle labels with tags by splitting on ":" and then extracting the base name
	label_part = label.split(":")[0].strip()

	# Extract the base label name without tags
	if "[USE ALONE]" in label_part:
	base_label = label_part.replace("[USE ALONE]", "").strip()
	elif "[DONT USE]" in label_part:
	base_label = label_part.replace("[DONT USE]", "").strip()
	else:
	base_label = label_part

	# Get the label index
	label_idx = name_to_label.get(base_label, 0)

	# Get RGB color from dictionary, then convert to hex (#RRGGBB)
	color = label_color_mapping[label_idx]
	hex_color = "#{:02x}{:02x}{:02x}".format(*color)

	# Define message types and hints for each label
	label_hints = {
	# Multi-Label Combinations That Work Well (info)
	1: ("info", "Urban Fabric: Tends to generate city-like structures and dense urban layouts. Works well with Industrial/Commercial."),
	2: ("info", "Industrial/Commercial: Often produces road networks and industrial complexes. Draw thin lines to simulate roads when using this label."),
	6: ("info", "Permanent Crops: Generates brown or green agricultural fields."),
	10: ("info", "Forests: Creates dense forested regions."),
	14: ("info", "Water: Primarily generates lakes (greenish color) and occasionally seas (bluish color). Combines well with Wetlands."),
	4: ("info", "Artificial Non-Agricultural Vegetated Areas: Complements urban layouts."),

	# Labels That Work Well Alone (info)
	7: ("info", "Pastures: Best used independently."),
	5: ("info", "Arable Land: Works well individually."),
	11: ("info", "Herbaceous Vegetation: Sometimes generates white snowy textures. Best used alone."),

	# Labels That Do Not Perform Well (warning)
	3: ("warning", "Mine/Dump/Construction: Limited generation quality due to insufficient data."),
	8: ("warning", "Complex & Mixed Cultivation Patterns: Inconsistent results."),
	9: ("warning", "Orchards: Often unreliable."),
	12: ("warning", "Open Spaces with Little or No Vegetation: Sparse or incoherent generation."),
	15: ("warning", "Clouds & Shadows: Limited performance due to insufficient data."),

	# Recommended Combinations (info)
	13: ("info", "Wetlands: Works best when combined with Water to generate more realistic transitions."),

	# Empty for other labels
	0: ("", ""),
	}

	# Get hint for current label
	hint_type, hint_text = label_hints.get(label_idx, ("", ""))

	# Return an update for the existing ImageEditor component and the hint
	return gr.ImageEditor(
	label="Semantic Drawing",
	container=True,
	brush=gr.Brush(
	colors=[f'#{c[0]:02x}{c[1]:02x}{c[2]:02x}' for c in label_color_mapping.values()],
	default_color=hex_color,
	color_mode='fixed',
	default_size=20
	),
	interactive=True
	), hex_color, hint_type, hint_text

	tfs_label = transforms.Compose([
	transforms.ToPILImage(),
	transforms.Resize((256, 256), interpolation=transforms.InterpolationMode.NEAREST),
	transforms.ToTensor(),
	# transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
	])
	def copy_output_to_input(output_image):
	return output_image

	# Define the generate_image function with conditional decorator
	@spaces_decorator
	def generate_image(input_image, semantic_drawing, num_imgs):
	"""
	Generate image using the model with adjustable prob_mask parameter.
	"""
	print("num_imgs", num_imgs)
	tfs = transforms.Compose([
	transforms.ToPILImage(),
	transforms.Resize((256, 256)),
	transforms.ToTensor(),
	transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
	])
	diffusion_steps = 1000
	print(diffusion_steps)
	img = tfs(input_image).unsqueeze(0)
	# Convert semantic_drawing to numpy array if it's not already
	# 256x256x3
	raw_label = semantic_drawing["composite"][:,:,:3].copy()
	# Convert RGB colors to semantic indexes
	semantic_label = np.zeros((raw_label.shape[0], raw_label.shape[1]), dtype=np.uint8)
	for label_idx, color in label_color_mapping.items():
	# Find pixels that match this color exactly
	mask = np.all(raw_label == color, axis=2)
	semantic_label[mask] = label_idx

	# Check if any pixels couldn't be mapped to a label
	unmapped = np.all(np.logical_not(np.any(
	[np.all(raw_label == color, axis=2) for color in label_color_mapping.values()],
	axis=0
	)), axis=0)

	if np.any(unmapped):
	raise ValueError(f"Found colors in the semantic drawing that don't match any label color")
	semantic_label= semantic_label.reshape(1, 1, 256, 256)
	input_label = th.FloatTensor(1, 19, 256, 256).zero_()
	input_image = th.tensor(np.where(semantic_label == 0, img, 0))
	input_label[0, -3:, :, :] = input_image
	input_semantics = input_label.scatter_(1, th.tensor(semantic_label).long(), 1.0)
	model_kwargs = {'y': input_semantics.tile(num_imgs, 1, 1, 1),}
	model_kwargs['s'] = args.s
	print(args.s)

	sample_fn = (
	diffusion.p_sample_loop if not args.use_ddim else diffusion.ddim_sample_loop
	)

	output_imgs = []
	sample = sample_fn(
	model,
	(num_imgs, 3, 256, 256),
	clip_denoised=args.clip_denoised,
	model_kwargs=model_kwargs,
	progress=True,
	)
	sample = (sample + 1) / 2.0
	sample = sample.cpu().permute(0, 2, 3, 1)
	for num in range(num_imgs):
	output_img = sample.cpu().numpy()[num]
	output_imgs.append(output_img)

	return output_imgs # return twice: for output and new input

	def put_clicked_into_input(evt: gr.SelectData):
	"""
	evt.value -> value associated with the clicked item
	(a NumPy array or a list, gradio versions differ)
	evt.index -> zero-based index of the thumbnail that was clicked
	"""
	clicked = evt.value
	image = clicked['image']['path']
	if isinstance(clicked, list): # safety for older gradio versions
	clicked = clicked[evt.index]
	return image # single image (NumPy array)

	def put_selected_into_input(chosen_img):
	print('chosen_img', chosen_img)
	if chosen_img is None:
	raise gr.Error("Please click an image in the gallery first.")
	return chosen_img

	with gr.Blocks() as demo:
	gr.Markdown("# Satellite Image Semantic Inpainting with DDPM")

	# Add recommendations and best practices
	with gr.Accordion("Usage Guidelines & Best Practices", open=True):
	gr.Markdown("""
	### General Recommendations

	- Generate Multiple Samples: Generate 2-4 samples to explore different variations. Not all outputs are visually satisfying, so multiple samples increase your chance of getting high-quality results.

	### Best Practices

	- Use Few Classes at a Time: Start with 1-2 semantic classes, draw them on the mask, and generate images.
	- Iterative Editing: Select the best image from the generated samples as your new base, then add 1-2 more semantic classes and repeat.
	- Preserve Context: Using only 1-2 new mask types at a time helps maintain better visual and semantic context.
	- Avoid Large Masks: Drawing very large masks disrupts spatial coherence. Prefer smaller, targeted edits for better realism.

	### Label Indications
	- Labels marked with [USE ALONE] work best independently.
	- Labels marked with [DONT USE] generally produce poor results.
	- Unmarked labels can be combined for interesting effects.
	""")

	# Create modified label names with tags
	modified_label_names = {}
	for k, v in label_to_name.items():
	if k in [7, 5, 11]: # Labels that work well alone
	modified_label_names[k] = f"{v} [USE ALONE]"
	elif k in [3, 8, 9, 12, 15]: # Labels that don't perform well
	modified_label_names[k] = f"{v} [DONT USE]"
	else:
	modified_label_names[k] = v

	with gr.Row():
	with gr.Column(scale=1):
	input_image = gr.Image(
	label="Input Image",
	type="numpy",
	container=True,
	height=350
	)
	with gr.Column(scale=1):
	semantic_drawing = gr.ImageEditor(
	label="Semantic Drawing",
	container=True,
	brush=gr.Brush(
	colors=[f'#{c[0]:02x}{c[1]:02x}{c[2]:02x}' for c in label_color_mapping.values()],
	default_color='#000000',
	color_mode='fixed',
	default_size=20
	),
	interactive=True,
	height=350
	)

	with gr.Row(equal_height=True):
	with gr.Column(scale=1):
	label_dropdown = gr.Dropdown(
	choices=[f"{modified_label_names[k]}: {label_color_mapping[k]}" for k in label_color_mapping],
	label="Select Label",
	value=f"{modified_label_names[0]}: {label_color_mapping[0]}",
	)
	with gr.Column(scale=1):
	color_display = gr.ColorPicker(
	label="Selected Color",
	value="#000000",
	interactive=False,
	container=True
	)

	with gr.Row():
	num_imgs_slider = gr.Slider(
	minimum=1, maximum=4, value=1, step=1,
	label="Images to generate"
	)
	with gr.Row():
	generate_btn = gr.Button("Generate Image", size="large")


	with gr.Row():
	output_image = gr.Gallery(
	label="Generated images",
	height="auto", columns=[4], # 4 columns looks nice up to 8 images
	preview=False,
	)

	selected_img = gr.State(value=None)
	with gr.Row():
	use_as_input_btn = gr.Button("Use Output as New Input")

	# Example images
	gr.Examples(
	examples=[
	["examples/example1.png", create_drawing_canvas()],
	["examples/example2.png", create_drawing_canvas()],
	["examples/example3.png", create_drawing_canvas()],
	["examples/example4.png", "./examples/example4_mask.png"],
	["examples/example5.png", "./examples/example5_mask.png"],
	["examples/example6.png", "./examples/example6_mask.png"],
	["examples/example7.png", "./examples/example7_mask.png"],
	["examples/example8.png", "./examples/example8_mask.png"],
	],
	inputs=[input_image, semantic_drawing]
	)

	# Connect the generate button
	generate_btn.click(
	fn=generate_image,
	inputs=[input_image, semantic_drawing, num_imgs_slider],
	outputs=output_image
	)

	output_image.select(
	fn=put_clicked_into_input,
	inputs=None, # the gallery value is the selected image, not the list
	outputs=selected_img # gr.Image expects a single ndarray → OK
	)


	use_as_input_btn.click(
	fn=put_selected_into_input,
	inputs=selected_img, # comes from our hidden State
	outputs=input_image
	)



	# Add hint components (hidden from UI but used for handling hints)
	hint_type = gr.Textbox(visible=False)
	hint_text = gr.Textbox(visible=False)

	# Function to display appropriate notification based on hint type and text
	def show_notification(hint_type, hint_text):
	if hint_type == "warning":
	return gr.Warning(hint_text)
	elif hint_type == "info":
	return gr.Info(hint_text)
	return None

	# Connect the dropdown to update functions
	label_dropdown.change(
	fn=update_drawing_color,
	inputs=label_dropdown,
	outputs=[semantic_drawing, color_display, hint_type, hint_text]
	).then(
	fn=show_notification,
	inputs=[hint_type, hint_text],
	outputs=None
	)

	if __name__ == "__main__":
	demo.launch()