132 lines
4.8 KiB
Python
132 lines
4.8 KiB
Python
import os, mmcv, cv2, json
|
|
import numpy as np
|
|
import torch
|
|
from pathlib import Path
|
|
from PIL import Image
|
|
from tqdm import tqdm
|
|
from mmseg.apis import init_segmentor, inference_segmentor
|
|
from shapely.geometry import Polygon, mapping
|
|
import sys
|
|
|
|
_DETECT_DIR = Path(__file__).parent # mx15hdi/Detect/
|
|
_MX15HDI_DIR = _DETECT_DIR.parent # mx15hdi/
|
|
|
|
|
|
def load_model():
|
|
"""서버 시작 시 1회 호출. 로드된 모델 객체를 반환한다."""
|
|
# 우선순위: 환경변수 DEVICE > GPU 자동감지 > CPU 폴백
|
|
env_device = os.environ.get('DEVICE', '').strip()
|
|
if env_device:
|
|
device = env_device
|
|
elif torch.cuda.is_available():
|
|
device = 'cuda:0'
|
|
else:
|
|
device = 'cpu'
|
|
print(f'[Inference] 사용 device: {device}')
|
|
|
|
config = str(_DETECT_DIR / 'V7_SPECIAL.py')
|
|
checkpoint = str(_DETECT_DIR / 'epoch_165.pth')
|
|
model = init_segmentor(config, checkpoint, device=device)
|
|
model.PALETTE = [
|
|
[0, 0, 0], # background
|
|
[0, 0, 204], # black
|
|
[180, 180, 180], # brown
|
|
[255, 255, 0], # rainbow
|
|
[178, 102, 255] # silver
|
|
]
|
|
return model
|
|
|
|
|
|
def blend_images(original_img, color_mask, alpha=0.6):
|
|
"""
|
|
Blend original image and color mask with alpha transparency.
|
|
Inputs: numpy arrays HWC uint8
|
|
"""
|
|
blended = cv2.addWeighted(original_img, 1 - alpha, color_mask, alpha, 0)
|
|
return blended
|
|
|
|
|
|
def run_inference(model, file_id: str, write_files: bool = False) -> dict:
|
|
"""
|
|
사전 로드된 모델로 file_id 폴더 내 이미지를 세그멘테이션한다.
|
|
|
|
Args:
|
|
model: load_model()로 로드된 모델 객체.
|
|
file_id: 처리할 이미지 폴더명.
|
|
write_files: True이면 Detect/result/ 와 Detect/Mask_result/ 에 중간 파일 저장.
|
|
False이면 디스크 쓰기 생략 (기본값).
|
|
|
|
Returns:
|
|
inference_cache: {image_filename: {'blended': ndarray, 'mask': ndarray, 'ext': str}}
|
|
이 값을 run_georeference()에 전달하면 중간 파일 읽기를 생략할 수 있다.
|
|
"""
|
|
img_path = str(_MX15HDI_DIR / 'Metadata' / 'Image' / 'Original_Images' / file_id)
|
|
output_folder = str(_DETECT_DIR / 'result' / file_id)
|
|
mask_folder = str(_DETECT_DIR / 'Mask_result' / file_id)
|
|
|
|
if not os.path.exists(img_path):
|
|
raise FileNotFoundError(f"이미지 폴더가 존재하지 않습니다: {img_path}")
|
|
|
|
if write_files:
|
|
os.makedirs(output_folder, exist_ok=True)
|
|
os.makedirs(mask_folder, exist_ok=True)
|
|
|
|
image_files = [
|
|
f for f in os.listdir(img_path)
|
|
if f.lower().endswith(('.jpg', '.jpeg', '.png', '.bmp', '.tif', '.tiff'))
|
|
]
|
|
|
|
# palette_array는 이미지마다 동일하므로 루프 외부에서 1회 생성
|
|
palette_array = np.array(model.PALETTE, dtype=np.uint8)
|
|
|
|
inference_cache = {}
|
|
|
|
for image_file in tqdm(image_files, desc="Processing images"):
|
|
image_path = os.path.join(img_path, image_file)
|
|
image_name, image_ext = os.path.splitext(image_file)
|
|
image_ext = image_ext.lower()
|
|
|
|
# 이미지를 1회만 읽어 inference와 blending 모두에 재사용
|
|
img_bgr = cv2.imread(image_path)
|
|
img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
|
|
|
|
# 이미 로드된 배열을 inference_segmentor에 전달 (경로 전달 시 내부에서 재읽기 발생)
|
|
result = inference_segmentor(model, img_bgr)
|
|
seg_map = result[0]
|
|
|
|
# Create color mask from palette
|
|
color_mask = palette_array[seg_map]
|
|
|
|
# blended image
|
|
blended = blend_images(img_rgb, color_mask, alpha=0.6)
|
|
blended_bgr = cv2.cvtColor(blended, cv2.COLOR_RGB2BGR)
|
|
|
|
# mask — numpy 슬라이싱으로 cv2.cvtColor 호출 1회 제거
|
|
mask_bgr = color_mask[:, :, ::-1].copy()
|
|
|
|
# 결과를 메모리 캐시에 저장 (georeference 단계에서 재사용)
|
|
# mask는 palette 원본(RGB) 그대로 저장 — Oilshape의 class_colors가 RGB 기준이므로 BGR로 저장 시 색상 매칭 실패
|
|
inference_cache[image_file] = {
|
|
'blended': blended_bgr,
|
|
'mask': color_mask,
|
|
'ext': image_ext,
|
|
}
|
|
|
|
if write_files:
|
|
cv2.imwrite(
|
|
os.path.join(output_folder, f"{image_name}{image_ext}"),
|
|
blended_bgr,
|
|
[cv2.IMWRITE_JPEG_QUALITY, 85]
|
|
)
|
|
cv2.imwrite(os.path.join(mask_folder, f"{image_name}{image_ext}"), mask_bgr)
|
|
|
|
return inference_cache
|
|
|
|
|
|
if __name__ == '__main__':
|
|
if len(sys.argv) < 2:
|
|
raise ValueError("파라미터가 제공되지 않았습니다. 폴더 이름을 명령줄 인자로 입력해주세요.")
|
|
_model = load_model()
|
|
# CLI 단독 실행 시에는 중간 파일도 디스크에 저장
|
|
run_inference(_model, sys.argv[1], write_files=True)
|