house_diffusion/rplanhg_datasets.py

import math
import random
import torch as th

from PIL import Image, ImageDraw
import blobfile as bf
from mpi4py import MPI
import numpy as np
from torch.utils.data import DataLoader, Dataset
from glob import glob
import json
import os
import cv2 as cv
from tqdm import tqdm
from shapely import geometry as gm
from shapely.ops import unary_union
from collections import defaultdict
import copy


def load_rplanhg_data(
        batch_size,
        analog_bit,
        target_set,
        set_name='train',
):
    """
    For a dataset, create a generator over (shapes, kwargs) pairs.
    """
    # set_name = 'train'
    set_name = 'eval'
    print(f"loading {set_name} of target set {target_set}")
    deterministic = False if set_name == 'train' else True
    dataset = RPlanhgDataset(set_name, analog_bit, target_set)
    if deterministic:
        loader = DataLoader(
            dataset, batch_size=batch_size, shuffle=False, num_workers=2, drop_last=False
        )
    else:
        loader = DataLoader(
            dataset, batch_size=batch_size, shuffle=True, num_workers=2, drop_last=False
        )
    while True:
        yield from loader


def make_non_manhattan(poly, polygon, house_poly):
    dist = abs(poly[2] - poly[0])
    direction = np.argmin(dist)
    center = poly.mean(0)
    min = poly.min(0)
    max = poly.max(0)

    tmp = np.random.randint(3, 7)
    new_min_y = center[1] - (max[1] - min[1]) / tmp
    new_max_y = center[1] + (max[1] - min[1]) / tmp
    if center[0] < 128:
        new_min_x = min[0] - (max[0] - min[0]) / np.random.randint(2, 5)
        new_max_x = center[0]
        poly1 = [[min[0], min[1]], [new_min_x, new_min_y], [new_min_x, new_max_y], [min[0], max[1]], [max[0], max[1]],
                 [max[0], min[1]]]
    else:
        new_min_x = center[0]
        new_max_x = max[0] + (max[0] - min[0]) / np.random.randint(2, 5)
        poly1 = [[min[0], min[1]], [min[0], max[1]], [max[0], max[1]], [new_max_x, new_max_y], [new_max_x, new_min_y],
                 [max[0], min[1]]]

    new_min_x = center[0] - (max[0] - min[0]) / tmp
    new_max_x = center[0] + (max[0] - min[0]) / tmp
    if center[1] < 128:
        new_min_y = min[1] - (max[1] - min[1]) / np.random.randint(2, 5)
        new_max_y = center[1]
        poly2 = [[min[0], min[1]], [min[0], max[1]], [max[0], max[1]], [max[0], min[1]], [new_max_x, new_min_y],
                 [new_min_x, new_min_y]]
    else:
        new_min_y = center[1]
        new_max_y = max[1] + (max[1] - min[1]) / np.random.randint(2, 5)
        poly2 = [[min[0], min[1]], [min[0], max[1]], [new_min_x, new_max_y], [new_max_x, new_max_y], [max[0], max[1]],
                 [max[0], min[1]]]
    p1 = gm.Polygon(poly1)
    iou1 = house_poly.intersection(p1).area / p1.area
    p2 = gm.Polygon(poly2)
    iou2 = house_poly.intersection(p2).area / p2.area
    if iou1 > 0.9 and iou2 > 0.9:
        return poly
    if iou1 < iou2:
        return poly1
    else:
        return poly2


get_bin = lambda x, z: [int(y) for y in format(x, 'b').zfill(z)]
get_one_hot = lambda x, z: np.eye(z)[min(x, z - 1)]


