Spaces:
Sleeping
Sleeping
| """GameTerrain — веб-демо генерации игрового рельефа (Hugging Face Spaces). | |
| Пользователь выбирает тип рельефа (flat/hilly/mountain), перепад высот и seed -> | |
| Conditional VAE генерирует heightmap -> скачивание в форматах для игр и 3D: | |
| - 16-bit PNG heightmap (Unity Terrain / Unreal Landscape принимают нативно) | |
| - .obj и .glb меш (Blender, игровые движки, 3D-редакторы) | |
| - превью-рендер PNG (matplotlib hillshade + палитра) | |
| - интерактивное 3D-превью прямо в браузере (Gradio Model3D) | |
| Опционально — гидравлическая эрозия (droplet) для реализма. | |
| Превью — matplotlib (лёгкое, работает на бесплатном HF CPU). Blender в демо НЕ | |
| тащим (тяжёлый, нет GPU). CVAE задаёт управляемую макроструктуру; эрозия и профиль | |
| высоты — слой реализма поверх. | |
| Запуск локально: python app.py | |
| Веса CVAE — из ./weights/cvae_best.pt или с HF Hub (env CVAE_MODEL_REPO). | |
| """ | |
| from __future__ import annotations | |
| import os | |
| import math | |
| import tempfile | |
| import numpy as np | |
| import gradio as gr | |
| import torch | |
| from PIL import Image | |
| # --- gradio 4.44 fix: schema-parser падает на additionalProperties=True (bool) --- | |
| # get_api_info() -> json_schema_to_python_type делает `"const" in schema`, а schema | |
| # иногда bool -> "TypeError: argument of type 'bool' is not iterable" на роуте /info, | |
| # из-за чего demo.launch() считает localhost недоступным и контейнер падает. | |
| # Делаем парсер терпимым к булевой схеме. | |
| import gradio_client.utils as _gcu # noqa: E402 | |
| if hasattr(_gcu, "_json_schema_to_python_type"): | |
| _gcu_orig_jstp = _gcu._json_schema_to_python_type | |
| def _gcu_safe_jstp(schema, defs=None): | |
| if isinstance(schema, bool): | |
| return "Any" | |
| return _gcu_orig_jstp(schema, defs) | |
| _gcu._json_schema_to_python_type = _gcu_safe_jstp | |
| if hasattr(_gcu, "get_type"): | |
| _gcu_orig_get_type = _gcu.get_type | |
| def _gcu_safe_get_type(schema): | |
| if not isinstance(schema, dict): | |
| return "Any" | |
| return _gcu_orig_get_type(schema) | |
| _gcu.get_type = _gcu_safe_get_type | |
| # --- /fix --- | |
| from src.models.cvae import ConvCVAE | |
| from src.postprocess.terrain_ops import hydraulic_erosion, normalize01 | |
| MODEL_REPO = os.environ.get("CVAE_MODEL_REPO", "") # напр. "username/gameterrain-cvae" | |
| LOCAL_WEIGHTS = "weights/cvae_best.pt" | |
| TERRAIN = ["flat", "hilly", "mountain"] | |
| TERRAIN_TO_ID = {t: i for i, t in enumerate(TERRAIN)} | |
| DEVICE = "cuda" if torch.cuda.is_available() else "cpu" | |
| # гамма формы по типу: gamma>1 -> равнина/пологое; gamma<1 -> острые пики | |
| PROFILE_GAMMA = {"flat": 2.2, "hilly": 1.7, "mountain": 0.75} | |
| # дефолтный перепад высот по типу (м) — для подсказки в UI | |
| DEFAULT_M = {"flat": 80.0, "hilly": 300.0, "mountain": 1000.0} | |
| # нормировка масштаба по умолчанию (перекроется статистикой из чекпойнта) | |
| SCALE_STATS = {"log1p_mean": 5.054, "log1p_std": 1.317} | |
| # ------------------------- модель ------------------------- | |
| def load_model(): | |
| """ConvCVAE с весами из чекпойнта (ключ 'model_state'). Возвращает (model, scale_stats).""" | |
| path = LOCAL_WEIGHTS | |
| if MODEL_REPO: | |
| try: | |
| from huggingface_hub import hf_hub_download | |
| path = hf_hub_download(repo_id=MODEL_REPO, filename="cvae_best.pt") | |
| except Exception as e: | |
| print("HF Hub load failed, fallback to local:", e) | |
| ck = torch.load(path, map_location="cpu", weights_only=False) | |
| m = (ck.get("config", {}) or {}).get("model", {}) if isinstance(ck, dict) else {} | |
| model = ConvCVAE( | |
| in_channels=int(m.get("in_channels", 1)), | |
| image_size=int(m.get("image_size", 256)), | |
| latent_dim=int(m.get("latent_dim", 128)), | |
| base_channels=int(m.get("base_channels", 32)), | |
| channel_multipliers=tuple(m.get("channel_multipliers", [1, 2, 4, 8, 16])), | |
| num_classes=int(m.get("num_classes", 3)), | |
| class_embed_dim=int(m.get("class_embed_dim", 16)), | |
| scale_embed_dim=int(m.get("scale_embed_dim", 8)), | |
| ) | |
| state = ck.get("model_state", ck.get("model", ck)) if isinstance(ck, dict) else ck | |
| model.load_state_dict(state) # strict: архитектура совпадает с обученной | |
| model.eval().to(DEVICE) | |
| stats = ck.get("scale_stats", SCALE_STATS) if isinstance(ck, dict) else SCALE_STATS | |
| return model, stats | |
| try: | |
| MODEL, SCALE_STATS = load_model() | |
| MODEL_OK = True | |
| print("CVAE загружена.") | |
| except Exception as e: | |
| print("МОДЕЛЬ НЕ ЗАГРУЖЕНА:", e) | |
| MODEL, MODEL_OK = None, False | |
| def meters_to_scale_norm(meters: float) -> float: | |
| """перепад высот (м) -> нормированное условие масштаба (как при обучении: log1p+станд.).""" | |
| return (math.log1p(max(float(meters), 0.0)) - SCALE_STATS["log1p_mean"]) / (SCALE_STATS["log1p_std"] + 1e-8) | |
| # ------------------------- генерация ------------------------- | |
| def generate_heightmap(terrain_type, height_m, seed): | |
| """CVAE -> heightmap [H,W] в [-1,1] (условие: тип + масштаб в метрах).""" | |
| torch.manual_seed(int(seed)) | |
| cls = TERRAIN_TO_ID[terrain_type] | |
| if MODEL_OK: | |
| sn = meters_to_scale_norm(height_m) | |
| hm = MODEL.sample(1, cls, float(sn), device=DEVICE).cpu().numpy()[0, 0] | |
| else: # заглушка для отладки UI, если веса не загрузились | |
| rng = np.random.default_rng(int(seed)) | |
| from scipy.ndimage import gaussian_filter | |
| hm = gaussian_filter(rng.standard_normal((256, 256)), 8) | |
| hm = hm / (np.abs(hm).max() + 1e-8) | |
| return hm.astype(np.float32) | |
| def apply_shape(hm, terrain_type): | |
| """Форма по типу (gamma): равнина пологая, горы острые. -> высота [0,1].""" | |
| h01 = (np.clip(hm, -1, 1) + 1) / 2 | |
| s = np.power(h01, PROFILE_GAMMA[terrain_type]) | |
| return normalize01(s).astype(np.float32) | |
| # ------------------------- экспорт форматов ------------------------- | |
| def save_heightmap_16bit(h01, path): | |
| """16-bit PNG (uint16) — нормированная высота для Unity Terrain / Unreal Landscape.""" | |
| arr = (np.clip(h01, 0, 1) * 65535.0).round().astype(np.uint16) | |
| Image.fromarray(arr, mode="I;16").save(path) | |
| return path | |
| def save_mesh(z_world, path, colorize=True): | |
| """Меш из поля высот (world units) -> .obj/.glb (trimesh), гладкие нормали. | |
| colorize=True кладёт ТЕКСТУРУ по высоте (трава->камень->снег) + матовый PBR- | |
| материал — её gr.Model3D (model-viewer) реально рисует цветной (вершинные | |
| цвета он игнорирует). colorize=False — обычный серый меш. | |
| """ | |
| import trimesh | |
| from trimesh.visual.material import PBRMaterial | |
| from PIL import Image | |
| H, W = z_world.shape | |
| xs, ys = np.meshgrid(np.linspace(0, 10, W), np.linspace(0, 10, H)) | |
| verts = np.stack([xs.ravel(), ys.ravel(), z_world.ravel()], axis=1) | |
| # грани сеткой (векторизованно) | |
| idx = np.arange(H * W).reshape(H, W) | |
| a = idx[:-1, :-1].ravel(); b = idx[:-1, 1:].ravel() | |
| c = idx[1:, 1:].ravel(); d = idx[1:, :-1].ravel() | |
| faces = np.concatenate([np.stack([a, b, c], 1), np.stack([a, c, d], 1)], 0) | |
| if colorize: | |
| from matplotlib.colors import LinearSegmentedColormap | |
| cmap = LinearSegmentedColormap.from_list("hero", HERO_COLORS) | |
| zn = (z_world - z_world.min()) / ((z_world.max() - z_world.min()) + 1e-9) | |
| tex = (cmap(zn)[..., :3] * 255).astype(np.uint8) # H×W RGB по высоте | |
| uu, vv = np.meshgrid(np.linspace(0, 1, W), np.linspace(0, 1, H)) | |
| uv = np.stack([uu.ravel(), vv.ravel()], axis=1) # вершина (i,j) -> свой тексель | |
| mat = PBRMaterial(name="terrain", baseColorTexture=Image.fromarray(tex), | |
| metallicFactor=0.0, roughnessFactor=1.0) | |
| visual = trimesh.visual.TextureVisuals(uv=uv, material=mat) | |
| mesh = trimesh.Trimesh(vertices=verts, faces=faces, visual=visual, process=False) | |
| else: | |
| mesh = trimesh.Trimesh(vertices=verts, faces=faces, process=False) | |
| mesh.fix_normals() | |
| _ = mesh.vertex_normals # посчитать гладкие нормали | |
| mesh.export(path) | |
| return path | |
| def save_glb(s01, height_m, path, colorize=True): | |
| """GLB для gr.Model3D и движков: **Y-вверх** (стандарт glTF), центрирован в (0,·,0), | |
| с вертикальной экзажерацией по типу — чтобы рельеф было видно сразу, без вращения. | |
| Почему отдельно от save_mesh: .obj отдаём Blender-у (Z-вверх), а glTF/Model3D ждёт | |
| Y-вверх. Если отдать Z-вверх меш, model-viewer показывает рельеф «стеной» лицом к | |
| камере (плоский квадрат) — это и был баг с пустым/неповорачиваемым 3D. | |
| """ | |
| import trimesh | |
| H, W = s01.shape | |
| zexag = 0.08 + 0.62 * min(float(height_m), 1100.0) / 1100.0 # высота резко по классу: | |
| yh = (np.clip(s01, 0, 1) * (8.0 * zexag)).astype(np.float32) # flat ~плоский, гора высокая | |
| ax = np.linspace(-4, 4, W); ay = np.linspace(-4, 4, H) | |
| gx, gy = np.meshgrid(ax, ay) | |
| verts = np.stack([gx.ravel(), yh.ravel(), gy.ravel()], axis=1) # [X, Y=высота, Z] | |
| idx = np.arange(H * W).reshape(H, W) | |
| a = idx[:-1, :-1].ravel(); b = idx[:-1, 1:].ravel() | |
| c = idx[1:, 1:].ravel(); d = idx[1:, :-1].ravel() | |
| faces = np.concatenate([np.stack([a, b, c], 1), np.stack([a, c, d], 1)], 0) | |
| # doubleSided=True — иначе model-viewer отбраковывает задние грани и поверхность | |
| # видна насквозь с одной стороны (сверху прозрачно, снизу плотно). Двусторонний | |
| # материал рисует обе стороны -> рельеф плотный под любым углом. | |
| from trimesh.visual.material import PBRMaterial | |
| if colorize: | |
| from matplotlib.colors import LinearSegmentedColormap | |
| cmap = LinearSegmentedColormap.from_list("hero", HERO_COLORS) | |
| zn = (yh - yh.min()) / ((yh.max() - yh.min()) + 1e-9) | |
| tex = (cmap(zn)[..., :3] * 255).astype(np.uint8) # цвет по высоте (текстура) | |
| uu, vv = np.meshgrid(np.linspace(0, 1, W), np.linspace(0, 1, H)) | |
| uv = np.stack([uu.ravel(), vv.ravel()], axis=1) | |
| mat = PBRMaterial(name="terrain", baseColorTexture=Image.fromarray(tex), | |
| metallicFactor=0.0, roughnessFactor=1.0, doubleSided=True) | |
| visual = trimesh.visual.TextureVisuals(uv=uv, material=mat) | |
| else: | |
| mat = PBRMaterial(name="terrain_raw", baseColorFactor=[0.62, 0.62, 0.62, 1.0], | |
| metallicFactor=0.0, roughnessFactor=1.0, doubleSided=True) | |
| visual = trimesh.visual.TextureVisuals(material=mat) | |
| mesh = trimesh.Trimesh(vertices=verts, faces=faces, visual=visual, process=False) | |
| mesh.fix_normals() | |
| _ = mesh.vertex_normals | |
| mesh.export(path) | |
| return path | |
| def render_preview(z_world, path): | |
| """Лёгкое превью: высотная палитра + hillshade (matplotlib).""" | |
| import matplotlib | |
| matplotlib.use("Agg") | |
| import matplotlib.pyplot as plt | |
| from matplotlib.colors import LightSource, LinearSegmentedColormap | |
| cmap = LinearSegmentedColormap.from_list("t", [ | |
| "#2f5d3a", "#5a7d44", "#8a8159", "#9c8569", "#cfc9c0", "#ffffff"]) | |
| ls = LightSource(azdeg=315, altdeg=45) | |
| rgb = ls.shade(z_world, cmap=cmap, vert_exag=2.0, blend_mode="soft") | |
| fig, ax = plt.subplots(figsize=(6, 6)) | |
| ax.imshow(rgb); ax.axis("off") | |
| fig.savefig(path, dpi=130, bbox_inches="tight", pad_inches=0) | |
| plt.close(fig) | |
| return path | |
| # палитра рельефа: трава -> луга -> камень -> снег | |
| HERO_COLORS = ["#244a30", "#3f6b3a", "#6f7d4a", "#8a8159", "#aa9c84", "#d9d3c9", "#ffffff"] | |
| def render_side(s01, height_m, terrain_type, path): | |
| """Кинематографичный кадр СБОКУ (как фотография гор) — статичный, не крутилка. | |
| Низкий угол камеры + градиент неба + высотная палитра с затенением; вертикальная | |
| драматизация по перепаду высот, чтобы пики вырастали на фоне неба, как на фото. | |
| """ | |
| import matplotlib | |
| matplotlib.use("Agg") | |
| import matplotlib.pyplot as plt | |
| from matplotlib.colors import LightSource, LinearSegmentedColormap | |
| terr = LinearSegmentedColormap.from_list("hero", HERO_COLORS) | |
| sky = LinearSegmentedColormap.from_list("sky", ["#5b79a6", "#9fb4d2", "#e2e9f1"]) | |
| z = s01[::2, ::2] # 128x128 — быстрее на CPU | |
| H, W = z.shape | |
| xs, ys = np.meshgrid(np.linspace(0, 1, W), np.linspace(0, 1, H)) | |
| ls = LightSource(azdeg=300, altdeg=18) # низкое солнце -> длинные тени | |
| rgb = ls.shade(z, cmap=terr, vert_exag=3.2, blend_mode="soft") | |
| zexag = 0.08 + 0.62 * min(float(height_m), 1100.0) / 1100.0 # flat плоский, гора высокая (как в меше) | |
| fig = plt.figure(figsize=(12.5, 6.2)) # крупный кинокадр | |
| sky_ax = fig.add_axes([0, 0, 1, 1]); sky_ax.axis("off") # небо-градиент позади | |
| grad = np.linspace(0, 1, 256).reshape(-1, 1) | |
| sky_ax.imshow(grad, aspect="auto", cmap=sky, origin="upper", extent=[0, 1, 0, 1]) | |
| ax = fig.add_axes([0, 0, 1, 1], projection="3d") # 3D поверх, прозрачно | |
| ax.patch.set_alpha(0.0) | |
| ax.plot_surface(xs, ys, z, facecolors=rgb, rstride=1, cstride=1, | |
| linewidth=0, antialiased=False, shade=False) | |
| ax.set_box_aspect((1, 1, zexag)) | |
| ax.set_axis_off() | |
| ax.view_init(elev=12, azim=-62) # НИЗКИЙ угол = вид сбоку | |
| fig.savefig(path, dpi=150, pad_inches=0) | |
| plt.close(fig) | |
| return path | |
| # ------------------------- основной колбэк ------------------------- | |
| def run(terrain_type, seed, use_erosion, erosion_strength, height_m): | |
| # высота — авто по типу, если в «тонкой настройке» не задана вручную (0) | |
| if not height_m or float(height_m) <= 0: | |
| height_m = DEFAULT_M[terrain_type] | |
| hm = generate_heightmap(terrain_type, height_m, seed) # CVAE [-1,1] | |
| s01 = apply_shape(hm, terrain_type) # форма по типу -> [0,1] | |
| note = "" | |
| if use_erosion: | |
| droplets = int(8000 + float(erosion_strength) * 20000) # 8k..28k — умеренно для CPU | |
| s01 = hydraulic_erosion(s01, num_droplets=droplets, seed=int(seed), lifetime=35, | |
| erode_speed=0.45, sediment_capacity=7.0, min_slope=0.001, | |
| erosion_radius=2, deposit_speed=0.25) | |
| s01 = normalize01(s01).astype(np.float32) | |
| note = f" · эрозия on" | |
| vmax = (float(height_m) / 1000.0) * 4.0 # высота меша в world units | |
| z_world = s01 * vmax | |
| tmp = tempfile.mkdtemp() | |
| png16 = save_heightmap_16bit(s01, os.path.join(tmp, "heightmap_16bit.png")) | |
| obj = save_mesh(z_world, os.path.join(tmp, "terrain.obj"), colorize=True) # Blender (Z-вверх) | |
| glb = save_glb(s01, height_m, os.path.join(tmp, "terrain.glb"), colorize=True) # цветная крутилка (Y-вверх) | |
| glb_raw = save_glb(s01, height_m, os.path.join(tmp, "terrain_raw.glb"), colorize=False) # без покраски | |
| hero = render_side(s01, height_m, terrain_type, os.path.join(tmp, "hero_mountains.png")) # вид на горы | |
| preview = render_preview(z_world, os.path.join(tmp, "preview.png")) # карта высот | |
| info = f"**{terrain_type}** · перепад ~{int(height_m)} м · seed {int(seed)}{note}" | |
| if not MODEL_OK: | |
| info += " \n⚠️ веса не загружены — показана заглушка" | |
| return glb, glb_raw, hero, preview, info, [png16, obj, glb, glb_raw, hero, preview] | |
| # ------------------------- интерфейс ------------------------- | |
| THEME = gr.themes.Soft(primary_hue="emerald", secondary_hue="blue", neutral_hue="slate") | |
| with gr.Blocks(theme=THEME, title="GameTerrain — генератор игрового рельефа") as demo: | |
| gr.Markdown("# 🏔️ GameTerrain") | |
| with gr.Row(equal_height=False): | |
| with gr.Column(scale=4): | |
| terrain = gr.Radio(TERRAIN, value="mountain", label="Тип рельефа") | |
| seed = gr.Slider(0, 9999, value=42, step=1, label="Seed — крути для нового варианта") | |
| erosion = gr.Checkbox(value=False, label="🌊 Эрозия — добавить реализм (русла, гребни)") | |
| with gr.Accordion("⚙️ Тонкая настройка (необязательно)", open=False): | |
| height_m = gr.Slider(0, 2500, value=0, step=10, | |
| label="Перепад высот, м · 0 = авто по типу") | |
| er_strength = gr.Slider(0.0, 1.0, value=0.5, step=0.1, label="Сила эрозии") | |
| btn = gr.Button("✨ Сгенерировать рельеф", variant="primary", size="lg") | |
| info = gr.Markdown() | |
| with gr.Column(scale=6): | |
| with gr.Row(): | |
| model3d = gr.Model3D(label="🌀 3D — цветная (вращай)", height=270, | |
| camera_position=(-25, 62, None), clear_color=(0.90, 0.93, 0.98, 1.0)) | |
| model3d_raw = gr.Model3D(label="🌀 3D — без покраски (вращай)", height=270, | |
| camera_position=(-25, 62, None), clear_color=(0.90, 0.93, 0.98, 1.0)) | |
| hero_img = gr.Image(label="вид сбоку на рельеф", type="filepath", height=380) | |
| with gr.Row(): | |
| preview_img = gr.Image(label="🗺️ Карта высот (hillshade)", type="filepath", height=200) | |
| files = gr.File(label="⬇️ Скачать (PNG·OBJ·GLB)", file_count="multiple", height=200) | |
| btn.click(run, [terrain, seed, erosion, er_strength, height_m], | |
| [model3d, model3d_raw, hero_img, preview_img, info, files]) | |
| if __name__ == "__main__": | |
| demo.launch(server_name="0.0.0.0", server_port=7860) | |