goheung/app/AI_modules/water_body_segmentation/model.py
2026-02-02 19:07:53 +09:00

119 lines
3.5 KiB
Python

# DeepLabV3+ Model
import tensorflow as tf
from keras.layers import Input
from keras.layers import Conv2D
from keras.layers import Concatenate
from keras.layers import UpSampling2D
from keras.models import Model
from tensorflow.keras.applications import ResNet50
from .config import IMAGE_SIZE, LEARNING_RATE, LLF_FILTERS, RESNET_WEIGHTS
from .layers import ConvBlock, AtrousSpatialPyramidPooling
def create_model(image_size=IMAGE_SIZE, learning_rate=LEARNING_RATE,
compile_model=True, in_channels=4):
"""
DeepLabV3+ 모델 생성
Args:
image_size: 입력 이미지 크기
learning_rate: 학습률
compile_model: True면 모델 컴파일 포함
in_channels: 입력 채널 수 (4 = R,G,B,MNDWI)
Returns:
DeepLabV3+ Keras Model
"""
# Input Layer
InputL = Input(shape=(image_size, image_size, in_channels), name="InputLayer")
# 4채널 → 3채널 변환 (ResNet50은 3채널만 지원)
if in_channels != 3:
x = Conv2D(3, kernel_size=1, padding='same', name="ChannelReduce")(InputL)
else:
x = InputL
# Backbone: ResNet50 (pretrained on ImageNet)
# 별도 input_shape로 생성 후 중간 레이어 출력을 추출
resnet50_base = ResNet50(
include_top=False,
weights=RESNET_WEIGHTS,
input_shape=(image_size, image_size, 3)
)
# ResNet50의 중간 출력을 가져오기 위한 멀티출력 모델 생성
resnet50 = Model(
inputs=resnet50_base.input,
outputs={
'deep': resnet50_base.get_layer('conv4_block6_2_relu').output,
'low': resnet50_base.get_layer('conv2_block3_2_relu').output,
},
name="ResNet50-Backbone"
)
features_dict = resnet50(x)
# ASPP Phase - Deep CNN features
DCNN = features_dict['deep']
ASPP = AtrousSpatialPyramidPooling(DCNN)
ASPP = UpSampling2D(
size=(image_size // 4 // ASPP.shape[1], image_size // 4 // ASPP.shape[2]),
name="AtrousSpatial"
)(ASPP)
# Low-Level Features (LLF) Phase
LLF = features_dict['low']
LLF = ConvBlock(filters=LLF_FILTERS, kernel_size=1, name="LLF-ConvBlock")(LLF)
# Decoder: Combine ASPP and LLF
combined = Concatenate(axis=-1, name="Combine-LLF-ASPP")([ASPP, LLF])
features = ConvBlock(name="Top-ConvBlock-1")(combined)
features = ConvBlock(name="Top-ConvBlock-2")(features)
# Upsample to original size
upsample = UpSampling2D(
size=(image_size // features.shape[1], image_size // features.shape[1]),
interpolation='bilinear',
name="Top-UpSample"
)(features)
# Output Mask (1 channel for binary segmentation)
PredMask = Conv2D(
1,
kernel_size=3,
strides=1,
padding='same',
activation='sigmoid',
use_bias=False,
name="OutputMask"
)(upsample)
# Build Model
model = Model(InputL, PredMask, name="DeepLabV3-Plus")
# Compile if requested
if compile_model:
optimizer = tf.keras.optimizers.Adam(learning_rate=learning_rate)
model.compile(loss='binary_crossentropy', optimizer=optimizer)
return model
def load_model(model_path):
"""
저장된 모델 로드
Args:
model_path: 모델 파일 경로 (.h5)
Returns:
로드된 Keras Model
"""
from keras.models import load_model as keras_load_model
from .layers import ConvBlock
custom_objects = {
'ConvBlock': ConvBlock
}
return keras_load_model(model_path, custom_objects=custom_objects)