이미지 복원 대회 실험을 통해 본 AI 기반 reconstruction 문제

 

손상 영역 복원과 색상화 베이스라인 구축 경험을 ultra-high-resolution fMRI reconstruction 관점에서 해석하며

최근 참여한 데이콘 이미지 재복원 대회에서는 흑백이며 일부 영역이 손상된 이미지를 입력으로 받아, 원본에 가까운 컬러 이미지를 복원하는 모델을 설계하였다.

 

표면적으로는 일반적인 비전 대회처럼 보일 수 있지만, 실제로 문제를 분석해보면 이는 단순 색상화(colorization)를 넘어 불완전하고 저품질인 관측값으로부터 구조와 시각적 정보를 동시에 복원하는 inverse problem에 가깝다. 이러한 관점은 내가 관심을 갖고 지원했던 Forschungszentrum Jülich의 ultra-high-resolution fMRI reconstruction 프로젝트와도 기술적으로 중요한 접점을 가진다고 느꼈다.

 

Jülich의 프로젝트는 7T MRI 기반의 초고해상도 fMRI에서 발생하는 낮은 SNR, 구조 왜곡, 공간적 fidelity 저하 문제를 해결하기 위해 deep learning 기반 reconstruction 및 super-resolution 기법을 적용하는 것을 목표로 한다. 내가 이번 대회에서 직접 다룬 데이터는 자연영상이었지만, 문제를 접근하는 방식 자체는 의료영상 reconstruction 연구에서 중요한 여러 요소와 닮아 있었다. 특히 손실된 정보의 복원, 국소 영역 복원의 중요성, 정량 지표 기반의 성능 평가, 구조적 유사성 보존을 고려한 모델 설계라는 점에서 좋은 예행연습이 되었다.

---

1. 문제를 reconstruction 관점에서 재해석하기

이번 대회의 입력은 단순한 grayscale 이미지가 아니었다. 입력 이미지는

• 흑백으로 변환되어 있고
• 일부 영역은 검은 원, 사각형, 다각형 등으로 손상되어 있으며
• 정답은 원래의 컬러 이미지

형태였다. 즉 모델은 두 가지 문제를 동시에 풀어야 했다.

 

첫째, 관측되지 않은 손실 영역을 복원해야 한다.
둘째, grayscale로 축소된 정보를 바탕으로 가능한 원본에 가까운 색 분포를 재구성해야 한다.

이 문제를 reconstruction 관점에서 보면, 이는 불완전한 관측 x 로부터 원신호에 해당하는 y 를 추정하는 문제로 볼 수 있다. 물론 자연영상과 fMRI는 도메인이 다르지만, 원신호를 직접 관측할 수 없고 제한된 입력으로부터 더 풍부한 표현을 복원해야 한다는 점에서는 공통점이 있다.

 

Jülich 프로젝트 역시 초고해상도 laminar fMRI에서 낮은 SNR과 reconstruction artefact를 극복하면서, anatomical reference와 잘 정합되는 구조적으로 정밀한 이미지를 얻는 것이 핵심이다. 즉 내가 대회에서 다룬 문제는 비록 자연영상 기반이지만, information loss가 존재하는 입력으로부터 구조와 품질을 복원하는 학습형 reconstruction pipeline을 구축해보는 경험이었다고 정리할 수 있다.

---

2. 데이터 탐색을 통해 확인한 핵심: local defect가 중요한 문제

초기에는 데이터 구조를 먼저 점검하였다. 학습 이미지와 정답 이미지는 pair로 제공되었고, 입력은 대부분 512×512 grayscale 이미지였다. 직접 시각화해본 결과, 손상 영역은 랜덤 노이즈가 아니라 비교적 큰 연결 성분을 이루는 검은 기하학적 도형 형태를 보였다.

이 점은 매우 중요했다. 왜냐하면 이번 문제에서 중요한 것은 단순히 전체 이미지를 보기 좋게 만드는 것이 아니라, 손실 영역이라는 특정 local region을 얼마나 정확히 복원하느냐였기 때문이다. 실제 평가 지표 역시 전체 SSIM뿐 아니라 손실 영역 SSIM을 별도로 반영하고 있었기 때문에, global reconstruction만으로는 충분하지 않았다.

 

