app.py

import gradio as gr
import pandas as pd
import numpy as np
import torch as th
from house_diffusion import dist_util
from house_diffusion.script_util import (
    create_model_and_diffusion,
)
from PIL import Image

import io
import drawSvg as drawsvg
import cairosvg
from tqdm import tqdm

import webcolors
import tempfile
from pathlib import Path
import shutil
import os


ROOM_CLASS = {
    'Living Room': 1, 'Kitchen': 2, 'Bedroom': 3, 'Bathroom': 4,
    'Balcony': 5, 'Entrance': 6, 'Dining Room': 7, 'Study Room': 8,
    'Storage': 10, 'Front Door': 11, 'Unknown': 13, 'Interior Door': 12
}
ROOM_CATEGORIES = {
    'Living Room': 1, 'Kitchen': 2, 'Bedroom': 3, 'Bathroom': 4,
    'Balcony': 5, 'Entrance': 6, 'Dining Room': 7, 'Study Room': 8,
    'Storage': 10, 'Front Door': 11, 'Other': 13
}


def save_samples(
        sample, ext, model_kwargs,
        tmp_count, num_room_types,
        # save_gif=False,
        save_gif=True,
        door_indices=[11, 12, 13], ID_COLOR=None,
        is_syn=False, draw_graph=False, save_svg=False, metrics=False):
    prefix = 'syn_' if is_syn else ''
    graph_errors = []

    print(sample.shape)

    if not save_gif:
        sample = sample[-1:]
    for i in tqdm(range(sample.shape[1])):
        resolution = 256
        images = []
        images2 = []
        images3 = []
        for k in range(sample.shape[0]):
            draw_color = drawsvg.Drawing(resolution, resolution, displayInline=False)
            draw_color.append(drawsvg.Rectangle(0, 0, resolution, resolution, fill='white'))
            polys = []
            types = []
            for j, point in (enumerate(sample[k][i])):
                if model_kwargs[f'{prefix}src_key_padding_mask'][i][j] == 1:
                    continue
                point = point.cpu().data.numpy()
                if j == 0:
                    poly = []
                if j > 0 and (model_kwargs[f'{prefix}room_indices'][i, j] != model_kwargs[f'{prefix}room_indices'][
                    i, j - 1]).any():
                    polys.append(poly)
                    types.append(c)
                    poly = []
                pred_center = False
                if pred_center:
                    point = point / 2 + 1
                    point = point * resolution // 2
                else:
                    point = point / 2 + 0.5
                    point = point * resolution
                poly.append((point[0], point[1]))
                c = np.argmax(model_kwargs[f'{prefix}room_types'][i][j - 1].cpu().numpy())
            polys.append(poly)
            types.append(c)
            for poly, c in zip(polys, types):
                if c in door_indices or c == 0:
                    continue
                room_type = c
                c = webcolors.hex_to_rgb(ID_COLOR[c])
                draw_color.append(
                    drawsvg.Lines(*np.array(poly).flatten().tolist(), close=True, fill=ID_COLOR[room_type],
                                  fill_opacity=1.0, stroke='black', stroke_width=1))

            for poly, c in zip(polys, types):
                if c not in door_indices:
                    continue
                room_type = c
                c = webcolors.hex_to_rgb(ID_COLOR[c])

                # TODO --------------------------------------------------------------------------------------
                # https://github.com/sakmalh/house_diffusion
                line_lengths = [np.linalg.norm(np.array(poly[i]) - np.array(poly[(i + 1) % len(poly)])) for i in
                                range(len(poly))]

                if metrics:
                    text_size = 5
                    for z, length in enumerate(line_lengths):
                        # Calculate the mid-point of the line segment
                        midpoint = ((poly[z][0] + poly[(z + 1) % len(poly)][0]) / 2,
                                    (poly[z][1] + poly[(z + 1) % len(poly)][1]) / 2)

                        # Calculate x and y differences
                        x_diff = poly[z][0] - poly[(z + 1) % len(poly)][0]
                        y_diff = poly[z][1] - poly[(z + 1) % len(poly)][1]

                        # Determine text position adjustments based on differences
                        if int(y_diff) != 0:
                            if y_diff > 0:
                                text_x = midpoint[0] + text_size
                                text_y = midpoint[1]

                                draw_color.append(drawsvg.Line(
                                    text_x, text_y + text_size,  # Start point at the text label
                                            poly[z][0] + text_size, poly[z][1],  # End point at the polygon endpoint
                                    stroke='black',
                                    stroke_width=1
                                ))

                                draw_color.append(drawsvg.Line(
                                    text_x, text_y - text_size,  # Start point at the text label
                                            poly[(z + 1) % len(poly)][0] + text_size, poly[(z + 1) % len(poly)][1],
                                    # End point at the polygon endpoint
                                    stroke='black',
                                    stroke_width=1
                                ))
                            else:
                                text_x = midpoint[0] - text_size
                                text_y = midpoint[1]

                                draw_color.append(drawsvg.Line(
                                    text_x, text_y - text_size,  # Start point at the text label
                                            poly[z][0] - text_size, poly[z][1],  # End point at the polygon endpoint
                                    stroke='black',
                                    stroke_width=1
                                ))

                                draw_color.append(drawsvg.Line(
                                    text_x, text_y + text_size,  # Start point at the text label
                                            poly[(z + 1) % len(poly)][0] - text_size, poly[(z + 1) % len(poly)][1],
                                    # End point at the polygon endpoint
                                    stroke='black',
                                    stroke_width=1
                                ))
                        else:
                            if x_diff > 0:
                                text_x = midpoint[0]
                                text_y = midpoint[1] - text_size

                                draw_color.append(drawsvg.Line(
                                    text_x + text_size, text_y,  # Start point at the text label
                                    poly[z][0], poly[z][1] - text_size,  # End point at the polygon endpoint
                                    stroke='black',
                                    stroke_width=1
                                ))

                                draw_color.append(drawsvg.Line(
                                    text_x - text_size, text_y,  # Start point at the text label
                                    poly[(z + 1) % len(poly)][0], poly[(z + 1) % len(poly)][1] - text_size,
                                    # End point at the polygon endpoint
                                    stroke='black',
                                    stroke_width=1
                                ))
                            else:
                                text_x = midpoint[0]
                                text_y = midpoint[1] + text_size

                                draw_color.append(drawsvg.Line(
                                    text_x - text_size, text_y,  # Start point at the text label
                                    poly[z][0], poly[z][1] + text_size,  # End point at the polygon endpoint
                                    stroke='black',
                                    stroke_width=1
                                ))

                                draw_color.append(drawsvg.Line(
                                    text_x + text_size, text_y,  # Start point at the text label
                                    poly[(z + 1) % len(poly)][0], poly[(z + 1) % len(poly)][1] + text_size,
                                    # End point at the polygon endpoint
                                    stroke='black',
                                    stroke_width=1
                                ))

                        # Add the text label to the SVG
                        draw_color.append(
                            drawsvg.Text(
                                f'{int(abs(length))}',  # Format the length to two decimal places
                                text_size,
                                text_x, text_y,
                                fill='black',
                                text_anchor='middle',
                                alignment_baseline='middle'
                            )
                        )

                draw_color.append(
                    drawsvg.Lines(*np.array(poly).flatten().tolist(), close=True, fill=ID_COLOR[room_type],
                                  fill_opacity=1.0, stroke='black', stroke_width=1))

            if k == sample.shape[0] - 1 or True:
                if save_svg:
                    # draw_color.saveSvg(f'outputs/{ext}/{tmp_count + i}c_{k}_{ext}.svg')
                    return draw_color
                else:
                    Image.open(io.BytesIO(cairosvg.svg2png(draw_color.asSvg()))).save(
                        f'outputs/{ext}/{tmp_count + i}c_{ext}.png')

        # if save_gif:
        #     imageio.mimwrite(f'outputs/gif/{tmp_count + i}.gif', images, fps=10, loop=1)
        #     imageio.mimwrite(f'outputs/gif/{tmp_count + i}_v2.gif', images2, fps=10, loop=1)
        #     imageio.mimwrite(f'outputs/gif/{tmp_count + i}_v3.gif', images3, fps=10, loop=1)
    return graph_errors


