s3d_floorplan_eval/S3DLoader/S3DLoader.py

import torch
from torch.utils.data import Dataset, DataLoader
import torch.utils.data.distributed
import os
import cv2
import json

from S3DLoader.s3d_utils import *
from S3DLoader.poly_utils import *


class S3DLoader(object):
    def __init__(self, args, mode, generate_input_candidates=False):
        self.mode = mode
        self.seed = 8978
        np.random.seed(seed=self.seed)

        if hasattr(args, 'network_mode'):
            self.function_mode = args.network_mode
        else:
            self.function_mode = "S"

        if hasattr(args, 'batch_size'):
            self.batch_size = args.batch_size
        else:
            self.batch_size = 1

        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        print('Selected device is:', device)
        self.device = device

        if mode == 'train':
            self.dataset = self.create_dataset(args, mode, generate_input_candidates)
            self.augment = True

            self.data = DataLoader(self.dataset, self.batch_size,
                                   drop_last=True,
                                   collate_fn=self.collate_fn,
                                   shuffle=True)

            self.sample_n = len(self.dataset)

        elif mode == 'online_eval' or mode == 'test':
            self.dataset = self.create_dataset(args, mode, generate_input_candidates)
            self.augment = False
            # self.batch_size = 4

            self.sample_n = len(self.dataset)

            self.data = DataLoader(self.dataset, self.batch_size,
                                   drop_last=True,
                                   collate_fn=self.collate_fn)


        elif mode == 'test':
            self.dataset = self.create_dataset(args, mode)
            self.augment = False
            self.batch_size = 1

            self.sample_n = 20

            self.data = DataLoader(self.dataset, self.batch_size,
                                   num_workers=1,
                                   drop_last=True,
                                   collate_fn=self.collate_fn)

        # elif mode == 'test':
        #     self.dataset = self.create_dataset(args, mode)
        #     self.augment = False
        #
        #     self.data = DataLoader(self.dataset,
        #                            1,
        #                            shuffle=False,
        #                            num_workers=1)
        #     self.sample_n = 20

        else:
            print('mode should be one of \'train, test, online_eval\'. Got {}'.format(mode))

    def collate_fn(self, samples):

        # wall_maps = [torch.tensor(s["wall_map"][None,:,:,None], device=self.device) for s in samples]
        room_maps = [torch.tensor(s["room_map"][None,:,:,None], device=self.device) for s in samples]
        input_maps = [torch.tensor(s["input_map"][None], device=self.device) for s in samples]
        scores = [torch.tensor(s["score"][None], device=self.device) for s in samples]

        torch_sample = {}
        torch_sample["room_map"] = torch.cat(room_maps, dim=0)
        # torch_sample["wall_map"] = torch.cat(wall_maps, dim=0)
        torch_sample["input_map"] = torch.cat(input_maps, dim=0)
        torch_sample["score"] = torch.cat(scores, dim=0)


        for key, value in torch_sample.items():
            assert torch.all(torch_sample[key] == torch_sample[key])
            assert torch.all(torch.logical_not(torch.isinf(torch_sample[key])))

        return torch_sample

    def create_dataset(self, args, mode, generate_input_candidates):
        #dataset_path = "../Structured3D/montefloor_data"
        self.args = args
        dataset_path = args.dataset_path

        if mode == "train":
            scenes_path = os.path.join(dataset_path, "train")

            dataset = S3DDataset(args, scenes_path, None,
                                 num_scenes=3000, generate_input_candidates=generate_input_candidates, mode=mode)

        elif mode == "online_eval":
            scenes_path = os.path.join(dataset_path, "val")

            dataset = S3DDataset(args, scenes_path, None,
                                 num_scenes=250, generate_input_candidates=generate_input_candidates, mode=mode)
        elif mode == "test":
            scenes_path = os.path.join(dataset_path, "test")
            # scenes_path = os.path.join(dataset_path, "val")

            dataset = S3DDataset(args, scenes_path, None,
                                 num_scenes=250, generate_input_candidates=generate_input_candidates, mode=mode)

        return dataset

    def load_sample(self, sample_batch):
        """
        Identity function. Everything is already loaded in Dataset class for Structured 3D
        :param sample_batch:
        :return:
        """
        return sample_batch