이러한 문제 설정은 fMRI reconstruction과도 사고방식 측면에서 연결된다. ultra-high-resolution fMRI에서는 전체 이미지 품질도 중요하지만, 진짜 중요한 것은 작고 미세한 cortical layer 수준의 신호가 구조적으로 왜곡되지 않고 보존되는가이다. 즉, 전체 평균 품질만 높아서는 부족하고, 특정 국소 구조(local structure)의 fidelity를 얼마나 잘 지키는지가 핵심이 된다. 이번 대회에서 손상 영역 복원 성능을 별도로 의식하게 된 것도, 결국 이런 global quality와 local fidelity를 분리해서 보는 관점을 훈련하는 과정이었다.

---

3. baseline 설계: 1채널 입력에서 3채널 출력을 복원하는 U-Net

초기 베이스라인으로는 가장 단순한 encoder-decoder 계열의 U-Net을 사용하였다. 입력은 grayscale 1채널, 출력은 RGB 3채널로 두었고, 우선은 파이프라인이 정상 동작하는지 확인하기 위해 소규모 subset과 낮은 해상도로 빠른 sanity check 실험을 수행하였다.

이 설정은 성능 최적화보다는 다음을 확인하는 데 목적이 있었다.

• 데이터 로딩 및 pair 매칭이 안정적인가
• 입력과 출력 차원이 올바르게 연결되는가
• 학습 loss가 안정적으로 감소하는가
• inference와 submission pipeline이 끝까지 동작하는가

실험 결과, loss는 안정적으로 감소했고 예측 결과도 생성되었다. 즉 기본적인 supervised reconstruction 파이프라인은 정상적으로 구축되었다고 볼 수 있었다. 다만 초기 결과는 매우 흐렸고, 색상은 전체적인 평균값에 수렴하는 경향이 강했다. 이는 저해상도, 적은 학습 데이터, 짧은 epoch, 그리고 L1 loss 중심의 학습이라는 조건에서 충분히 예상 가능한 결과였다.

이 과정은 Jülich 프로젝트에서의 deep learning reconstruction framework를 직접 구현했다고 볼 수는 없지만, 적어도 입력 손실이 있는 데이터를 대상으로 encoder-decoder 기반 복원 모델을 설계하고, 학습-검증-추론-정량 평가 파이프라인을 직접 구성했다는 점에서 관련성이 있다. 특히 의료영상 reconstruction에서도 단순 pixelwise loss만 사용하면 oversmoothing이 발생하기 쉽다는 점을 감안하면, 이번 baseline의 한계를 통해 reconstruction loss 설계의 중요성을 실감할 수 있었다.

---

4. 왜 단순 global loss만으로는 부족했는가

초기 실험을 진행하면서 가장 분명하게 드러난 한계는, 모델이 손상 영역 자체를 충분히 복원하지 못한다는 점이었다. 하늘, 나무, 도로처럼 넓은 배경의 색은 어느 정도 예측하기 시작했지만, 검은 원이나 사각형처럼 손실된 부분은 그대로 남거나 흐려지는 수준에 그쳤다.

 

이 현상은 손실 함수 관점에서 자연스럽다. 전체 이미지에 대해 단순 L1 loss를 주면, 손상 영역은 이미지의 일부 픽셀에 불과하기 때문에 모델은 우선적으로 전체 평균 오차를 줄이는 방향으로 학습된다. 즉 넓은 배경 영역의 색을 대략 맞추는 것이 손실을 줄이는 데 더 유리하고, 상대적으로 면적이 작은 손상 부위를 정확하게 복원하는 동기는 약해진다.

 

이 문제는 ultra-high-resolution fMRI reconstruction에서도 충분히 중요한 시사점을 준다. 예를 들어 laminar-level fMRI에서는 전체적인 이미지가 매끈해 보이는 것보다, 실제로는 cortical layer를 구분할 수 있는 미세 구조와 signal profile을 보존하는 것이 훨씬 더 중요하다. 따라서 reconstruction 모델은 global denoising이나 평균적 fidelity만 높이는 방향이 아니라, 정보 손실이 치명적인 영역을 더 민감하게 복원할 수 있도록 설계되어야 한다. 이번 대회에서 local region에 더 큰 중요도를 두는 loss가 필요하다고 느낀 것도 바로 이런 이유였다.

---

5. pseudo-mask 기반 weighted loss 아이디어

