test.py

# -*- coding: utf-8 -*-
import geopandas as gpd
import pandas as pd
import rasterio
from rasterio.windows import from_bounds
import numpy as np
import torch
import torchvision.models as models
import segmentation_models_pytorch as smp
import albumentations as A
from albumentations.pytorch import ToTensorV2
from tqdm import tqdm
import os
import warnings

# 경고 메시지 숨기기 (TIFF 등)
warnings.filterwarnings("ignore")

# --- 1. 설정 (Configuration) ---
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {DEVICE}")

# [필수 수정] 파일 경로
VECTOR_FILE = r"./data/SF_2025_Filtered_by_2022.shp"  # 검증할 건물 벡터(SHP)
STACK_FILE  = r"./data/training_data/K5_Stack.tif"    # SAR 이미지 스택 (band1:TC, band2:SLC)
MASK_FILE   = r"./data/module_area.tif"               # 전체 영역/풋프린트 마스크

# [필수 수정] 모델 경로
MODEL_A_PATH = "./pth_folder/best_model_unet.pth"    # 모듈 A (면적; 세그멘테이션)
MODEL_B_PATH = "./pth_folder/best_model_height.pth"  # 모듈 B (높이; 회귀)

# [필수 수정] SHP 컬럼 이름
COL_ID     = "mblr"   # 건물 고유 ID
COL_HEIGHT = "hgt_maxcm"     # '허가 높이' 대용 (cm 단위)

# === 임계값(판정 기준) : 첨부 이미지 규칙 반영 ===
# 면적 임계값(건물별 동적) = (실측면적 × (1 - IoU평균)) + 유의미 최소 증축 면적
MEAN_IOU_A         = 0.59  # [수정] 모듈 A 검증 평균 IoU (예: 0.80). 본인 실험치로 교체
MIN_SIG_AREA_M2    = 2.0   # [수정] 유의미 최소 증축 면적(예: 2~3 m² 권장)

# 높이 임계값(고정) = RMSE(1.98m) + 최소 층고(2.29m) = 4.27m
RMSE_HEIGHT_M        = 1.98
MIN_STOREY_HEIGHT_M  = 2.29
TH_HEIGHT_FIXED      = RMSE_HEIGHT_M + MIN_STOREY_HEIGHT_M  # 4.27 m

# 칩(타일) 설정
CHIP_SIZE_METERS = 80   # 건물 중심 기준 잘라낼 실제 길이(m)
IMG_SIZE         = 256  # 모델 입력 해상도
# ---------------------------------

# --- 2. 모델 로드 ---
def load_models():
    print("모델 로딩 중...")
    # [모듈 A] U-Net (Segmentation)
    model_a = smp.Unet(
        encoder_name="resnet34",
        encoder_weights=None,
        in_channels=1,
        classes=1,
        activation='sigmoid'
    ).to(DEVICE)
    try:
        model_a.load_state_dict(torch.load(MODEL_A_PATH, map_location=DEVICE))
        model_a.eval()
    except Exception as e:
        print(f"Error loading Module A: {e}")
        exit()

    # [모듈 B] ResNet34 (Regression; 입력 2채널: SLC+Mask → 1개 스칼라 높이)
    model_b = models.resnet34(weights=None)
    model_b.conv1 = torch.nn.Conv2d(2, 64, kernel_size=7, stride=2, padding=3, bias=False)
    model_b.fc    = torch.nn.Linear(model_b.fc.in_features, 1)
    model_b.to(DEVICE)
    try:
        model_b.load_state_dict(torch.load(MODEL_B_PATH, map_location=DEVICE))
        model_b.eval()
    except Exception as e:
        print(f"Error loading Module B: {e}")
        exit()

    print("✅ 두 AI 모델 로드 완료.")
    return model_a, model_b

# --- 3. 전처리/유틸 ---
transform_common = A.Compose([A.Resize(IMG_SIZE, IMG_SIZE), ToTensorV2()])

def get_chips(src_stack, src_mask, geom, chip_size_m):
    """건물 중심 좌표 기준으로 SAR 및 마스크 칩을 잘라냅니다."""
    center = geom.centroid
    half = chip_size_m / 2.0
    left, right  = center.x - half, center.x + half
    bottom, top  = center.y - half, center.y + half
    window = from_bounds(left, bottom, right, top, src_stack.transform)

    chip_tc  = src_stack.read(1, window=window)  # 모듈 A용 (TC)
    chip_slc = src_stack.read(2, window=window)  # 모듈 B용 (SLC)
    chip_mask = src_mask.read(1, window=window)  # 모듈 B용 (Mask)
    return chip_tc, chip_slc, chip_mask

def prepare_inputs(chip_tc, chip_slc, chip_mask):
    """잘라낸 칩을 모델 입력 텐서로 변환."""
    # 백분위수 기반 정규화
    def normalize(arr):
        if arr.size > 0 and np.ptp(arr) > 1e-6:
            p1, p99 = np.percentile(arr, (1, 99))
            arr = np.clip(arr, p1, p99)
            if p99 > p1:
                return (arr - p1) / (p99 - p1)
        return np.zeros_like(arr)

    norm_tc  = normalize(chip_tc)
    norm_slc = normalize(chip_slc)

    # [A 입력] 1채널 → (H,W,1) → Resize+ToTensor → (1,1,H,W)
    if norm_tc.ndim == 2:
        norm_tc = np.expand_dims(norm_tc, axis=-1)
    else:
        norm_tc = norm_tc.transpose(1, 2, 0)
    input_a = transform_common(image=norm_tc)['image'].unsqueeze(0).to(DEVICE).float()

    # [B 입력] 2채널(SLC+Mask) → (H,W,2) → Resize+ToTensor → (1,2,H,W)
    if norm_slc.ndim == 3:
        norm_slc = norm_slc.squeeze(0)
    if chip_mask.ndim == 3:
        chip_mask = chip_mask.squeeze(0)
    stacked_b = np.stack([norm_slc, chip_mask], axis=-1)
    input_b = transform_common(image=stacked_b)['image'].unsqueeze(0).to(DEVICE).float()

    return input_a, input_b

