app.py

#!/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()