def function_test(org_graphs, corners, room_type):
    get_one_hot = lambda x, z: np.eye(z)[x]
    max_num_points = 100

    house = []
    corner_bounds = []
    num_points = 0

    for i, room in enumerate(room_type):
        # Adding conditions
        num_room_corners = corners[i]
        rtype = np.repeat(np.array([get_one_hot(room, 25)]), num_room_corners, 0)
        room_index = np.repeat(np.array([get_one_hot(len(house) + 1, 32)]), num_room_corners, 0)
        corner_index = np.array([get_one_hot(x, 32) for x in range(num_room_corners)])
        # Src_key_padding_mask
        padding_mask = np.repeat(1, num_room_corners)
        padding_mask = np.expand_dims(padding_mask, 1)
        # Generating corner bounds for attention masks
        connections = np.array([[i, (i + 1) % num_room_corners] for i in range(num_room_corners)])
        connections += num_points
        corner_bounds.append([num_points, num_points + num_room_corners])
        num_points += num_room_corners
        room = np.concatenate((np.zeros([num_room_corners, 2]), rtype, corner_index, room_index,
                               padding_mask, connections), 1)
        house.append(room)

    house_layouts = np.concatenate(house, 0)
    padding = np.zeros((max_num_points - len(house_layouts), 94))
    gen_mask = np.ones((max_num_points, max_num_points))
    gen_mask[:len(house_layouts), :len(house_layouts)] = 0
    house_layouts = np.concatenate((house_layouts, padding), 0)

    door_mask = np.ones((max_num_points, max_num_points))
    self_mask = np.ones((max_num_points, max_num_points))
    for i, room in enumerate(room_type):
        if room == 1:
            living_room_index = i
            break
    for i in range(len(corner_bounds)):
        is_connected = False
        for j in range(len(corner_bounds)):
            if i == j:
                self_mask[corner_bounds[i][0]:corner_bounds[i][1], corner_bounds[j][0]:corner_bounds[j][1]] = 0
            elif any(np.equal([i, 1, j], org_graphs).all(1)) or any(np.equal([j, 1, i], org_graphs).all(1)):
                door_mask[corner_bounds[i][0]:corner_bounds[i][1], corner_bounds[j][0]:corner_bounds[j][1]] = 0
                is_connected = True
        if not is_connected:
            door_mask[corner_bounds[i][0]:corner_bounds[i][1],
            corner_bounds[living_room_index][0]:corner_bounds[living_room_index][1]] = 0

    syn_houses = house_layouts
    syn_door_masks = door_mask
    syn_self_masks = self_mask
    syn_gen_masks = gen_mask

    syn_graph = np.concatenate((org_graphs, np.zeros([200 - len(org_graphs), 3])), 0)

    cond = {
        'syn_door_mask': syn_door_masks,
        'syn_self_mask': syn_self_masks,
        'syn_gen_mask': syn_gen_masks,
        'syn_room_types': syn_houses[:, 2:2 + 25],
        'syn_corner_indices': syn_houses[:, 2 + 25:2 + 57],
        'syn_room_indices': syn_houses[:, 2 + 57:2 + 89],
        'syn_src_key_padding_mask': 1 - syn_houses[:, 2 + 89],
        'syn_connections': syn_houses[:, 2 + 90:2 + 92],
        'syn_graph': syn_graph,
    }

    return cond


