fix: the pretrained model weights are not loaded.

app.py

@@ -4,7 +4,7 @@
 # @Author: Haozhe Xie
 # @Date:   2024-03-02 16:30:00
 # @Last Modified by: Haozhe Xie
-# @Last Modified at: 2024-03-03 12:21:20
+# @Last Modified at: 2024-03-03 16:01:01
 # @Email:  root@haozhexie.com
 
 import gradio as gr
@@ -51,6 +51,7 @@ def get_models(file_name):
     if torch.cuda.is_available():
         model = torch.nn.DataParallel(model).cuda().eval()
 
+    model.load_state_dict(ckpt["gancraft_g"], strict=False)
     return model
 
 
@@ -60,15 +61,17 @@ def get_city_layout():
     return hf, seg
 
 
-def get_generated_city(radius, altitude, azimuth):
+def get_generated_city(radius, altitude, azimuth, map_center):
     # The import must be done after CUDA extension compilation
     import citydreamer.inference
 
     return citydreamer.inference.generate_city(
         get_generated_city.fgm,
         get_generated_city.bgm,
-        get_generated_city.hf,
-        get_generated_city.seg,
+        get_generated_city.hf.copy(),
+        get_generated_city.seg.copy(),
+        map_center,
+        map_center,
         radius,
         altitude,
         azimuth,
@@ -89,6 +92,7 @@ def main(debug):
             gr.Slider(128, 512, value=320, step=5, label="Camera Radius (m)"),
             gr.Slider(256, 512, value=384, step=5, label="Camera Altitude (m)"),
             gr.Slider(0, 360, value=180, step=5, label="Camera Azimuth (°)"),
+            gr.Slider(1440, 6752, value=2656, step=5, label="Map Center (px)"),
         ],
         [gr.Image(type="numpy", label="Generated City")],
         title=title,

citydreamer/inference.py

@@ -4,7 +4,7 @@
 # @Author: Haozhe Xie
 # @Date:   2024-03-02 16:30:00
 # @Last Modified by: Haozhe Xie
-# @Last Modified at: 2024-03-03 12:10:18
+# @Last Modified at: 2024-03-03 15:59:00
 # @Email:  root@haozhexie.com
 
 import copy
@@ -53,7 +53,7 @@ CONSTANTS = {
 }
 
 
-def generate_city(fgm, bgm, hf, seg, radius, altitude, azimuth):
+def generate_city(fgm, bgm, hf, seg, cx, cy, radius, altitude, azimuth):
     cam_pos = get_orbit_camera_position(radius, altitude, azimuth)
     seg, building_stats = get_instance_seg_map(seg)
     # Generate latent codes
@@ -63,15 +63,6 @@ def generate_city(fgm, bgm, hf, seg, radius, altitude, azimuth):
         bgm.module.cfg.NETWORK.GANCRAFT.STYLE_DIM,
         bgm.output_device,
     )
-    # Random choose the center of the patch
-    cy = (
-        np.random.randint(seg.shape[0] - CONSTANTS["EXTENDED_VOL_SIZE"])
-        + CONSTANTS["EXTENDED_VOL_SIZE"] // 2
-    )
-    cx = (
-        np.random.randint(seg.shape[1] - CONSTANTS["EXTENDED_VOL_SIZE"])
-        + CONSTANTS["EXTENDED_VOL_SIZE"] // 2
-    )
     # Generate local image patch of the height field and seg map
     part_hf, part_seg = get_part_hf_seg(hf, seg, cx, cy, CONSTANTS["EXTENDED_VOL_SIZE"])
     # Generate local image patch of the height field and seg map
@@ -98,9 +89,10 @@ def generate_city(fgm, bgm, hf, seg, radius, altitude, azimuth):
         bg_z,
         building_zs,
     )
-    return ((img.cpu().numpy().squeeze().transpose((1, 2, 0)) / 2 + 0.5) * 255).astype(
+    img = ((img.cpu().numpy().squeeze().transpose((1, 2, 0)) / 2 + 0.5) * 255).astype(
         np.uint8
     )
+    return img
 
 
 def get_orbit_camera_position(radius, altitude, azimuth):