이후에는 입력 grayscale 이미지에서 pseudo-mask를 추정하고, 손실 영역에 더 큰 가중치를 부여하는 방향을 고려하였다. 구체적으로는 아주 어두운 픽셀을 thresholding한 뒤 connected components와 최소 면적 필터를 적용하여, 손상 영역에 가까운 pseudo-mask를 구성하는 접근이었다.

 

물론 단순 threshold만으로는 그림자나 원래 어두운 물체까지 같이 잡히는 문제가 있었지만, 이 실험은 중요한 교훈을 주었다. reconstruction 문제에서는 정답이 직접 주어지지 않는 경우에도 도메인 priors를 활용해 중요한 영역을 추정하고, 이를 학습 objective에 반영하는 방식이 필요하다는 점이다.

 

이를 바탕으로 이후에는 전체 L1 loss에 더해 pseudo-mask 영역의 오차에 높은 가중치를 부여하는 weighted loss를 설계하고자 하였다.

---

6. 성능 개선 실험: 해상도와 데이터 양을 높였을 때

이후에는 이미지 크기를 256으로 높이고, subset 크기를 늘려 학습을 다시 진행하였다. 그 결과 초기 초경량 baseline에 비해 분명한 개선이 나타났다. 하늘은 더 파랗게, 식생은 더 초록색 계열로 복원되었고, 전체적으로 grayscale에 가까웠던 출력이 보다 자연스러운 색 분포를 가지게 되었다.

 

import os
import zipfile
import random
import numpy as np
import pandas as pd
import cv2
from tqdm import tqdm
from PIL import Image

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader, random_split, Subset


# =========================
# 설정
# =========================
ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
TRAIN_CSV = os.path.join(ROOT, "train.csv")
TEST_CSV = os.path.join(ROOT, "test.csv")
TRAIN_INPUT_DIR = os.path.join(ROOT, "train_input")
TRAIN_GT_DIR = os.path.join(ROOT, "train_gt")
TEST_INPUT_DIR = os.path.join(ROOT, "test_input")

OUTPUT_DIR = os.path.join(ROOT, "outputs")
CKPT_DIR = os.path.join(OUTPUT_DIR, "checkpoints_fast")
SUBMISSION_DIR = os.path.join(OUTPUT_DIR, "submissions_fast")

os.makedirs(OUTPUT_DIR, exist_ok=True)
os.makedirs(CKPT_DIR, exist_ok=True)
os.makedirs(SUBMISSION_DIR, exist_ok=True)

DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")

IMG_SIZE = 256
BATCH_SIZE = 8
EPOCHS = 5
LR = 1e-4
SEED = 42

SUBSET_SIZE = 5000
VAL_SIZE = 500

torch.manual_seed(SEED)
np.random.seed(SEED)
random.seed(SEED)


# =========================
# Dataset
# =========================
class DaconImageDataset(Dataset):
    def __init__(self, csv_path, input_dir, gt_dir=None, img_size=128, is_test=False):
        self.df = pd.read_csv(csv_path)
        self.input_dir = input_dir
        self.gt_dir = gt_dir
        self.img_size = img_size
        self.is_test = is_test

    def __len__(self):
        return len(self.df)

    def _load_input(self, filename):
        path = os.path.join(self.input_dir, filename)
        try:
            img = Image.open(path).convert("L")
        except Exception as e:
            raise FileNotFoundError(f"입력 이미지 로드 실패: {path} | {e}")

        img = img.resize((self.img_size, self.img_size))
        img = np.array(img, dtype=np.float32) / 255.0
        img = np.expand_dims(img, axis=0)  # (1, H, W)
        return torch.tensor(img, dtype=torch.float32)

    def _load_gt(self, filename):
        path = os.path.join(self.gt_dir, filename)
        try:
            img = Image.open(path).convert("RGB")
        except Exception as e:
            raise FileNotFoundError(f"GT 이미지 로드 실패: {path} | {e}")

        img = img.resize((self.img_size, self.img_size))
        img = np.array(img, dtype=np.float32) / 255.0
        img = np.transpose(img, (2, 0, 1))  # (3, H, W)
        return torch.tensor(img, dtype=torch.float32)

    def __getitem__(self, idx):
        input_rel = self.df.iloc[idx]["input_image_path"]
        input_name = os.path.basename(input_rel)
        x = self._load_input(input_name)

        if self.is_test:
            return x, input_name

        gt_rel = self.df.iloc[idx]["gt_image_path"]
        gt_name = os.path.basename(gt_rel)
        y = self._load_gt(gt_name)
        return x, y


