#!/usr/bin/env python3 # -*- coding: utf-8 -*- import gradio as gr import torch import numpy as np import cv2 from transformers import SegformerImageProcessor, SegformerForSemanticSegmentation from PIL import Image import json from scipy import ndimage from scipy.ndimage import binary_opening, binary_closing, sobel, binary_dilation, median_filter import zipfile import io import os import tempfile import shutil import warnings warnings.filterwarnings('ignore') # グローバル変数 device = "cuda" if torch.cuda.is_available() else "cpu" processor = None model = None # カテゴリ定義 ADE20K_TO_CITY_MAPPING = { 'road': 'road', 'street': 'road', 'path': 'road', 'sidewalk': 'road', 'building': 'building_c', 'house': 'building_a', 'skyscraper': 'building_e', 'highrise': 'building_e', 'tower': 'building_e', 'office': 'building_d', 'shop': 'building_b', 'store': 'building_b', 'apartment': 'building_c', 'hotel': 'building_d', 'tree': 'forest', 'plant': 'forest', 'palm': 'forest', 'grass': 'park', 'field': 'park', 'flower': 'park', 'water': 'water', 'sea': 'water', 'river': 'water', 'lake': 'water', 'earth': 'bare_land', 'sand': 'bare_land', 'ground': 'bare_land', 'parking lot': 'infrastructure', 'stadium': 'building_d', } CITY_CATEGORIES = { 'road': {'label': '道路', 'color': (128, 64, 128), 'height': 0, 'semantic_id': 0}, 'forest': {'label': '森林', 'color': (34, 139, 34), 'height': 1.5, 'semantic_id': 1}, 'park': {'label': '公園/緑地', 'color': (144, 238, 144), 'height': 0.5, 'semantic_id': 2}, 'water': {'label': '水域', 'color': (30, 144, 255), 'height': 0, 'semantic_id': 3}, 'building_a': {'label': '建物A（小）', 'color': (255, 200, 150), 'height': 0.6, 'semantic_id': 4}, 'building_b': {'label': '建物B（中小）', 'color': (255, 160, 122), 'height': 1.0, 'semantic_id': 5}, 'building_c': {'label': '建物C（中）', 'color': (240, 120, 90), 'height': 1.5, 'semantic_id': 6}, 'building_d': {'label': '建物D（中大）', 'color': (220, 80, 60), 'height': 2.2, 'semantic_id': 7}, 'building_e': {'label': '建物E（大）', 'color': (200, 40, 40), 'height': 3.0, 'semantic_id': 8}, 'bare_land': {'label': '空き地', 'color': (210, 180, 140), 'height': 0.1, 'semantic_id': 9}, 'infrastructure': {'label': 'インフラ', 'color': (100, 100, 100), 'height': 0.8, 'semantic_id': 10}, 'other': {'label': 'その他/境界', 'color': (80, 80, 80), 'height': 0, 'semantic_id': 11} } def load_model(): global processor, model if processor is None or model is None: processor = SegformerImageProcessor.from_pretrained("nvidia/segformer-b5-finetuned-ade-640-640") model = SegformerForSemanticSegmentation.from_pretrained("nvidia/segformer-b5-finetuned-ade-640-640").to(device) return processor, model def map_ade20k_to_city(class_id, id2label): if class_id not in id2label: return 'other' class_name = id2label[class_id].lower() for ade_name, city_cat in ADE20K_TO_CITY_MAPPING.items(): if ade_name in class_name: return city_cat if any(w in class_name for w in ['skyscraper', 'highrise', 'tower']): return 'building_e' elif any(w in class_name for w in ['office', 'hotel', 'commercial', 'stadium']): return 'building_d' elif any(w in class_name for w in ['building', 'apartment']): return 'building_c' elif any(w in class_name for w in ['shop', 'store', 'market']): return 'building_b' elif any(w in class_name for w in ['house', 'home', 'shed', 'hut']): return 'building_a' return 'other' def segment_with_tiling(image, processor, model, tile_size=320, overlap=64, use_tiling=True): h, w = image.shape[:2] if not use_tiling or (h <= tile_size and w <= tile_size): pil_image = Image.fromarray(image) inputs = processor(images=pil_image, return_tensors="pt").to(device) with torch.no_grad(): outputs = model(**inputs) logits = outputs.logits upsampled_logits = torch.nn.functional.interpolate( logits, size=image.shape[:2], mode="bilinear", align_corners=False ) return upsampled_logits.argmax(dim=1)[0].cpu().numpy() stride = tile_size - overlap num_tiles_h = (h - overlap) // stride + (1 if (h - overlap) % stride > 0 else 0) num_tiles_w = (w - overlap) // stride + (1 if (w - overlap) % stride > 0 else 0) votes = np.zeros((h, w, 150), dtype=np.float32) counts = np.zeros((h, w), dtype=np.float32) for i in range(num_tiles_h): for j in range(num_tiles_w): y_start = i * stride x_start = j * stride y_end = min(y_start + tile_size, h) x_end = min(x_start + tile_size, w) tile = image[y_start:y_end, x_start:x_end] pil_tile = Image.fromarray(tile) inputs = processor(images=pil_tile, return_tensors="pt").to(device) with torch.no_grad(): outputs = model(**inputs) logits = outputs.logits upsampled = torch.nn.functional.interpolate( logits, size=tile.shape[:2], mode="bilinear", align_corners=False ) probs = torch.nn.functional.softmax(upsampled, dim=1)[0].cpu().numpy() votes[y_start:y_end, x_start:x_end] += probs.transpose(1, 2, 0) counts[y_start:y_end, x_start:x_end] += 1 counts = np.maximum(counts, 1) final_votes = votes / counts[:, :, np.newaxis] return final_votes.argmax(axis=2) def create_colored_segmentation(seg_map): h, w = seg_map.shape colored = np.zeros((h, w, 3), dtype=np.uint8) for cat_name, cat_info in CITY_CATEGORIES.items(): mask = seg_map == cat_info['semantic_id'] colored[mask] = cat_info['color'] return colored def detect_boundaries(segmentation_map, thickness=5): edges_h = np.abs(sobel(segmentation_map.astype(float), axis=0)) > 0 edges_v = np.abs(sobel(segmentation_map.astype(float), axis=1)) > 0 boundaries = edges_h | edges_v if thickness > 1: kernel = np.ones((thickness, thickness), dtype=bool) boundaries = binary_dilation(boundaries, structure=kernel) other_id = CITY_CATEGORIES['other']['semantic_id'] segmentation_map[boundaries] = other_id return segmentation_map, boundaries def create_3d_mesh(segments, image, resolution=2): height, width = image.shape[:2] meshes_data = [] for idx, segment in enumerate(segments): segmentation = segment['segmentation'] bbox = segment['bbox'] x, y, w, h = bbox if w < 3 or h < 3: continue segment_area = segmentation[y:y+h, x:x+w] segment_image = image[y:y+h, x:x+w] if not segment_area.any(): continue vertices = [] faces = [] colors = [] step = resolution building_height = segment['height'] * 0.5 for sy in range(0, segment_area.shape[0] - step, step): for sx in range(0, segment_area.shape[1] - step, step): if not segment_area[sy, sx]: continue world_x = (x + sx - width/2) * 0.1 world_z = (y + sy - height/2) * 0.1 base_idx = len(vertices) vertices.extend([ [float(world_x), float(building_height), float(world_z)], [float(world_x + step*0.1), float(building_height), float(world_z)], [float(world_x + step*0.1), float(building_height), float(world_z + step*0.1)], [float(world_x), float(building_height), float(world_z + step*0.1)] ]) vertices.extend([ [float(world_x), 0.0, float(world_z)], [float(world_x + step*0.1), 0.0, float(world_z)], [float(world_x + step*0.1), 0.0, float(world_z + step*0.1)], [float(world_x), 0.0, float(world_z + step*0.1)] ]) if sy < segment_image.shape[0] and sx < segment_image.shape[1]: color = segment_image[sy, sx] / 255.0 color_list = [float(color[0]), float(color[1]), float(color[2])] else: color_list = [0.5, 0.5, 0.5] wall_color = [c * 0.7 for c in color_list] colors.extend([color_list] * 4 + [wall_color] * 4) faces.extend([ [base_idx, base_idx+1, base_idx+2], [base_idx, base_idx+2, base_idx+3] ]) if len(vertices) > 0: meshes_data.append({ 'id': int(idx), 'category': str(segment['category']), 'label': str(segment['label']), 'semantic_id': int(segment['semantic_id']), 'vertices': vertices, 'faces': faces, 'colors': colors, 'center': [ float((x + w/2 - width/2) * 0.1), float(segment['height'] * 0.5), float((y + h/2 - height/2) * 0.1) ], 'bbox': [int(x), int(y), int(w), int(h)], 'area': float(segment['area']), 'height': float(segment['height']) }) return meshes_data def process_image(image, max_size, tile_size, tile_overlap, min_area, mesh_res, apply_morphology, morph_kernel, detect_bound, bound_thickness, apply_smoothing, smooth_iter, use_tiling): # モデルロード processor, model = load_model() # 画像前処理 if isinstance(image, str): original_image = cv2.imread(image) original_image = cv2.cvtColor(original_image, cv2.COLOR_BGR2RGB) else: original_image = np.array(image) original_height, original_width = original_image.shape[:2] if max(original_height, original_width) > max_size: scale_factor = max_size / max(original_height, original_width) new_width = int(original_width * scale_factor) new_height = int(original_height * scale_factor) resized_image = cv2.resize(original_image, (new_width, new_height), interpolation=cv2.INTER_LANCZOS4) else: resized_image = original_image.copy() # セグメンテーション predicted_seg = segment_with_tiling(resized_image, processor, model, tile_size, tile_overlap, use_tiling) # クラスマッピング city_segmentation = np.zeros(predicted_seg.shape, dtype=np.uint8) id2label = model.config.id2label for class_id in np.unique(predicted_seg): city_category = map_ade20k_to_city(class_id, id2label) semantic_id = CITY_CATEGORIES[city_category]['semantic_id'] mask = predicted_seg == class_id city_segmentation[mask] = semantic_id # クラス平滑化 if apply_smoothing: for _ in range(smooth_iter): city_segmentation = median_filter(city_segmentation, size=3) # 境界検出 if detect_bound: city_segmentation, boundary_mask = detect_boundaries(city_segmentation, bound_thickness) # セグメント抽出 segments_data = [] segment_id = 0 for cat_name, cat_info in CITY_CATEGORIES.items(): semantic_id = cat_info['semantic_id'] mask = city_segmentation == semantic_id if not mask.any(): continue if apply_morphology: kernel = np.ones((morph_kernel, morph_kernel), dtype=bool) mask = binary_opening(mask, structure=kernel) mask = binary_closing(mask, structure=kernel) labeled, num_features = ndimage.label(mask) for i in range(1, num_features + 1): segment_mask = labeled == i area = np.sum(segment_mask) if area < min_area: continue rows, cols = np.where(segment_mask) if len(rows) == 0: continue y_min, y_max = rows.min(), rows.max() x_min, x_max = cols.min(), cols.max() segments_data.append({ 'id': segment_id, 'category': cat_name, 'label': cat_info['label'], 'semantic_id': semantic_id, 'color': cat_info['color'], 'height': cat_info['height'], 'segmentation': segment_mask, 'bbox': [int(x_min), int(y_min), int(x_max - x_min), int(y_max - y_min)], 'area': int(area) }) segment_id += 1 # 3Dメッシュ生成 meshes = create_3d_mesh(segments_data, resized_image, mesh_res) # メタデータ metadata = { 'version': '2.1', 'total_segments': len(meshes), 'categories': {} } for mesh in meshes: cat = mesh['category'] if cat not in metadata['categories']: metadata['categories'][cat] = {'label': mesh['label'], 'count': 0} metadata['categories'][cat]['count'] += 1 # 可視化 colored_seg = create_colored_segmentation(city_segmentation) overlay = (resized_image.astype(np.float32) * 0.5 + colored_seg.astype(np.float32) * 0.5).astype(np.uint8) # JSONファイル作成 output_data = {'metadata': metadata, 'meshes': meshes} json_str = json.dumps(output_data, ensure_ascii=False, indent=2) # ZIPファイル作成（一時ファイルとして保存） import tempfile import shutil temp_dir = tempfile.mkdtemp() zip_path = os.path.join(temp_dir, 'city_3d_output.zip') with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zip_file: zip_file.writestr('city_3d_model.json', json_str) # セグメンテーション画像 _, buffer = cv2.imencode('.png', cv2.cvtColor(colored_seg, cv2.COLOR_RGB2BGR)) zip_file.writestr('segmentation_result.png', buffer.tobytes()) # 統計情報 stats = f"総セグメント数: {len(meshes)}\n\n" for cat, info in metadata['categories'].items(): stats += f"{info['label']}: {info['count']}個\n" return colored_seg, overlay, stats, zip_path # Gradio UI with gr.Blocks(title="3D City Map Generator") as demo: gr.Markdown("# 🏙️ 3D City Map Generator") gr.Markdown("航空写真から3D都市マップを生成します（Segformer B5モデル使用）") with gr.Row(): with gr.Column(): input_image = gr.Image(label="航空写真をアップロード", type="numpy") with gr.Accordion("⚙️ 詳細設定", open=False): max_size = gr.Slider(640, 2048, value=800, step=64, label="最大画像サイズ") use_tiling = gr.Checkbox(value=True, label="タイル分割処理を使用") tile_size = gr.Slider(120, 640, value=320, step=40, label="タイルサイズ") tile_overlap = gr.Slider(32, 128, value=64, step=8, label="タイル重複") min_area = gr.Slider(20, 200, value=32, step=4, label="最小セグメント面積") mesh_res = gr.Slider(1, 4, value=3, step=1, label="メッシュ解像度") apply_morphology = gr.Checkbox(value=True, label="モルフォロジー処理") morph_kernel = gr.Slider(3, 9, value=7, step=2, label="モルフォロジーカーネル") detect_bound = gr.Checkbox(value=True, label="境界検出") bound_thickness = gr.Slider(1, 5, value=5, step=1, label="境界の太さ") apply_smoothing = gr.Checkbox(value=True, label="クラス平滑化") smooth_iter = gr.Slider(1, 3, value=2, step=1, label="平滑化反復回数") process_btn = gr.Button("🚀 3Dマップ生成", variant="primary") with gr.Column(): seg_output = gr.Image(label="セグメンテーション結果") overlay_output = gr.Image(label="オーバーレイ") stats_output = gr.Textbox(label="統計情報", lines=10) download_output = gr.File(label="📥 3Dモデルをダウンロード (ZIP)") process_btn.click( fn=process_image, inputs=[input_image, max_size, tile_size, tile_overlap, min_area, mesh_res, apply_morphology, morph_kernel, detect_bound, bound_thickness, apply_smoothing, smooth_iter, use_tiling], outputs=[seg_output, overlay_output, stats_output, download_output] ) gr.Markdown(""" ### 使い方 1. 航空写真をアップロード 2. 必要に応じてパラメータを調整 3. 「3Dマップ生成」をクリック 4. ZIPファイルをダウンロードして、JSONファイルをBlenderなどで使用 ### パラメータ説明 - **最大画像サイズ**: 大きいほど精度向上（処理時間増加） - **タイル分割**: 大きな画像の精度向上に重要 - **最小セグメント面積**: 増やすとノイズ削減 - **メッシュ解像度**: 増やすとファイルサイズ減少 - **境界検出**: 建物と道路の混合を防ぐ """) if __name__ == "__main__": demo.launch()