Spaces:
Running on Zero
Running on Zero
File size: 4,590 Bytes
331600c 83d5461 331600c 83d5461 331600c 83d5461 46b4a9e 83d5461 46b4a9e 83d5461 46b4a9e 83d5461 46b4a9e 83d5461 46b4a9e 83d5461 46b4a9e 83d5461 46b4a9e 83d5461 46b4a9e 83d5461 f424e40 83d5461 46b4a9e 83d5461 46b4a9e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 | import glob
import logging
import os
import pickle
import ssl
import subprocess
import sys
import urllib.request
logging.basicConfig(format="[%(levelname)s] %(asctime)s %(message)s", level=logging.INFO)
# Project-local CUDA extension wheels (diff_gaussian_rasterization,
# grid_encoder, voxlib_ext) live in ./wheels/ and are installed at app
# startup because the HF Spaces build phase only mounts requirements.txt,
# not the wheels directory, so they cannot be referenced from there.
_wheel_dir = os.path.join(os.path.dirname(__file__), "wheels")
for _whl in sorted(glob.glob(os.path.join(_wheel_dir, "*.whl"))):
logging.info("Installing project-local wheel: %s", os.path.basename(_whl))
subprocess.check_call([sys.executable, "-m", "pip", "install", "--no-deps", _whl])
import gradio as gr # noqa: E402
import numpy as np # noqa: E402
import spaces # noqa: E402
import torch # noqa: E402
from PIL import Image # noqa: E402
ssl._create_default_https_context = ssl._create_unverified_context
sys.path.append(os.path.join(os.path.dirname(__file__), "gaussiancity"))
def get_models(file_name):
import gaussiancity.generator
if not os.path.exists(file_name):
urllib.request.urlretrieve(
f"https://huggingface.co/hzxie/gaussian-city/resolve/main/{file_name}",
file_name,
)
ckpt = torch.load(file_name, map_location="cpu", weights_only=False)
model = gaussiancity.generator.Generator(
ckpt["cfg"].NETWORK.GAUSSIAN,
n_classes=ckpt["cfg"].DATASETS.GOOGLE_EARTH.N_CLASSES,
proj_size=ckpt["cfg"].DATASETS.GOOGLE_EARTH.PROJ_SIZE,
)
model = torch.nn.DataParallel(model).cuda().eval()
model.load_state_dict(ckpt["gaussian_g"], strict=False)
return model
def get_city_layout():
import gaussiancity.inference
if os.path.exists("assets/NYC.pkl"):
with open("assets/NYC.pkl", "rb") as fp:
layout = pickle.load(fp)
else:
td_hf = np.array(Image.open("assets/NYC-HghtFld.png")).astype(np.int32)
# Fix: nonzero is not supported for tensors with more than INT_MAX elements
td_hf[td_hf > 500] = 500
bu_hf = np.zeros_like(td_hf)
seg_map = np.array(Image.open("assets/NYC-SegMap.png").convert("P")).astype(np.int32)
ins_map = gaussiancity.inference.get_instance_seg_map(seg_map.copy())
pts_map = gaussiancity.inference.get_point_map(seg_map)
layout = {
"TD_HF": td_hf,
"BU_HF": bu_hf,
"SEG": seg_map,
"INS": ins_map,
"PTS": pts_map,
}
with open("assets/NYC.pkl", "wb") as fp:
pickle.dump(layout, fp)
if os.path.exists("assets/CENTERS.pkl"):
with open("assets/CENTERS.pkl", "rb") as fp:
centers = pickle.load(fp)
else:
centers = gaussiancity.inference.get_centers(layout["INS"], layout["TD_HF"])
with open("assets/CENTERS.pkl", "wb") as fp:
pickle.dump(centers, fp)
layout["CTR"] = centers
return layout
logging.basicConfig(format="[%(levelname)s] %(asctime)s %(message)s", level=logging.INFO)
logging.info("Loading pretrained models...")
fgm = get_models("GaussianCity-Fgnd.pth")
bgm = get_models("GaussianCity-Bgnd.pth")
logging.info("Loading New York city layout to RAM...")
city_layout = get_city_layout()
@spaces.GPU
def get_generated_city(radius, altitude, azimuth, map_center):
import gaussiancity.inference
return gaussiancity.inference.generate_city(
fgm.to("cuda"),
bgm.to("cuda"),
city_layout,
map_center,
map_center,
radius,
altitude,
azimuth,
)
def main():
title = "Generative Gaussian Splatting for Unbounded 3D City Generation"
with open("README.md", "r") as f:
markdown = f.read()
desc = markdown[markdown.rfind("---") + 3 :]
with open("ARTICLE.md", "r") as f:
arti = f.read()
app = gr.Interface(
get_generated_city,
[
gr.Slider(256, 960, value=768, step=4, label="Camera Radius (m)"),
gr.Slider(256, 960, value=768, step=4, label="Camera Altitude (m)"),
gr.Slider(0, 360, value=210, step=5, label="Camera Azimuth (°)"),
gr.Slider(1024, 7168, value=3570, step=4, label="Map Center (px)"),
],
[gr.Image(type="numpy", label="Generated City")],
title=title,
description=desc,
article=arti,
flagging_mode="never",
)
app.queue()
app.launch()
if __name__ == "__main__":
main()
|