class S3DDataset(Dataset):
    def __init__(self, options, scenes_path, score_gen, num_scenes, generate_input_candidates, mode):
        print("Creating Structured3D Dataset with %d scenes..." % num_scenes)
        self.options = options
        self.score_gen = None

        self.mode = mode

        self.scenes_path = scenes_path
        self.floor_data_folder_name = ""

        self.scenes_list = os.listdir(scenes_path)
        self.scenes_list.sort()

        inv_scenes = ["scene_01155", "scene_01852", "scene_01192", "scene_01816"]
        self.scenes_list = [s for s in self.scenes_list if s not in inv_scenes]
        self.scenes_list = self.scenes_list[:num_scenes]


        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.device = device

        self.gen_input_candidates = generate_input_candidates

    def __getitem__(self, item):
        scene_name = self.scenes_list[item]
        sample = self.load_scene(scene_name)

        return sample

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

    def load_density_map(self, sp):
        """
        Load density map

        :param sp:
        :return:
        """
        density_path = os.path.join(sp, self.floor_data_folder_name, "density.png")
        density_map = cv2.imread(density_path, cv2.IMREAD_ANYCOLOR | cv2.IMREAD_ANYDEPTH) / 255.

        if self.gen_input_candidates:
            thresh = np.maximum(np.random.random(), 0.8)
            density_map = np.minimum(density_map, thresh) / thresh

        if self.mode != "test":
            pow = np.random.random()
            pow = (1.5 - 1.) * (pow - 1) + 1.5
            density_map = density_map ** pow


        return density_map.astype(np.float32)

    def load_annotation(self, sp):
        """
        Load annotation dict

        :param sp:
        :return:
        :rtype: dict
        """
        anno_path = os.path.join(sp, self.floor_data_folder_name, "annotation_3d.json")
        with open(anno_path, "r") as f:
            anno_dict = json.load(f)

        return anno_dict

    def load_scene(self, scene_name):
        """
        Load scene

        :param scene_name:
        :return:
        """

        def cvt_tmp_sample_to_torch():
            torch_sample = {}

            room_map = torch.tensor(np.array(sample['room_map']), device=self.device)[None]
            # room_map = kornia.morphology.dilation(room_map[:, None], kernel=torch.ones((5, 5), device=self.device))[:,0]

            torch_sample['room_map'] = room_map

            if 'input_map' in sample.keys():
                torch_sample['input_map'] = torch.tensor(np.array(sample['input_map']), device=self.device)[None]
                torch_sample['cand_inst'] = torch.tensor(np.array(sample['cand_inst']), device=self.device)[None]
                torch_sample['cand_confidence'] = torch.tensor(np.array(sample['cand_confidence']), device=self.device)[
                    None]

            else:
                torch_sample['density_map'] = torch.tensor(np.array(sample['density_map']), device=self.device)[None]
            torch_sample['wall_map'] = torch.tensor(np.array(sample['wall_map']), device=self.device)[None]
            # torch_sample['room_map'] = torch.tensor(np.array(sample['room_map']), device=self.device)[None]
            torch_sample['polygons_list'] = [torch.tensor(poly, device=self.device)[None] for poly in sample['polygons_list']]

            return torch_sample

        sp = os.path.join(self.scenes_path, scene_name)
        sample = {}
        sample["scene_name"] = scene_name

        scene_anno = self.load_annotation(sp)

        # density_map = torch.tensor(np.array(density_map))[None]
        density_map = self.load_density_map(sp)

        self.generate_room_map(sample, scene_anno, density_map)

        sample['density_map'] = density_map

        # import pdb; pdb.set_trace()
        for key, value in sample.items():
            assert np.all(value == value), "%s contains NaN" % key

        # import matplotlib.pyplot as plt
        # plt.figure()
        # plt.subplot(131)
        # plt.title(scene_name)
        # plt.imshow(density_map)
        # plt.subplot(132)
        # plt.imshow(sample["room_map"])
        # plt.subplot(133)
        # # plt.imshow(sample["input_map"][:,:,1])
        # # plt.imshow(sample["cand_inst"][:,:,0])
        # plt.show()

        return sample

    def generate_room_map(self, sample, annos, density_map):
        """

        :param density_map:
        :param sample:
        :param annos:
        :return:
        """

        h, w = density_map.shape

        polys = parse_floor_plan_polys(annos)

        room_map, polygons_list = generate_floorplan(annos, polys, h, w, ignore_types=['outwall', 'door', 'window'], constant_color=False, shuffle=self.gen_input_candidates)

        room_map = cv2.dilate(room_map, np.ones((5,5)))


        wall_map, _ = generate_floorplan(annos, polys, h, w, ignore_types=[], include_types=['outwall'], constant_color=True)
        wall_map *= (room_map == 0)

        sample['room_map'] = room_map.astype(np.float32)
        sample['wall_map'] = wall_map.astype(np.float32)

        sample['polygons_list'] = polygons_list

    def generate_density(self, points, width=256, height=256):
        image_res_tensor = torch.tensor([width, height], device=self.device).reshape(1, 1, 2)

        coordinates = torch.round(points[:, :, :2] * image_res_tensor)
        coordinates = torch.minimum(torch.maximum(coordinates, torch.zeros_like(image_res_tensor)),
                                    image_res_tensor - 1).type(torch.cuda.LongTensor)

        density = torch.zeros((self.batch_size, height, width), dtype=torch.float, device=self.device)

        for i in range(self.batch_size):
            unique_coordinates, counts = torch.unique(coordinates[i], return_counts=True, dim=0)

            density[i, unique_coordinates[:, 1], unique_coordinates[:, 0]] = counts.type(torch.cuda.FloatTensor)
            density[i] = density[i] / torch.max(density[i])

        return density