CMT: demo.py

코드를 이해하기 위한 지식

mask = cv2.imread(mask_fn)[..., 0] / 255. 에서 [..., 0]에 대한 설명
[..., 0] 구문은 Python의 NumPy 라이브러리나 PyTorch에서 사용되는 고급 인덱싱 기법 중 하나입니다. 이 구문은 다차원 배열에서 특정 차원의 모든 요소를 선택하면서, 동시에 다른 차원에서는 특정 인덱스의 요소만을 선택하고자 할 때 사용됩니다.
예를 들어, cv2.imread로 이미지를 읽었을 때 반환되는 배열의 형태는 일반적으로 (높이, 너비, 채널)입니다. 여기서 채널은 BGR 순서로 색상을 나타냅니다(OpenCV는 기본적으로 BGR 포맷을 사용). [..., 0]는 이 배열에서 모든 높이와 너비에 대해 첫 번째 채널(B 채널)만을 선택하라는 의미입니다. ...는 "모든 다른 차원"을 의미하는 Ellipsis 객체로, 여기서는 모든 높이와 너비를 대상으로 합니다.
그러나 마스크 이미지의 경우, 일반적으로 색상 채널이 하나뿐인 그레이스케일 이미지를 사용합니다. 이 경우에도 [..., 0] 구문을 사용할 수 있지만, 실제로는 단일 채널만 존재하기 때문에 [..., 0]이 선택하는 것은 그레이스케일 값 자체가 됩니다. 즉, (높이, 너비, 1) 형태의 배열에서 첫 번째 채널(그리고 유일한 채널)을 선택하여 (높이, 너비) 형태의 2D 배열로 만듭니다.

gt = torch.Tensor(gt_)[None].permute(0, 3, 1, 2).to(device, dtype=torch.float32)에서 [None]을 통한 차원 추가
[None] 구문은 NumPy나 PyTorch에서 새로운 축(차원)을 배열에 추가하기 위해 사용됩니다. 이 방법은 배열의 형태를 변경할 때 주로 사용되며, 추가된 차원에는 특정 "값"이 할당되는 것이 아니라 배열의 구조가 변경됩니다.
예를 들어, (높이, 너비, 채널) 형태의 3D 배열이 있을 때, 이 배열 앞에 [None]을 사용하면 배열의 형태가 (1, 높이, 너비, 채널)로 변경됩니다. 여기서 1은 새로 추가된 차원의 크기를 나타내며, 이 차원을 통해 예를 들어 배치 처리를 위한 차원을 추가할 수 있습니다. 즉, 원래의 데이터는 변경되지 않고, 배열의 구조만이 확장됩니다.
이러한 방식은 특히 딥러닝에서 모델에 데이터를 입력할 때 유용하게 사용됩니다. 많은 딥러닝 프레임워크는 데이터를 (배치 크기, 채널, 높이, 너비) 형태로 요구하는데, 단일 이미지를 처리할 때도 이 형태를 유지하기 위해 [None]을 사용하여 배치 차원을 추가할 수 있습니다.

import warnings # 경고 메시지를 제어하는 라이브러리
from utils import * # 모델 파라미터 불러오는 기능
from tqdm import tqdm # 진행 상황을 표시해주는 모듈
import argparse, os, cv2, glob # glob: 파일 경로명의 패턴 매칭에 사용되는 모듈
from network.network_pro import Inpaint # 인페인트 모듈 임포트

# 아규먼트 받기
warnings.filterwarnings('ignore') # 모든 경고메시지 출력 방지
parser = argparse.ArgumentParser(description="Official Pytorch Code for K. Ko and C.-S. Kim, Continuously Masked Transformer for Image Inpainting, ICCV 2023", usage='use "%(prog)s --help" for more information', formatter_class=argparse.RawTextHelpFormatter)
parser.add_argument('--ckpt', required=True, help='Path for the pretrained model')
parser.add_argument('--img_path', default="./samples/test_img", help='''Path for directory of images. Please note that the file name should be same with that of its corresponding mask''')
parser.add_argument('--mask_path', default="./samples/test_mask", help='''Path for directory of masks.''')
parser.add_argument('--output_path', default="./samples/results", help='Path for saving inpainted images')
args = parser.parse_args()

