Spaces:
Runtime error
Runtime error
| import matplotlib.pyplot as plt | |
| import numpy as np | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| import csv | |
| import time | |
| import random | |
| import cv2 | |
| import os | |
| import torch | |
| from tqdm import tqdm | |
| from argparse import ArgumentParser | |
| if __name__ == '__main__': | |
| parser = ArgumentParser() | |
| parser.add_argument('--terrain', type=str, required=True, help='directory path to terrain dataset') | |
| parser.add_argument('--outdir', type=str, required=True) | |
| assert os.path.exists("./scenedreamer_released.pt") | |
| pcg_asset = torch.load("./scenedreamer_released.pt", map_location='cpu') | |
| args = parser.parse_args() | |
| terrain_dir = args.terrain | |
| outdir = args.outdir | |
| sample_height = 256 | |
| sample_size = 1024 | |
| os.makedirs(outdir, exist_ok=True) | |
| trees_models = pcg_asset['assets'] | |
| # can be customized | |
| biome_trees_dict = { | |
| 'desert': [], | |
| 'savanna': [5], | |
| 'twoodland': [1, 7], | |
| 'tundra': [], | |
| 'seasonal forest': [1, 2], | |
| 'rainforest': [1, 2, 3], | |
| 'temp forest': [4], | |
| 'temp rainforest': [0, 3], | |
| 'boreal': [5,6,7], | |
| 'water': [], | |
| } | |
| biome2mclabels = torch.tensor([28, 9, 8, 1, 9, 8, 9, 8, 30, 26], dtype=torch.int32) | |
| biome_names = list(biome_trees_dict.keys()) | |
| chunk_grid_x, chunk_grid_y = torch.meshgrid(torch.arange(sample_size), torch.arange(sample_size)) | |
| terrain_list = os.listdir(terrain_dir) | |
| for world in tqdm(terrain_list): | |
| voxel_t = torch.zeros(sample_height, sample_size, sample_size).to(torch.int32) | |
| current_dir = os.path.join(terrain_dir, world) | |
| height_map = np.load(os.path.join(current_dir, 'biome_rivers_height.npy')) | |
| height_map[height_map < 0] = 0 | |
| height_map = ((height_map - height_map.min()) / (1 - height_map.min()) * (sample_height - 1)).astype(np.int16) | |
| semantic_map = cv2.imread(os.path.join(current_dir, 'biome_rivers_labels.png'), 0) | |
| tree_map = cv2.imread(os.path.join(current_dir, 'biome_trees_dist.png'), 0) | |
| total_size = height_map.shape[0] | |
| crop_pos_x, crop_pos_y = np.random.randint(0, total_size - sample_size, size=2) | |
| org_height_map = height_map[crop_pos_x: crop_pos_x + sample_size, crop_pos_y: crop_pos_y + sample_size].astype(int) | |
| chunk_height_map = torch.from_numpy(org_height_map)[None, ...] | |
| chunk_semantic_map = semantic_map[crop_pos_x: crop_pos_x + sample_size, crop_pos_y: crop_pos_y + sample_size] | |
| chunk_trees_map = tree_map[crop_pos_x: crop_pos_x + sample_size, crop_pos_y: crop_pos_y + sample_size] | |
| org_semantic_map = torch.from_numpy(chunk_semantic_map.copy()) | |
| org_semantic_map[chunk_trees_map != 255] = 10 | |
| chunk_semantic_map = biome2mclabels[torch.from_numpy(chunk_semantic_map)[None, ...].long().contiguous()] | |
| voxel_t = voxel_t.scatter_(0, chunk_height_map, chunk_semantic_map) | |
| for preproc_step in range(8): | |
| voxel_t = voxel_t.scatter(0, torch.clip(chunk_height_map + preproc_step + 1, 0, sample_height - 1), chunk_semantic_map) | |
| chunk_height_map = chunk_height_map + 8 | |
| chunk_height_map = chunk_height_map[0] | |
| boundary_detect = 50 | |
| for biome_id in range(biome2mclabels.shape[0]): | |
| tree_pos_mask = (chunk_trees_map == biome_id) | |
| tree_pos_x = chunk_grid_x[tree_pos_mask] | |
| tree_pos_y = chunk_grid_y[tree_pos_mask] | |
| tree_pos_h = chunk_height_map[tree_pos_mask] | |
| assert len(tree_pos_x) == len(tree_pos_y) | |
| selected_trees = biome_trees_dict[biome_names[biome_id]] | |
| if len(selected_trees) == 0: | |
| continue | |
| for idx in range(len(tree_pos_x)): | |
| if tree_pos_x[idx] < boundary_detect or tree_pos_x[idx] > sample_size - boundary_detect or tree_pos_y[idx] < boundary_detect or tree_pos_y[idx] > sample_size - boundary_detect or tree_pos_h[idx] > sample_height - boundary_detect: | |
| # FIXME: hack, to avoid out of index near the boundary | |
| continue | |
| tree_id = random.choice(selected_trees) | |
| tmp = voxel_t[tree_pos_h[idx]: tree_pos_h[idx] + trees_models[tree_id].shape[0], tree_pos_x[idx]: tree_pos_x[idx] + trees_models[tree_id].shape[1], tree_pos_y[idx]: tree_pos_y[idx] + trees_models[tree_id].shape[2]] | |
| tmp_mask = (tmp == 0) | |
| try: | |
| voxel_t[tree_pos_h[idx]: tree_pos_h[idx] + trees_models[tree_id].shape[0], tree_pos_x[idx]: tree_pos_x[idx] + trees_models[tree_id].shape[1], tree_pos_y[idx]: tree_pos_y[idx] + trees_models[tree_id].shape[2]][tmp_mask] = trees_models[tree_id][tmp_mask] | |
| except: | |
| print('height?', tree_pos_h[idx]) | |
| print(tmp_mask.shape) | |
| print(tmp.shape) | |
| print(trees_models[tree_id].shape) | |
| print(voxel_t.shape) | |
| print(tree_id) | |
| raise NotImplementedError | |
| trans_mat = torch.eye(4) # Transform voxel to world | |
| # Generate heightmap for camera trajectory generation | |
| m, h = torch.max((torch.flip(voxel_t, [0]) != 0).int(), dim=0, keepdim=False) | |
| heightmap = voxel_t.shape[0] - 1 - h | |
| heightmap[m == 0] = 0 # Special case when the whole vertical column is empty | |
| voxel_t = voxel_t.numpy() | |
| voxel_value = voxel_t[voxel_t != 0] | |
| voxel_x, voxel_y, voxel_z = np.where(voxel_t != 0) | |
| current_height_map = (chunk_height_map / (sample_height - 1))[None, None, ...] | |
| current_semantic_map = F.one_hot(org_semantic_map.to(torch.int64)).to(torch.float).permute(2, 0, 1)[None, ...] | |
| semantic_map = torch.argmax(current_semantic_map, dim=1) | |
| print('semantic map after one hot and argmax', torch.unique(semantic_map, return_counts=True)) | |
| print(current_height_map.shape) | |
| print(current_semantic_map.shape) | |
| print(heightmap.shape) | |
| print(voxel_t.shape) | |
| print(voxel_value.shape) | |
| print(voxel_x.shape) | |
| print(voxel_y.shape) | |
| print(voxel_z.shape) | |
| print(voxel_z.dtype) | |
| voxel_sparse = np.stack([voxel_x, voxel_y, voxel_z, voxel_value]) | |
| print(voxel_sparse.shape) | |
| current_outdir = os.path.join(outdir, world) | |
| os.makedirs(current_outdir, exist_ok=True) | |
| np.save(os.path.join(current_outdir, 'voxel_sparse.npy'), voxel_sparse.astype(np.int16)) | |
| np.save(os.path.join(current_outdir, 'height_map.npy'), current_height_map.numpy()) | |
| np.save(os.path.join(current_outdir, 'semantic_map.npy'), current_semantic_map.numpy()) | |
| np.save(os.path.join(current_outdir, 'hmap_mc.npy'), heightmap.numpy()) | |