# =========================
# 가벼운 U-Net
# =========================
class DoubleConv(nn.Module):
    def __init__(self, in_ch, out_ch):
        super().__init__()
        self.block = nn.Sequential(
            nn.Conv2d(in_ch, out_ch, 3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(out_ch, out_ch, 3, padding=1),
            nn.ReLU(inplace=True),
        )

    def forward(self, x):
        return self.block(x)


class SmallUNet(nn.Module):
    def __init__(self, in_channels=1, out_channels=3, base_ch=32):
        super().__init__()

        self.enc1 = DoubleConv(in_channels, base_ch)          # 32
        self.enc2 = DoubleConv(base_ch, base_ch * 2)          # 64
        self.enc3 = DoubleConv(base_ch * 2, base_ch * 4)      # 128
        self.enc4 = DoubleConv(base_ch * 4, base_ch * 8)      # 256

        self.pool = nn.MaxPool2d(2)
        self.bottleneck = DoubleConv(base_ch * 8, base_ch * 16)  # 512

        self.up4 = nn.ConvTranspose2d(base_ch * 16, base_ch * 8, kernel_size=2, stride=2)
        self.dec4 = DoubleConv(base_ch * 16, base_ch * 8)

        self.up3 = nn.ConvTranspose2d(base_ch * 8, base_ch * 4, kernel_size=2, stride=2)
        self.dec3 = DoubleConv(base_ch * 8, base_ch * 4)

        self.up2 = nn.ConvTranspose2d(base_ch * 4, base_ch * 2, kernel_size=2, stride=2)
        self.dec2 = DoubleConv(base_ch * 4, base_ch * 2)

        self.up1 = nn.ConvTranspose2d(base_ch * 2, base_ch, kernel_size=2, stride=2)
        self.dec1 = DoubleConv(base_ch * 2, base_ch)

        self.final = nn.Conv2d(base_ch, out_channels, kernel_size=1)

    def forward(self, x):
        e1 = self.enc1(x)
        e2 = self.enc2(self.pool(e1))
        e3 = self.enc3(self.pool(e2))
        e4 = self.enc4(self.pool(e3))

        b = self.bottleneck(self.pool(e4))

        d4 = self.up4(b)
        d4 = torch.cat([d4, e4], dim=1)
        d4 = self.dec4(d4)

        d3 = self.up3(d4)
        d3 = torch.cat([d3, e3], dim=1)
        d3 = self.dec3(d3)

        d2 = self.up2(d3)
        d2 = torch.cat([d2, e2], dim=1)
        d2 = self.dec2(d2)

        d1 = self.up1(d2)
        d1 = torch.cat([d1, e1], dim=1)
        d1 = self.dec1(d1)

        return torch.sigmoid(self.final(d1))


# =========================
# Train / Valid
# =========================
def train_one_epoch(model, loader, criterion, optimizer):
    model.train()
    total_loss = 0.0

    pbar = tqdm(loader, desc="Train", leave=False)
    for x, y in pbar:
        x = x.to(DEVICE, non_blocking=True)
        y = y.to(DEVICE, non_blocking=True)

        optimizer.zero_grad()
        pred = model(x)
        loss = criterion(pred, y)
        loss.backward()
        optimizer.step()

        total_loss += loss.item()
        pbar.set_postfix(loss=f"{loss.item():.4f}")

    return total_loss / len(loader)


@torch.no_grad()
def valid_one_epoch(model, loader, criterion):
    model.eval()
    total_loss = 0.0

    pbar = tqdm(loader, desc="Valid", leave=False)
    for x, y in pbar:
        x = x.to(DEVICE, non_blocking=True)
        y = y.to(DEVICE, non_blocking=True)

        pred = model(x)
        loss = criterion(pred, y)
        total_loss += loss.item()

    return total_loss / len(loader)


# =========================
# Inference + zip
# =========================
@torch.no_grad()
def inference_and_zip(model, test_loader, save_dir, zip_path):
    model.eval()
    os.makedirs(save_dir, exist_ok=True)

    for x, file_name in tqdm(test_loader, desc="Inference"):
        x = x.to(DEVICE, non_blocking=True)

        pred = model(x)                      # (1, 3, H, W)
        pred = pred.squeeze(0).cpu().numpy()
        pred = np.transpose(pred, (1, 2, 0))  # HWC
        pred = (pred * 255.0).clip(0, 255).astype(np.uint8)

        pred_bgr = cv2.cvtColor(pred, cv2.COLOR_RGB2BGR)
        cv2.imwrite(os.path.join(save_dir, file_name[0]), pred_bgr)

    with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zipf:
        for fname in sorted(os.listdir(save_dir)):
            if fname.lower().endswith(".png"):
                zipf.write(os.path.join(save_dir, fname), arcname=fname)


# =========================
# Main
# =========================
def main():
    print(f"DEVICE: {DEVICE}")
    print(f"ROOT: {ROOT}")

    full_dataset = DaconImageDataset(
        csv_path=TRAIN_CSV,
        input_dir=TRAIN_INPUT_DIR,
        gt_dir=TRAIN_GT_DIR,
        img_size=IMG_SIZE,
        is_test=False
    )

    # 일부만 사용
    subset_indices = list(range(min(SUBSET_SIZE, len(full_dataset))))
    full_dataset = Subset(full_dataset, subset_indices)

    train_size = len(full_dataset) - VAL_SIZE
    val_size = VAL_SIZE

    train_dataset, val_dataset = random_split(
        full_dataset,
        [train_size, val_size],
        generator=torch.Generator().manual_seed(SEED)
    )

    train_loader = DataLoader(
        train_dataset,
        batch_size=BATCH_SIZE,
        shuffle=True,
        num_workers=0,
        pin_memory=True
    )

    val_loader = DataLoader(
        val_dataset,
        batch_size=BATCH_SIZE,
        shuffle=False,
        num_workers=0,
        pin_memory=True
    )

    test_dataset = DaconImageDataset(
        csv_path=TEST_CSV,
        input_dir=TEST_INPUT_DIR,
        gt_dir=None,
        img_size=IMG_SIZE,
        is_test=True
    )

    test_loader = DataLoader(
        test_dataset,
        batch_size=1,
        shuffle=False,
        num_workers=0,
        pin_memory=True
    )

    model = SmallUNet(in_channels=1, out_channels=3, base_ch=32).to(DEVICE)

    criterion = nn.L1Loss()
    optimizer = optim.AdamW(model.parameters(), lr=LR)

    best_val_loss = float("inf")

    for epoch in range(EPOCHS):
        print(f"\n===== Epoch {epoch+1}/{EPOCHS} =====")

        train_loss = train_one_epoch(model, train_loader, criterion, optimizer)
        val_loss = valid_one_epoch(model, val_loader, criterion)

        print(f"train_loss: {train_loss:.4f}")
        print(f"val_loss  : {val_loss:.4f}")

        ckpt_path = os.path.join(CKPT_DIR, f"epoch_{epoch+1}.pth")
        torch.save({
            "epoch": epoch + 1,
            "model_state_dict": model.state_dict(),
            "optimizer_state_dict": optimizer.state_dict(),
            "val_loss": val_loss
        }, ckpt_path)

        if val_loss < best_val_loss:
            best_val_loss = val_loss
            best_path = os.path.join(CKPT_DIR, "best_model.pth")
            torch.save(model.state_dict(), best_path)
            print(f"Best model saved: {best_path}")

    best_path = os.path.join(CKPT_DIR, "best_model.pth")
    model.load_state_dict(torch.load(best_path, map_location=DEVICE))

    pred_dir = os.path.join(SUBMISSION_DIR, "best_pred")
    zip_path = os.path.join(SUBMISSION_DIR, "submission_fast.zip")

    inference_and_zip(model, test_loader, pred_dir, zip_path)
    print(f"\n제출 zip 생성 완료: {zip_path}")


if __name__ == "__main__":
    main()

하지만 여전히 손상 영역 복원은 충분하지 않았다. 검은 원이나 사각형은 약간 흐려질 뿐 실제 배경 구조로 자연스럽게 메워지지 않는 경우가 많았다. 이 결과는 매우 중요했다. 왜냐하면 단순한 colorization 성능 향상만으로는 대회 점수의 핵심인 손상 영역 복원 성능을 확보할 수 없음을 보여주었기 때문이다.

 

이 경험은 Jülich 프로젝트의 목적과도 연결된다. 그 프로젝트는 단순히 이미지를 보기 좋게 만드는 것이 아니라, ultra-high-resolution fMRI에서 SNR과 structural fidelity를 함께 높이고, 더 나아가 super-resolution까지 적용하여 laminar-level response characterization을 개선하는 것을 목표로 한다. 내가 이번 실험에서 느낀 것은, reconstruction 문제에서 단순 global improvement만으로는 충분하지 않으며, 최종 목적이 되는 구조적·기능적 정밀도를 확보하려면 그에 맞는 loss와 evaluation을 별도로 설계해야 한다는 점이었다.

---

7. mri재복원 프로젝트와의 연결

이번 대회 실험이 곧바로 fMRI reconstruction의 직접적인 경험을 의미하지는 않는다. MRI physics, k-space, EPI/EPIK sequence, SPM/FSL/ANTs 기반 neuroimaging workflow 등은 별도의 전문적 학습이 필요한 영역이다. 그러나 내가 이번 실험을 통해 분명히 체득한 것은, 불완전한 입력으로부터 구조를 보존하면서 정보를 복원하는 deep learning pipeline을 설계하고, 그 한계를 정량 지표와 시각화를 통해 분석하는 방식이다.

 

 

참고자료 

https://dacon.io/competitions/official/236420/overview/description

이미지 색상화 및 손실 부분 복원 AI 경진대회 - DACON

https://velog.io/@lighthouse97/UNet%EC%9D%98-%EC%9D%B4%ED%95%B4

UNet의 이해

 

 

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, Dataset
import cv2
import numpy as np
import os
import zipfile
from skimage.metrics import structural_similarity as ssim
from tqdm import tqdm

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

class UNet(nn.Module):
    def __init__(self):
        super(UNet, self).__init__()
        self.enc1 = self.conv_block(3, 64)
        self.enc2 = self.conv_block(64, 128)
        self.enc3 = self.conv_block(128, 256)
        self.enc4 = self.conv_block(256, 512)
        self.enc5 = self.conv_block(512, 1024)

        self.up1 = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
        self.dec1 = self.conv_block(1024 + 512, 512)

        self.up2 = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
        self.dec2 = self.conv_block(512 + 256, 256)

        self.up3 = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
        self.dec3 = self.conv_block(256 + 128, 128)

        self.up4 = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
        self.dec4 = self.conv_block(128 + 64, 64)

        self.final = nn.Conv2d(64, 3, kernel_size=1)

    def conv_block(self, in_channels, out_channels):
        return nn.Sequential(
            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),
            nn.ReLU(inplace=True)
        )

    def forward(self, x):
        e1 = self.enc1(x)
        e2 = self.enc2(nn.MaxPool2d(2)(e1))
        e3 = self.enc3(nn.MaxPool2d(2)(e2))
        e4 = self.enc4(nn.MaxPool2d(2)(e3))
        e5 = self.enc5(nn.MaxPool2d(2)(e4))

        d1 = self.dec1(torch.cat([self.up1(e5), e4], dim=1))
        d2 = self.dec2(torch.cat([self.up2(d1), e3], dim=1))
        d3 = self.dec3(torch.cat([self.up3(d2), e2], dim=1))
        d4 = self.dec4(torch.cat([self.up4(d3), e1], dim=1))

        return torch.sigmoid(self.final(d4))