# --- 4. 메인 분석 ---
model_a, model_b = load_models()

print("데이터 파일 로딩 중...")
try:
    gdf = gpd.read_file(VECTOR_FILE)
    src_stack = rasterio.open(STACK_FILE)
    src_mask  = rasterio.open(MASK_FILE)
except Exception as e:
    print(f"Error: 데이터 파일 열기 실패: {e}")
    exit()

# 좌표계 통일 (SHP → 래스터 CRS)
if gdf.crs != src_stack.crs:
    print(f"🔄 좌표계 변환 중... ({gdf.crs} -> {src_stack.crs})")
    gdf = gdf.to_crs(src_stack.crs)
    print("✅ 좌표계 변환 완료.")

results = []
print(f"🚀 총 {len(gdf)}개 건물에 대한 종합 정밀 분석 시작...")

first_error_reported = False
area_per_px_resized = (CHIP_SIZE_METERS / IMG_SIZE) ** 2  # 80m 칩을 256px로 본 면적/픽셀

for _, row in tqdm(gdf.iterrows(), total=len(gdf), desc="Analyzing"):
    try:
        # 1) 칩 추출
        chip_tc, chip_slc, chip_mask_raw = get_chips(src_stack, src_mask, row.geometry, CHIP_SIZE_METERS)

        # 경계 벗어난 경우 스킵
        if chip_tc.size == 0 or chip_slc.size == 0 or chip_mask_raw.size == 0:
            continue
        if chip_tc.shape[0] == 0 or chip_tc.shape[1] == 0:
            continue

        # 2) 입력 텐서
        input_a, input_b = prepare_inputs(chip_tc, chip_slc, chip_mask_raw)

        # 3) 모델 추론
        with torch.no_grad():
            # [모듈 A] 면적 마스크
            pred_mask = model_a(input_a)            # (1,1,H,W) in [0,1]
            pred_pixels = torch.sum(pred_mask > 0.5).item()
            # 면적 환산: (예측 픽셀 수) × (리사이즈 좌표계에서의 픽셀 면적)
            pred_area_m2 = float(pred_pixels) * area_per_px_resized

            # [모듈 B] 높이(m) 회귀
            pred_height_m = float(model_b(input_b).item())

        # 4) 실측(법적) 값
        legal_area   = float(row.geometry.area)         # CRS가 m기반이면 m²
        legal_height = float(row[COL_HEIGHT]) / 100.0   # cm → m

        # 5) 차이
        diff_area   = pred_area_m2 - legal_area
        diff_height = pred_height_m - legal_height

        # 6) 임계값 계산 (이미지 규칙)
        # 면적 임계: 실측×(1-mean IoU) + 유의미 최소 면적
        th_area   = max(0.0, legal_area * (1.0 - MEAN_IOU_A)) + MIN_SIG_AREA_M2
        # 높이 임계: 고정 4.27 m
        th_height = TH_HEIGHT_FIXED

        # 7) 판정
        status = "정상"
        if diff_area > th_area and diff_height > th_height:
            status = "복합 불법 증축 의심"
        elif diff_area > th_area:
            status = "수평 증축 의심 (면적)"
        elif diff_height > th_height:
            status = "수직 증축 의심 (높이)"

        # 8) 결과 누적
        res_dict = row.to_dict()
        res_dict.update({
            'status': status,
            'AI_area': round(pred_area_m2, 1),
            'AI_hgt': round(pred_height_m, 1),
            'diff_area': round(diff_area, 1),
            'diff_hgt': round(diff_height, 1),
            'th_area': round(th_area, 1),
            'th_hgt': round(th_height, 2),
            'mean_iou_used': MEAN_IOU_A,
            'min_sig_area_m2': MIN_SIG_AREA_M2
        })
        results.append(res_dict)

    except Exception as e:
        if not first_error_reported:
            # 처음 1건만 자세히 보고 (원하면 주석 해제)
            # print(f"\n[First Error] ID={row.get(COL_ID)} :: {e}")
            first_error_reported = True
        continue

# --- 5. 결과 저장 ---
src_stack.close()
src_mask.close()

if len(results) > 0:
    print("분석 완료. 결과 파일 생성 중...")
    gdf_results = gpd.GeoDataFrame(results, geometry='geometry', crs=gdf.crs)

    output_file = "final_analysis_complete.geojson"
    # GeoJSON은 EPSG:4326 선호 → 변환 저장
    gdf_results.to_crs(epsg=4326).to_file(output_file, driver="GeoJSON")

    print("=" * 30)
    print(f"✅ 모든 작업 완료! 최종 결과 파일: {output_file}")
    print(f"📊 분석 요약 (총 {len(gdf_results)}개 건물):")
    try:
        print(gdf_results['status'].value_counts())
    except Exception:
        pass
else:
    print("\n❌ 오류: 분석된 건물이 없습니다. 입력 데이터와 좌표계를 다시 확인하세요.")