def create_layout(graphs, corners, room_type, metrics=False, use_ddim=True, ddim_steps=100, num_samples=4):
    model_path = "ckpt/model250000.pt"
    steps = f"ddim{ddim_steps}"
    args = {
        "input_channels": 18,
        "condition_channels": 89,
        "num_channels": 512,
        "out_channels": 2,
        "dataset": "rplan",
        "use_checkpoint": False,
        "use_unet": False,
        "learn_sigma": False,
        "diffusion_steps": 1000,
        "noise_schedule": "cosine",
        "timestep_respacing": steps,
        "use_kl": False,
        "predict_xstart": False,
        "rescale_timesteps": False,
        "rescale_learned_sigmas": False,
        "analog_bit": False,
        "target_set": -1,
        "set_name": "",
    }

    dist_util.setup_dist()
    model, diffusion = create_model_and_diffusion(
        args['input_channels'],
        args['condition_channels'],
        args['num_channels'],
        args['out_channels'],
        args['dataset'],
        args['use_checkpoint'],
        args['use_unet'],
        args['learn_sigma'],
        args['diffusion_steps'],
        args['noise_schedule'],
        args['timestep_respacing'],
        args['use_kl'],
        args['predict_xstart'],
        args['rescale_timesteps'],
        args['rescale_learned_sigmas'],
        args['analog_bit'],
        args['target_set'],
        args['set_name'],
    )
    model.load_state_dict(
        dist_util.load_state_dict(model_path, map_location="cpu")
    )
    model.to(dist_util.dev())
    model.eval()
    ID_COLOR = {1: '#EE4D4D', 2: '#C67C7B', 3: '#FFD274', 4: '#BEBEBE', 5: '#BFE3E8',
                6: '#7BA779', 7: '#E87A90', 8: '#FF8C69', 10: '#1F849B', 11: '#727171',
                13: '#785A67', 12: '#D3A2C7'}
    num_room_types = 14
    sample_fn = (diffusion.p_sample_loop if not use_ddim else diffusion.ddim_sample_loop)
    print(graphs, corners, room_type)
    model_kwargs = function_test(graphs, corners, room_type)
    for key in model_kwargs:
        # model_kwargs[key] = th.from_numpy(np.array([model_kwargs[key]])).cuda()
        model_kwargs[key] = th.from_numpy(np.array([model_kwargs[key]])).cpu()

    png_paths = []
    svg_paths = []
    for count in range(num_samples):
        sample = sample_fn(
            model,
            th.Size([1, 2, 100]),
            clip_denoised=True,
            model_kwargs=model_kwargs,
        )

        sample = sample.permute([0, 1, 3, 2])

        pred = save_samples(sample, 'pred', model_kwargs, count, num_room_types, ID_COLOR=ID_COLOR,
                            is_syn=True, draw_graph=False, save_svg=True, save_gif=False, metrics=metrics)

        temp_svg_file = tempfile.NamedTemporaryFile(delete=False, suffix=".svg")
        pred.saveSvg(temp_svg_file.name)
        png_file_name = temp_svg_file.name.split(".")[0].split("/")[-1]
        png_file_path = f'./generated_svgs/{png_file_name}.png'
        # print(temp_svg_file.name)
        # print(png_file_name)
        # print(png_file_path)

        Image.open(io.BytesIO(cairosvg.svg2png(pred.asSvg()))).save(png_file_path)

        output_dir = Path("./generated_svgs")
        output_dir.mkdir(parents=True, exist_ok=True)
        file_name = temp_svg_file.name.split("/")[-1]
        persistent_path = Path(f"{output_dir}/{file_name}")
        shutil.move(temp_svg_file.name, persistent_path)
        os.chmod(persistent_path, 0o644)

        svg_paths.append(str(persistent_path))
        png_paths.append(png_file_path)
        # print(str(persistent_path))

    return png_paths, svg_paths


