space_hackaton/tiling.py at main · wisejo/space_hackaton · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
import rasterio
from rasterio.windows import Window
import os
import numpy as np
from tqdm import tqdm

# --- 1. 설정: [경로 복원됨] ---
STACK_FILE = r"./data/training_data/K5_Stack.tif"
RASTERIZED_MASK_FILE = r"./data/module_area.tif"

CHIP_SIZE = 256
STRIDE = 128

# --- 2. 설정: [수정] 모듈 A를 위한 새 폴더 ---
OUT_DIR = r"./data/training_dataset_A"
DIR_MOD_A = os.path.join(OUT_DIR, "module_A_input")
DIR_MASK = os.path.join(OUT_DIR, "masks")
os.makedirs(DIR_MOD_A, exist_ok=True)
os.makedirs(DIR_MASK, exist_ok=True)
# ------------------------------------

chip_count = 0

print(f"'{STACK_FILE}'과 '{RASTERIZED_MASK_FILE}'을 사용하여 모듈 A 칩을 재생성합니다...")
print(f"저장 위치: {OUT_DIR}")

# --- 3. 거대 TIF 파일 열기 ---
with rasterio.open(STACK_FILE) as src_stack:
    with rasterio.open(RASTERIZED_MASK_FILE) as src_mask:

        assert src_stack.count >= 1, "스택 파일에 밴드가 없습니다."

        width = src_stack.width
        height = src_stack.height

        print(f"이미지 분할 시작 (Size: {width}x{height}, Stride: {STRIDE})...")

        for col in range(0, width - CHIP_SIZE + 1, STRIDE):
            for row in range(0, height - CHIP_SIZE + 1, STRIDE):

                window = Window(col, row, CHIP_SIZE, CHIP_SIZE)

                # [ ★★★ 수정된 부분 ★★★ ]
                # 1. 마스크 칩 읽기
                mask_chip = src_mask.read(1, window=window).astype(np.float32)

                # 2. (중요) 마스크를 0 또는 1로 강제 이진화(Binarization)
                # 0.5 이상이면 1(건물), 아니면 0(배경)
                mask_chip_binary = (mask_chip > 0.5).astype(np.uint8) # 0, 1 값만 갖도록
                # [ ★★★ 수정 끝 ★★★ ]

                if np.sum(mask_chip_binary) < (CHIP_SIZE * CHIP_SIZE * 0.05):
                    continue

                chip_a = src_stack.read(1, window=window)

                chip_meta = src_stack.meta.copy()
                chip_transform = src_stack.window_transform(window)
                chip_meta.update({
                    "height": chip_a.shape[0], "width": chip_a.shape[1],
                    "transform": chip_transform, "count": 1
                })

                mask_meta = src_mask.meta.copy()
                mask_meta.update({
                    "height": mask_chip_binary.shape[0],
                    "width": mask_chip_binary.shape[1],
                    "transform": chip_transform,
                    "count": 1,
                    "dtype": 'uint8' # [추가] 0/1 정수형으로 저장
                })

                chip_filename = f"chip_{chip_count}.tif"

                with rasterio.open(os.path.join(DIR_MOD_A, chip_filename), 'w', **chip_meta) as dst_a:
                    dst_a.write(chip_a, 1)

                # [수정] 이진화된 마스크를 저장
                with rasterio.open(os.path.join(DIR_MASK, chip_filename), 'w', **mask_meta) as dst_mask:
                    dst_mask.write(mask_chip_binary, 1)

                chip_count += 1

print("---" * 10)
print(f"✅ 작업 완료! 총 {chip_count}개의 모듈 A 칩을 생성했습니다.")