class PatchGANDiscriminator(nn.Module):
    def __init__(self, in_channels=3):
        super(PatchGANDiscriminator, self).__init__()
        self.model = nn.Sequential(
            nn.Conv2d(in_channels, 64, kernel_size=4, stride=2, padding=1),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(64, 128, kernel_size=4, stride=2, padding=1),
            nn.BatchNorm2d(128),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(128, 256, kernel_size=4, stride=2, padding=1),
            nn.BatchNorm2d(256),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(256, 512, kernel_size=4, stride=1, padding=1),
            nn.BatchNorm2d(512),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(512, 1, kernel_size=4, stride=1, padding=1),
            nn.Sigmoid()
        )

    def forward(self, x):
        return self.model(x)

class ImageDataset(Dataset):
    def __init__(self, input_dir, gt_dir, transform=None):
        self.input_dir = input_dir
        self.gt_dir = gt_dir
        self.input_images = sorted(os.listdir(input_dir))
        self.gt_images = sorted(os.listdir(gt_dir))
        self.transform = transform

    def __len__(self):
        return len(self.input_images)

    def __getitem__(self, idx):
        input_path = os.path.join(self.input_dir, self.input_images[idx])
        gt_path = os.path.join(self.gt_dir, self.gt_images[idx])
        input_image = cv2.imread(input_path)
        gt_image = cv2.imread(gt_path)
        if self.transform:
            input_image = self.transform(input_image)
            gt_image = self.transform(gt_image)
        return (
            torch.tensor(input_image).permute(2, 0, 1).float() / 255.0,
            torch.tensor(gt_image).permute(2, 0, 1).float() / 255.0
        )