rooms_data = []
edges_data = []


def generate_layout(metrics: bool, ddim_steps: int, num_samples: int):
    room_list = []
    room_corners = []
    living_room = 0
    front_door = False
    entrance = -1

    print(rooms_data)
    print(edges_data)

    for i, room in enumerate(rooms_data):
        room_list.append(ROOM_CLASS[room['room_type']])
        if room['num_corners'] != 0:
            room_corners.append(int(room['num_corners']))
        else:
            room_corners.append(4)

        if room['room_type'] == "Living Room":
            living_room = i

        elif room['room_type'] == "Entrance":
            entrance = i

        elif room['room_type'] == "Front Door":
            front_door = True

    edges = []
    for edge in edges_data:
        source_id = int(edge['room1_id'].split()[0])
        target_id = int(edge['room2_id'].split()[0])
        edges.append([source_id, 1, target_id])

        index = len(room_list)
        room_list.append(12)
        room_corners.append(4)
        edges.append([source_id, 1, index])
        edges.append([target_id, 1, index])

    if not front_door:
        room_list.append(11)
        room_corners.append(4)
        if entrance == -1:
            edges.append([len(room_list) - 1, 1, living_room])
        else:
            edges.append([len(room_list) - 1, 1, entrance])

    if np.sum(room_corners) > 99:
        return {"Error": "Number of Corners exceeded"}

    print(room_list, room_corners, edges)
    png_paths, svg_paths = create_layout(edges, room_corners, room_list, metrics=metrics, ddim_steps=ddim_steps,
                             num_samples=num_samples)

    png_color_guide = './color_guide.png'

    rooms_data.clear()
    edges_data.clear()

    rooms_df = pd.DataFrame(columns=["room_id", "room_type", "num_corners"])
    edges_df = pd.DataFrame(columns=["edge_id", "room1_id", "room2_id"])

    return png_paths, svg_paths, png_color_guide, rooms_df, edges_df