# 디렉토리 확인
assert os.path.exists(args.img_path), "Please check image path"
assert os.path.exists(args.mask_path), "Please check mask path"
if not os.path.exists(args.output_path):
    os.mkdir(args.output_path)

# Hyper parameters
device = torch.device('cuda')

# Pro
proposed = Inpaint() # 인페인팅 모델 초기화
proposed = load_checkpoint(args.ckpt, proposed) # 모델 파라미터 불러오기
proposed.eval().to(device) # eval로 만들고 GPU에 올리기
maskfn = glob.glob(os.path.join(args.mask_path, '*.*')) # 마스크 폴더에 있는 모든 파일 주소를 리스트로 저장
prog_bar = tqdm(maskfn) # tqdm은 순회할 수 있는 데이터를 받아 이를 시각화할 수 있는 모듈을 만드는 객체이다.
avg = 0. # 평균 PSNR 값을 계산하기 위한 변수를 초기화

for step, mask_fn in enumerate(prog_bar): # 마스크 파일 목록을 순회하며, 각 마스크 파일에 대해 인페인팅 작업을 수행. enumerate는 반복 가능한 데이터를 받아 번호와 함께 반환하는 함수
    fn = os.path.basename(mask_fn) # 현재 처리 중인 마스크 파일의 파일명을 추출
    
    # 이미지 불러오기
    gt_ = (cv2.imread(os.path.join(args.img_path, fn)) / 255.) * 2 - 1. # 대응하는 원본 이미지를 읽어들이고, 픽셀 값을 [-1, 1] 범위로 정규화 (일반적으로 신경망 모델이 -1과 1 사이의 입력 값에 대해 더 잘 학습하기 때문)
    mask = cv2.imread(mask_fn)[..., 0] / 255. # 마스크 이미지를 읽어들이고, 픽셀 값을 [0, 1] 범위로 정규화, 마스크는 그레이스케일 이미지로 가정, 마스크는 보통 이진 이미지로, 픽셀이 마스크되었는지 여부만 표시하기 때문에 [0, 1] 범위가 적합
    # cv2.imread() 함수로 흑백 데이터를 읽어오면 [높이, 너비]의 2차원을 가지는 넘파이 어레이로 받아온다.
    # gt_ = (256, 256, 3), mask = (256, 256)
    
    
    # 텐서 변환
    gt = torch.Tensor(gt_)[None].permute(0, 3, 1, 2).to(device, dtype=torch.float32) # 원본 이미지를 PyTorch 텐서로 변환하고, 배치 차원을 추가한 뒤, 채널 차원을 앞으로 이동시킨 후, GPU로 이동
    mask = torch.Tensor(mask)[None, None].to(device, dtype=torch.float32) # 마스크를 PyTorch 텐서로 변환하고, 배치 차원과 채널 차원을 추가한 후, GPU로 이동

    with torch.no_grad():
        out_pro = proposed(gt, mask) # 자동 미분을 비활성화한 상태에서 모델을 사용하여 인페인팅을 수행

    out_pro = torch.clip(out_pro, -1., 1.)*0.5 + 0.5 # 모델 입력을 [-1, 1]로 했기 때문에, 결과값도 [-1, 1]이다. 그러나 이미지를 시각화 하기 위해 이를 다시 [0, 1] 범위로 재정규화하고 아래처럼 255를 곱한다.
    out_pro = out_pro[0].permute(1, 2, 0).cpu().detach().numpy() * 255. # 출력 이미지를 CPU로 이동시키고, 채널 차원을 맨 뒤로 이동시킨 후, 넘파이 배열로 변환하고 [0, 255] 범위로 스케일링.
    score = psnr(out_pro, (gt_ * 0.5 + 0.5)*255.) # 출력 이미지와 원본 이미지 간의 PSNR 점수를 계산
    save_path_ = os.path.join(args.output_path, '{}').format(fn) # 출력 이미지를 저장할 경로를 설정
    cv2.imwrite(save_path_, out_pro) # 인페인팅된 이미지를 지정된 경로에 저장
    avg += score # 현재 이미지의 PSNR 점수를 평균 점수에 더함
    prog_bar.set_description("PSNR {}".format(avg / (step + 1))) # 프로그레스 바의 설명을 업데이트하여 현재까지의 평균 PSNR 값을 표시