generator = UNet().to(device)
discriminator = PatchGANDiscriminator().to(device)

adversarial_loss = nn.BCELoss()  
pixel_loss = nn.MSELoss()  

optimizer_G = optim.AdamW(generator.parameters(), lr=0.0001, betas=(0.5, 0.999))
optimizer_D = optim.AdamW(discriminator.parameters(), lr=0.0001, betas=(0.5, 0.999))

train_dataset = ImageDataset("/home/work/Dacon_Dataset/train_input", "/home/work/Dacon_Dataset/train_gt")
train_loader = DataLoader(train_dataset, batch_size=8, shuffle=True, num_workers=1, pin_memory=True)

epochs = 50
result_dir = "/home/work/Dacon_Dataset/result"
os.makedirs(result_dir, exist_ok=True)
checkpoint_path = "/home/work/Dacon_Dataset/checkpoint.pth"

for epoch in range(epochs):
    generator.train()
    discriminator.train()
    running_loss_G = 0.0
    running_loss_D = 0.0

    with tqdm(total=len(train_loader), desc=f"Epoch {epoch+1}/{epochs}", unit="batch") as pbar:
        for input_images, gt_images in train_loader:
            input_images, gt_images = input_images.to(device), gt_images.to(device)

            real_labels = torch.ones_like(discriminator(gt_images)).to(device)
            fake_labels = torch.zeros_like(discriminator(input_images)).to(device)

            optimizer_G.zero_grad()
            fake_images = generator(input_images)
            pred_fake = discriminator(fake_images)

            g_loss_adv = adversarial_loss(pred_fake, real_labels)
            g_loss_pixel = pixel_loss(fake_images, gt_images)
            g_loss = g_loss_adv + 100 * g_loss_pixel
            g_loss.backward()
            optimizer_G.step()

            optimizer_D.zero_grad()
            pred_real = discriminator(gt_images)
            loss_real = adversarial_loss(pred_real, real_labels)

            pred_fake = discriminator(fake_images.detach())
            loss_fake = adversarial_loss(pred_fake, fake_labels)

            d_loss = (loss_real + loss_fake) / 2
            d_loss.backward()
            optimizer_D.step()

            running_loss_G += g_loss.item()
            running_loss_D += d_loss.item()

            pbar.set_postfix(generator_loss=g_loss.item(), discriminator_loss=d_loss.item())
            pbar.update(1)

    print(f"Epoch [{epoch+1}/{epochs}] - Generator Loss: {running_loss_G / len(train_loader):.4f}, Discriminator Loss: {running_loss_D / len(train_loader):.4f}")

    test_input_dir = "/home/work/Dacon_Dataset/test_input"
    output_dir = f"output_images_epoch_{epoch+1}"
    os.makedirs(output_dir, exist_ok=True)
    with torch.no_grad():
        for img_name in sorted(os.listdir(test_input_dir)):
            img_path = os.path.join(test_input_dir, img_name)
            img = cv2.imread(img_path)
            input_tensor = torch.tensor(img).permute(2, 0, 1).unsqueeze(0).float().to(device) / 255.0
            output = generator(input_tensor).squeeze().permute(1, 2, 0).cpu().numpy() * 255.0
            output = output.astype(np.uint8)
            cv2.imwrite(os.path.join(output_dir, img_name), output)

    zip_filename = os.path.join(result_dir, f"epoch_{epoch+1}.zip")
    with zipfile.ZipFile(zip_filename, 'w') as zipf:
        for img_name in os.listdir(output_dir):
            zipf.write(os.path.join(output_dir, img_name), arcname=img_name)
    print(f"Epoch {epoch+1} results saved to {zip_filename}")

    torch.save({
        'epoch': epoch,
        'generator_state_dict': generator.state_dict(),
        'discriminator_state_dict': discriminator.state_dict(),
        'optimizer_G_state_dict': optimizer_G.state_dict(),
        'optimizer_D_state_dict': optimizer_D.state_dict()
    }, checkpoint_path)

generator.train()  
discriminator.train()