with gr.Blocks() as demo:
    gr.Markdown("## House Layout Generator")

    with gr.Row():
        room_type = gr.Dropdown(label="Room Type", choices=list(ROOM_CATEGORIES.keys()), value="Living Room")
        num_corners = gr.Number(label="Number of Corners", value=4)
        add_room_button = gr.Button("Add Room")

    with gr.Row():
        room1_id = gr.Dropdown(label="Room 1", choices=[], value=None)
        room2_id = gr.Dropdown(label="Room 2", choices=[], value=None)
        add_edge_button = gr.Button("Add Edge")

    rooms_table = gr.DataFrame(label="Rooms Table")
    edges_table = gr.DataFrame(label="Edges Table")

    metrics_toggle = gr.Checkbox(label="Include metrics", value=True)
    ddim_input = gr.Number(label="DDIM steps", value=100)
    num_sample = gr.Number(label="Number of samples", value=4)

    png_gallery = gr.Gallery(label="Layout PNG Outputs", columns=4)
    svg_files = gr.File(label="Layout SVG Outputs (higher quality)")
    png_color_guide = gr.Image(label="Color Guide")

    def add_room(room_type, num_corners):
        room_id = len(rooms_data)
        rooms_data.append({
            "room_id": room_id,
            "room_type": room_type,
            "num_corners": num_corners
        })
        return update_rooms_and_edges()


    def add_edge(room1_id, room2_id):
        edge_id = len(edges_data)
        edges_data.append({
            "edge_id": edge_id,
            "room1_id": room1_id,
            "room2_id": room2_id
        })
        return update_rooms_and_edges()


    def update_rooms_and_edges():
        rooms_df = pd.DataFrame(rooms_data, columns=["room_id", "room_type", "num_corners"])
        edges_df = pd.DataFrame(edges_data, columns=["edge_id", "room1_id", "room2_id"])
        room_options = [f"{room['room_id']} {room['room_type']}" for room in rooms_data]
        return rooms_df, edges_df, gr.update(choices=room_options, value=None), gr.update(choices=room_options,
                                                                                          value=None)


    generate_button = gr.Button("Generate Layout")
    # generate_button.click(generate_layout, inputs=[metrics_toggle, ddim_input, num_sample], outputs=[png_gallery, svg_files, png_color_guide])
    generate_button.click(
        generate_layout,
        inputs=[metrics_toggle, ddim_input, num_sample],
        outputs=[
            png_gallery,
            svg_files,
            png_color_guide,
            rooms_table,
            edges_table
        ]
    )
    add_room_button.click(add_room, inputs=[room_type, num_corners],
                          outputs=[rooms_table, edges_table, room1_id, room2_id])
    add_edge_button.click(add_edge, inputs=[room1_id, room2_id], outputs=[rooms_table, edges_table, room1_id, room2_id])


demo.launch()
# global demo
# demo.launch(share=True)