class RPlanhgDataset(Dataset):
    def __init__(self, set_name, analog_bit, target_set, non_manhattan=False):
        super().__init__()
        base_dir = '../datasets/rplan'
        self.non_manhattan = non_manhattan
        self.set_name = set_name
        self.analog_bit = analog_bit
        self.target_set = target_set
        self.subgraphs = []
        self.org_graphs = []
        self.org_houses = []
        max_num_points = 100
        if self.set_name == 'eval':
            cnumber_dist = np.load(f'processed_rplan/rplan_train_{target_set}_cndist.npz', allow_pickle=True)[
                'cnumber_dist'].item()
        if os.path.exists(f'processed_rplan/rplan_{set_name}_{target_set}.npz'):
            data = np.load(f'processed_rplan/rplan_{set_name}_{target_set}.npz', allow_pickle=True)
            self.graphs = data['graphs']
            self.houses = data['houses']
            self.door_masks = data['door_masks']
            self.self_masks = data['self_masks']
            self.gen_masks = data['gen_masks']
            self.num_coords = 2
            self.max_num_points = max_num_points
            cnumber_dist = np.load(f'processed_rplan/rplan_train_{target_set}_cndist.npz', allow_pickle=True)[
                'cnumber_dist'].item()
            if self.set_name == 'eval':
                data = np.load(f'processed_rplan/rplan_{set_name}_{target_set}_syn.npz', allow_pickle=True)
                self.syn_graphs = data['graphs']
                self.syn_houses = data['houses']
                self.syn_door_masks = data['door_masks']
                self.syn_self_masks = data['self_masks']
                self.syn_gen_masks = data['gen_masks']
        else:
            with open(f'{base_dir}/list.txt') as f:
                lines = f.readlines()
            cnt = 0

            # TODO
            failed_plans = []

            for line in tqdm(lines):
                #     cnt=cnt+1
                #     file_name = f'{base_dir}/{line[:-1]}'
                #     rms_type, fp_eds,rms_bbs,eds_to_rms=reader(file_name)
                #     fp_size = len([x for x in rms_type if x != 15 and x != 17])
                #     if self.set_name=='train' and fp_size == target_set:
                #             continue
                #     if self.set_name=='eval' and fp_size != target_set:
                #             continue
                #     a = [rms_type, rms_bbs, fp_eds, eds_to_rms]
                #     self.subgraphs.append(a)

                # for graph in tqdm(self.subgraphs):
                try:
                    cnt = cnt + 1
                    file_name = f'{base_dir}/{line[:-1]}'
                    rms_type, fp_eds, rms_bbs, eds_to_rms = reader(file_name)
                    fp_size = len([x for x in rms_type if x != 15 and x != 17])
                    if self.set_name == 'train' and fp_size == target_set:
                        continue
                    if self.set_name == 'eval' and fp_size != target_set:
                        continue
                    graph = [rms_type, rms_bbs, fp_eds, eds_to_rms]
                    rms_type = graph[0]
                    rms_bbs = graph[1]
                    fp_eds = graph[2]
                    eds_to_rms = graph[3]
                    rms_bbs = np.array(rms_bbs)
                    fp_eds = np.array(fp_eds)

                    # extract boundary box and centralize
                    tl = np.min(rms_bbs[:, :2], 0)
                    br = np.max(rms_bbs[:, 2:], 0)
                    shift = (tl + br) / 2.0 - 0.5
                    rms_bbs[:, :2] -= shift
                    rms_bbs[:, 2:] -= shift
                    fp_eds[:, :2] -= shift
                    fp_eds[:, 2:] -= shift
                    tl -= shift
                    br -= shift

                    # build input graph
                    graph_nodes, graph_edges, rooms_mks = self.build_graph(rms_type, fp_eds, eds_to_rms)

                    house = []
                    for room_mask, room_type in zip(rooms_mks, graph_nodes):
                        room_mask = room_mask.astype(np.uint8)
                        room_mask = cv.resize(room_mask, (256, 256), interpolation=cv.INTER_AREA)
                        contours, _ = cv.findContours(room_mask, cv.RETR_TREE, cv.CHAIN_APPROX_SIMPLE)
                        contours = contours[0]
                        house.append([contours[:, 0, :], room_type])
                    self.org_graphs.append(graph_edges)
                    self.org_houses.append(house)

                except IndexError:
                    # print(line)
                    failed_plans.append(line)
            print("failed: ", failed_plans)
            print("len: ", len(failed_plans))

            houses = []
            door_masks = []
            self_masks = []
            gen_masks = []
            graphs = []
            if self.set_name == 'train':
                cnumber_dist = defaultdict(list)

            if self.non_manhattan:
                for h, graph in tqdm(zip(self.org_houses, self.org_graphs), desc='processing dataset'):
                    # Generating non-manhattan Balconies
                    tmp = []
                    for i, room in enumerate(h):
                        if room[1] > 10:
                            continue
                        if len(room[0]) != 4:
                            continue
                        if np.random.randint(2):
                            continue
                        poly = gm.Polygon(room[0])
                        house_polygon = unary_union([gm.Polygon(room[0]) for room in h])
                        room[0] = make_non_manhattan(room[0], poly, house_polygon)

            for h, graph in tqdm(zip(self.org_houses, self.org_graphs), desc='processing dataset'):
                house = []
                corner_bounds = []
                num_points = 0
                for i, room in enumerate(h):
                    if room[1] > 10:
                        room[1] = {15: 11, 17: 12, 16: 13}[room[1]]
                    room[0] = np.reshape(room[0], [len(room[0]),
                                                   2]) / 256. - 0.5  # [[x0,y0],[x1,y1],...,[x15,y15]] and map to 0-1 - > -0.5, 0.5
                    room[0] = room[0] * 2  # map to [-1, 1]
                    if self.set_name == 'train':
                        cnumber_dist[room[1]].append(len(room[0]))
                    # Adding conditions
                    num_room_corners = len(room[0])
                    rtype = np.repeat(np.array([get_one_hot(room[1], 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((room[0], rtype, corner_index, room_index, padding_mask, connections), 1)
                    house.append(room)

                house_layouts = np.concatenate(house, 0)
                if len(house_layouts) > max_num_points:
                    continue
                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 in range(len(corner_bounds)):
                    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], graph).all(1)) or any(np.equal([j, 1, i], graph).all(1)):
                            door_mask[corner_bounds[i][0]:corner_bounds[i][1],
                            corner_bounds[j][0]:corner_bounds[j][1]] = 0
                houses.append(house_layouts)
                door_masks.append(door_mask)
                self_masks.append(self_mask)
                gen_masks.append(gen_mask)
                graphs.append(graph)
            self.max_num_points = max_num_points
            self.houses = houses
            self.door_masks = door_masks
            self.self_masks = self_masks
            self.gen_masks = gen_masks
            self.num_coords = 2
            self.graphs = graphs

            # --------------
            # graph_dict = {f'graph_{i}': graph for i, graph in enumerate(self.graphs)}
            for i, graph in enumerate(self.graphs):
                print(f"Graph {i}: shape = {np.shape(graph)}, type = {type(graph)}")

            # Save each graph individually within a dictionary
            # graph_dict = {f'graph_{i}': graph for i, graph in enumerate(self.graphs)}

            np.savez_compressed(f'processed_rplan/rplan_{set_name}_{target_set}', graphs=self.graphs,
                                houses=self.houses,
                                # np.savez_compressed(f'processed_rplan/rplan_{set_name}_{target_set}', **graph_dict, houses=self.houses,
                                door_masks=self.door_masks, self_masks=self.self_masks, gen_masks=self.gen_masks)
            if self.set_name == 'train':
                np.savez_compressed(f'processed_rplan/rplan_{set_name}_{target_set}_cndist', cnumber_dist=cnumber_dist)

            if set_name == 'eval':
                houses = []
                graphs = []
                door_masks = []
                self_masks = []
                gen_masks = []
                len_house_layouts = 0
                for h, graph in tqdm(zip(self.org_houses, self.org_graphs), desc='processing dataset'):
                    house = []
                    corner_bounds = []
                    num_points = 0
                    # num_room_corners_total = [cnumber_dist[room[1]][random.randint(0, len(cnumber_dist[room[1]])-1)] for room in h]
                    # while np.sum(num_room_corners_total)>=max_num_points:
                    #     num_room_corners_total = [cnumber_dist[room[1]][random.randint(0, len(cnumber_dist[room[1]])-1)] for room in h]
                    num_room_corners_total = []
                    for room in h:
                        room_type = room[1]
                        default_value = 4
                        if room_type in cnumber_dist and cnumber_dist[room_type]:
                            num_room_corners_total.append(
                                cnumber_dist[room_type][random.randint(0, len(cnumber_dist[room_type]) - 1)]
                            )
                        else:
                            # Handle the case where cnumber_dist[room_type] is missing or empty
                            print(f"Warning: No data found for room type {room_type}. Assigning default value.")
                            default_value = 4  # Assign a reasonable default value or handle accordingly
                            num_room_corners_total.append(default_value)

                    while np.sum(num_room_corners_total) >= max_num_points:
                        num_room_corners_total = []
                        for room in h:
                            room_type = room[1]
                            if room_type in cnumber_dist and cnumber_dist[room_type]:
                                num_room_corners_total.append(
                                    cnumber_dist[room_type][random.randint(0, len(cnumber_dist[room_type]) - 1)]
                                )
                        else:
                            num_room_corners_total.append(default_value)

                    for i, room in enumerate(h):
                        # Adding conditions
                        num_room_corners = num_room_corners_total[i]
                        rtype = np.repeat(np.array([get_one_hot(room[1], 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)
                    if np.sum([len(room[0]) for room in h]) > max_num_points:
                        continue
                    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(h):
                        if room[1] == 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], graph).all(1)) or any(np.equal([j, 1, i], graph).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

                    houses.append(house_layouts)
                    door_masks.append(door_mask)
                    self_masks.append(self_mask)
                    gen_masks.append(gen_mask)
                    graphs.append(graph)
                self.syn_houses = houses
                self.syn_door_masks = door_masks
                self.syn_self_masks = self_masks
                self.syn_gen_masks = gen_masks
                self.syn_graphs = graphs
                np.savez_compressed(f'processed_rplan/rplan_{set_name}_{target_set}_syn', graphs=self.syn_graphs,
                                    houses=self.syn_houses,
                                    door_masks=self.syn_door_masks, self_masks=self.syn_self_masks,
                                    gen_masks=self.syn_gen_masks)


    def __len__(self):
        return len(self.houses)

    def __getitem__(self, idx):
        # idx = int(idx//20)
        arr = self.houses[idx][:, :self.num_coords]
        graph = np.concatenate((self.graphs[idx], np.zeros([200 - len(self.graphs[idx]), 3])), 0)

        cond = {
            'door_mask': self.door_masks[idx],
            'self_mask': self.self_masks[idx],
            'gen_mask': self.gen_masks[idx],
            'room_types': self.houses[idx][:, self.num_coords:self.num_coords + 25],
            'corner_indices': self.houses[idx][:, self.num_coords + 25:self.num_coords + 57],
            'room_indices': self.houses[idx][:, self.num_coords + 57:self.num_coords + 89],
            'src_key_padding_mask': 1 - self.houses[idx][:, self.num_coords + 89],
            'connections': self.houses[idx][:, self.num_coords + 90:self.num_coords + 92],
            'graph': graph,
        }
        if self.set_name == 'eval':
            syn_graph = np.concatenate((self.syn_graphs[idx], np.zeros([200 - len(self.syn_graphs[idx]), 3])), 0)
            assert (graph == syn_graph).all(), idx
            cond.update({
                'syn_door_mask': self.syn_door_masks[idx],
                'syn_self_mask': self.syn_self_masks[idx],
                'syn_gen_mask': self.syn_gen_masks[idx],
                'syn_room_types': self.syn_houses[idx][:, self.num_coords:self.num_coords + 25],
                'syn_corner_indices': self.syn_houses[idx][:, self.num_coords + 25:self.num_coords + 57],
                'syn_room_indices': self.syn_houses[idx][:, self.num_coords + 57:self.num_coords + 89],
                'syn_src_key_padding_mask': 1 - self.syn_houses[idx][:, self.num_coords + 89],
                'syn_connections': self.syn_houses[idx][:, self.num_coords + 90:self.num_coords + 92],
                'syn_graph': syn_graph,
            })
        if self.set_name == 'train':
            #### Random Rotate
            rotation = random.randint(0, 3)
            if rotation == 1:
                arr[:, [0, 1]] = arr[:, [1, 0]]
                arr[:, 0] = -arr[:, 0]
            elif rotation == 2:
                arr[:, [0, 1]] = -arr[:, [1, 0]]
            elif rotation == 3:
                arr[:, [0, 1]] = arr[:, [1, 0]]
                arr[:, 1] = -arr[:, 1]

            ## To generate any rotation uncomment this

            # if self.non_manhattan:
            # theta = random.random()*np.pi/2
            # rot_mat = np.array([[np.cos(theta), -np.sin(theta), 0],
            # [np.sin(theta), np.cos(theta), 0]])
            # arr = np.matmul(arr,rot_mat)[:,:2]

            # Random Scale
            # arr = arr * np.random.normal(1., .5)

            # Random Shift
            # arr[:, 0] = arr[:, 0] + np.random.normal(0., .1)
            # arr[:, 1] = arr[:, 1] + np.random.normal(0., .1)

        if not self.analog_bit:
            arr = np.transpose(arr, [1, 0])
            return arr.astype(float), cond
        else:
            ONE_HOT_RES = 256
            arr_onehot = np.zeros((ONE_HOT_RES * 2, arr.shape[1])) - 1
            xs = ((arr[:, 0] + 1) * (ONE_HOT_RES / 2)).astype(int)
            ys = ((arr[:, 1] + 1) * (ONE_HOT_RES / 2)).astype(int)
            xs = np.array([get_bin(x, 8) for x in xs])
            ys = np.array([get_bin(x, 8) for x in ys])
            arr_onehot = np.concatenate([xs, ys], 1)
            arr_onehot = np.transpose(arr_onehot, [1, 0])
            arr_onehot[arr_onehot == 0] = -1
            return arr_onehot.astype(float), cond

    def make_sequence(self, edges):
        polys = []
        v_curr = tuple(edges[0][:2])
        e_ind_curr = 0
        e_visited = [0]
        seq_tracker = [v_curr]
        find_next = False
        while len(e_visited) < len(edges):
            if find_next == False:
                if v_curr == tuple(edges[e_ind_curr][2:]):
                    v_curr = tuple(edges[e_ind_curr][:2])
                else:
                    v_curr = tuple(edges[e_ind_curr][2:])
                find_next = not find_next
            else:
                # look for next edge
                for k, e in enumerate(edges):
                    if k not in e_visited:
                        if (v_curr == tuple(e[:2])):
                            v_curr = tuple(e[2:])
                            e_ind_curr = k
                            e_visited.append(k)
                            break
                        elif (v_curr == tuple(e[2:])):
                            v_curr = tuple(e[:2])
                            e_ind_curr = k
                            e_visited.append(k)
                            break

            # extract next sequence
            if v_curr == seq_tracker[-1]:
                polys.append(seq_tracker)
                for k, e in enumerate(edges):
                    if k not in e_visited:
                        v_curr = tuple(edges[0][:2])
                        seq_tracker = [v_curr]
                        find_next = False
                        e_ind_curr = k
                        e_visited.append(k)
                        break
            else:
                seq_tracker.append(v_curr)
        polys.append(seq_tracker)

        return polys

    def build_graph(self, rms_type, fp_eds, eds_to_rms, out_size=64):
        # create edges
        triples = []
        nodes = rms_type
        # encode connections
        for k in range(len(nodes)):
            for l in range(len(nodes)):
                if l > k:
                    is_adjacent = any([True for e_map in eds_to_rms if (l in e_map) and (k in e_map)])
                    if is_adjacent:
                        if 'train' in self.set_name:
                            triples.append([k, 1, l])
                        else:
                            triples.append([k, 1, l])
                    else:
                        if 'train' in self.set_name:
                            triples.append([k, -1, l])
                        else:
                            triples.append([k, -1, l])
        # get rooms masks
        eds_to_rms_tmp = []
        for l in range(len(eds_to_rms)):
            eds_to_rms_tmp.append([eds_to_rms[l][0]])
        rms_masks = []
        im_size = 256
        fp_mk = np.zeros((out_size, out_size))
        for k in range(len(nodes)):
            # add rooms and doors
            eds = []
            for l, e_map in enumerate(eds_to_rms_tmp):
                if (k in e_map):
                    eds.append(l)
            # draw rooms
            rm_im = Image.new('L', (im_size, im_size))
            dr = ImageDraw.Draw(rm_im)
            for eds_poly in [eds]:
                poly = self.make_sequence(np.array([fp_eds[l][:4] for l in eds_poly]))[0]
                poly = [(im_size * x, im_size * y) for x, y in poly]
                if len(poly) >= 2:
                    dr.polygon(poly, fill='white')
                else:
                    print("Empty room")
                    exit(0)
            rm_im = rm_im.resize((out_size, out_size))
            rm_arr = np.array(rm_im)
            inds = np.where(rm_arr > 0)
            rm_arr[inds] = 1.0
            rms_masks.append(rm_arr)
            if rms_type[k] != 15 and rms_type[k] != 17:
                fp_mk[inds] = k + 1
        # trick to remove overlap
        for k in range(len(nodes)):
            if rms_type[k] != 15 and rms_type[k] != 17:
                rm_arr = np.zeros((out_size, out_size))
                inds = np.where(fp_mk == k + 1)
                rm_arr[inds] = 1.0
                rms_masks[k] = rm_arr
        # convert to array
        nodes = np.array(nodes)
        triples = np.array(triples)
        rms_masks = np.array(rms_masks)
        return nodes, triples, rms_masks


def is_adjacent(box_a, box_b, threshold=0.03):
    x0, y0, x1, y1 = box_a
    x2, y2, x3, y3 = box_b
    h1, h2 = x1 - x0, x3 - x2
    w1, w2 = y1 - y0, y3 - y2
    xc1, xc2 = (x0 + x1) / 2.0, (x2 + x3) / 2.0
    yc1, yc2 = (y0 + y1) / 2.0, (y2 + y3) / 2.0
    delta_x = np.abs(xc2 - xc1) - (h1 + h2) / 2.0
    delta_y = np.abs(yc2 - yc1) - (w1 + w2) / 2.0
    delta = max(delta_x, delta_y)
    return delta < threshold


def reader(filename):
    with open(filename) as f:
        info = json.load(f)
        rms_bbs = np.asarray(info['boxes'])
        fp_eds = info['edges']
        rms_type = info['room_type']
        eds_to_rms = info['ed_rm']
        s_r = 0
        for rmk in range(len(rms_type)):
            if (rms_type[rmk] != 17):
                s_r = s_r + 1
        rms_bbs = np.array(rms_bbs) / 256.0
        fp_eds = np.array(fp_eds) / 256.0
        fp_eds = fp_eds[:, :4]
        tl = np.min(rms_bbs[:, :2], 0)
        br = np.max(rms_bbs[:, 2:], 0)
        shift = (tl + br) / 2.0 - 0.5
        rms_bbs[:, :2] -= shift
        rms_bbs[:, 2:] -= shift
        fp_eds[:, :2] -= shift
        fp_eds[:, 2:] -= shift
        tl -= shift
        br -= shift
        return rms_type, fp_eds, rms_bbs, eds_to_rms


if __name__ == '__main__':
    dataset = RPlanhgDataset('eval', False, 8)