import os
from glob import glob
import numpy as np
import torch
from torch.utils.data import Dataset
from PIL import Image
import matplotlib.pyplot as plt
from torchvision import transforms, datasets
import random
import cv2

class CustomDataset(Dataset):
    def __init__(self, list_imgs, list_masks, transform=None):
        self.list_imgs = list_imgs
        self.list_masks = list_masks
        self.transform = transform

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

    def __getitem__(self, index):
        img_path = self.list_imgs[index]
        mask_path = self.list_masks[index]

        img = cv2.imread(img_path, cv2.IMREAD_GRAYSCALE)
        mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)

        # 이미지 크기를 512x512로 변경
        img = cv2.resize(img, (512, 512), interpolation=cv2.INTER_LINEAR)
        mask = cv2.resize(mask, (512, 512), interpolation=cv2.INTER_NEAREST)

        img = img.astype(np.float32) / 255.0
        mask = mask.astype(np.float32) / 255.0

        if img.ndim == 2:
            img = img[:, :, np.newaxis]
        if mask.ndim == 2:
            mask = mask[:, :, np.newaxis]

        data = {'input': img, 'label': mask}

        if self.transform:
            data = self.transform(data)
        
        return data

def create_datasets(img_dir, mask_dir, train_ratio=0.7, val_ratio=0.2, transform=None):
    list_imgs = sorted(glob(os.path.join(img_dir, '**', '*.png'), recursive=True))
    list_masks = sorted(glob(os.path.join(mask_dir, '**', '*.png'), recursive=True))

    combined = list(zip(list_imgs, list_masks))
    random.shuffle(combined)
    list_imgs, list_masks = zip(*combined)

    num_imgs = len(list_imgs)
    num_train = int(num_imgs * train_ratio)
    num_val = int(num_imgs * val_ratio)

    train_set = CustomDataset(list_imgs[:num_train], list_masks[:num_train], transform)
    val_set = CustomDataset(list_imgs[num_train:num_train + num_val], list_masks[num_train:num_train + num_val], transform)
    test_set = CustomDataset(list_imgs[num_train + num_val:], list_masks[num_train + num_val:], transform)

    return train_set, val_set, test_set

# 사용 예시
dir_img = 'C:/Users/pinb/Desktop/imgs'
dir_mask = 'C:/Users/pinb/Desktop/masks'
train_set, val_set, test_set = create_datasets(dir_img, dir_mask)

# 첫 번째 이미지 확인 (예: train set)
data = train_set.__getitem__(7777)  # 이미지 불러오기

input_img = data['input']
label = data['label']

# 이미지 시각화
plt.subplot(121)
plt.imshow(input_img.reshape(input_img.shape[0], input_img.shape[1]), cmap='gray')
plt.title('Input Image')

plt.subplot(122)
plt.imshow(label.reshape(label.shape[0], label.shape[1]), cmap='gray')
plt.title('Label')

plt.show()

# 트렌스폼 구현하기
class ToTensor(object):
    # def __call__(self, data):
    #     label, input = data['label'], data['input']

    #     label = label.transpose((2, 0, 1)).astype(np.float32)
    #     input = input.transpose((2, 0, 1)).astype(np.float32)

    #     data = {'label': torch.from_numpy(label), 'input': torch.from_numpy(input)}

    #     return data
    def __call__(self, data):
        label, input = data['label'], data['input']

        # 이미지가 이미 그레이스케일이면 채널 차원 추가
        if label.ndim == 2:
            label = label[:, :, np.newaxis]
        if input.ndim == 2:
            input = input[:, :, np.newaxis]

        # 채널을 첫 번째 차원으로 이동
        label = label.transpose((2, 0, 1)).astype(np.float32)
        input = input.transpose((2, 0, 1)).astype(np.float32)

        data = {'label': torch.from_numpy(label), 'input': torch.from_numpy(input)}

        return data

class Normalization(object):
    def __init__(self, mean=0.5, std=0.5):
        self.mean = mean
        self.std = std

    def __call__(self, data):
        label, input = data['label'], data['input']

        input = (input - self.mean) / self.std

        data = {'label': label, 'input': input}

        return data

class RandomFlip(object):
    def __call__(self, data):
        label, input = data['label'], data['input']

        if np.random.rand() > 0.5:
            label = np.fliplr(label)
            input = np.fliplr(input)

        if np.random.rand() > 0.5:
            label = np.flipud(label)
            input = np.flipud(input)

        data = {'label': label, 'input': input}

        return data
    
# class Resize(object):
#     def __init__(self, output_size):
#         assert isinstance(output_size, (int, tuple))
#         self.output_size = output_size

#     def __call__(self, data):
#         label, input = data['label'], data['input']

#         h, w = input.shape[:2]
#         if isinstance(self.output_size, int):
#             if h > w:
#                 new_h, new_w = self.output_size * h / w, self.output_size
#             else:
#                 new_h, new_w = self.output_size, self.output_size * w / h
#         else:
#             new_h, new_w = self.output_size

#         new_h, new_w = int(new_h), int(new_w)

#         input = cv2.resize(input, (new_w, new_h))
#         label = cv2.resize(label, (new_w, new_h))

#         return {'label': label, 'input': input}

class Rotate(object):
    def __init__(self, angle_range):
        assert isinstance(angle_range, (tuple, list)) and len(angle_range) == 2
        self.angle_min, self.angle_max = angle_range

    def __call__(self, data):
        label, input = data['label'], data['input']

        # NumPy 배열로 변환 (필요한 경우)
        if not isinstance(input, np.ndarray):
            input = np.array(input)
        if not isinstance(label, np.ndarray):
            label = np.array(label)

        # (H, W, C) 형태를 (H, W)로 변경 (필요한 경우)
        if input.ndim == 3 and input.shape[2] == 1:
            input = input.squeeze(2)
        if label.ndim == 3 and label.shape[2] == 1:
            label = label.squeeze(2)

        # 랜덤 각도 선택 및 회전 적용
        angle = np.random.uniform(self.angle_min, self.angle_max)
        h, w = input.shape[:2]
        center = (w / 2, h / 2)
        rot_matrix = cv2.getRotationMatrix2D(center, angle, 1.0)
        input = cv2.warpAffine(input, rot_matrix, (w, h))
        label = cv2.warpAffine(label, rot_matrix, (w, h))

        return {'label': label, 'input': input}
    
# class Crop(object):
#     def __init__(self, output_size):
#         assert isinstance(output_size, (int, tuple))
#         if isinstance(output_size, int):
#             self.output_size = (output_size, output_size)
#         else:
#             assert len(output_size) == 2
#             self.output_size = output_size

#     def __call__(self, data):
#         label, input = data['label'], data['input']

#         h, w = input.shape[:2]
#         new_h, new_w = self.output_size

#         top = np.random.randint(0, h - new_h)
#         left = np.random.randint(0, w - new_w)

#         input = input[top: top + new_h, left: left + new_w]
#         label = label[top: top + new_h, left: left + new_w]

#         return {'label': label, 'input': input}

# 트랜스폼 잘 구현되었는지 확인
transform = transforms.Compose([Normalization(mean=0.5, std=0.5), RandomFlip(), Rotate(angle_range=(-90, 90)), ToTensor()])

dir_img = 'C:/Users/pinb/Desktop/imgs'
dir_mask = 'C:/Users/pinb/Desktop/masks'
train_set, val_set, test_set = create_datasets(img_dir=dir_img, mask_dir=dir_mask, transform=transform)


data = train_set.__getitem__(12599) # 한 이미지 불러오기
input = data['input']
label = data['label']

# 불러온 이미지 시각화
plt.subplot(122)
plt.hist(label.flatten(), bins=20)
plt.title('label')

plt.subplot(121)
plt.hist(input.flatten(), bins=20)
plt.title('input')

# 이미지 시각화
plt.subplot(121)
plt.imshow(input.squeeze(), cmap='gray')
plt.title('Input Image')

plt.subplot(122)
plt.imshow(label.squeeze(), cmap='gray')
plt.title('Label')

plt.tight_layout()
plt.show()

## 라이브러리 불러오기
import os
import numpy as np

import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter

import matplotlib.pyplot as plt

## 네트워크 구축하기
class UNet(nn.Module):
    def __init__(self):
        super(UNet, self).__init__()

        # Convolution + BatchNormalization + Relu 정의하기
        def CBR2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1, bias=True): 
            layers = []
            layers += [nn.Conv2d(in_channels=in_channels, out_channels=out_channels,
                                 kernel_size=kernel_size, stride=stride, padding=padding,
                                 bias=bias)]
            layers += [nn.BatchNorm2d(num_features=out_channels)]
            layers += [nn.ReLU()]

            cbr = nn.Sequential(*layers)

            return cbr

        # 수축 경로(Contracting path)
        self.enc1_1 = CBR2d(in_channels=1, out_channels=64)
        self.enc1_2 = CBR2d(in_channels=64, out_channels=64)

        self.pool1 = nn.MaxPool2d(kernel_size=2)

        self.enc2_1 = CBR2d(in_channels=64, out_channels=128)
        self.enc2_2 = CBR2d(in_channels=128, out_channels=128)

        self.pool2 = nn.MaxPool2d(kernel_size=2)

        self.enc3_1 = CBR2d(in_channels=128, out_channels=256)
        self.enc3_2 = CBR2d(in_channels=256, out_channels=256)

        self.pool3 = nn.MaxPool2d(kernel_size=2)

        self.enc4_1 = CBR2d(in_channels=256, out_channels=512)
        self.enc4_2 = CBR2d(in_channels=512, out_channels=512)

        self.pool4 = nn.MaxPool2d(kernel_size=2)

        self.enc5_1 = CBR2d(in_channels=512, out_channels=1024)

        # 확장 경로(Expansive path)
        self.dec5_1 = CBR2d(in_channels=1024, out_channels=512)

        self.unpool4 = nn.ConvTranspose2d(in_channels=512, out_channels=512,
                                          kernel_size=2, stride=2, padding=0, bias=True)

        self.dec4_2 = CBR2d(in_channels=2 * 512, out_channels=512)
        self.dec4_1 = CBR2d(in_channels=512, out_channels=256)

        self.unpool3 = nn.ConvTranspose2d(in_channels=256, out_channels=256,
                                          kernel_size=2, stride=2, padding=0, bias=True)

        self.dec3_2 = CBR2d(in_channels=2 * 256, out_channels=256)
        self.dec3_1 = CBR2d(in_channels=256, out_channels=128)

        self.unpool2 = nn.ConvTranspose2d(in_channels=128, out_channels=128,
                                          kernel_size=2, stride=2, padding=0, bias=True)

        self.dec2_2 = CBR2d(in_channels=2 * 128, out_channels=128)
        self.dec2_1 = CBR2d(in_channels=128, out_channels=64)

        self.unpool1 = nn.ConvTranspose2d(in_channels=64, out_channels=64,
                                          kernel_size=2, stride=2, padding=0, bias=True)

        self.dec1_2 = CBR2d(in_channels=2 * 64, out_channels=64)
        self.dec1_1 = CBR2d(in_channels=64, out_channels=64)

        self.fc = nn.Conv2d(in_channels=64, out_channels=1, kernel_size=1, stride=1, padding=0, bias=True)

        # self.enc1_1 = CBR2d(in_channels=1, out_channels=64)
        # self.pool1 = nn.MaxPool2d(kernel_size=2)

        # self.enc2_1 = CBR2d(in_channels=64, out_channels=128)
        # self.pool2 = nn.MaxPool2d(kernel_size=2)

        # self.enc3_1 = CBR2d(in_channels=128, out_channels=256)
        # self.pool3 = nn.MaxPool2d(kernel_size=2)

        # self.enc4_1 = CBR2d(in_channels=256, out_channels=512)
        # self.pool4 = nn.MaxPool2d(kernel_size=2)

        # self.enc5_1 = CBR2d(in_channels=512, out_channels=1024)

        # # 확장 경로(Expansive path)의 깊이 감소
        # self.dec5_1 = CBR2d(in_channels=1024, out_channels=512)
        # self.unpool4 = nn.ConvTranspose2d(in_channels=512, out_channels=512, kernel_size=2, stride=2)

        # self.dec4_1 = CBR2d(in_channels=512 + 512, out_channels=256)
        # self.unpool3 = nn.ConvTranspose2d(in_channels=256, out_channels=256, kernel_size=2, stride=2)

        # self.dec3_1 = CBR2d(in_channels=256 + 256, out_channels=128)
        # self.unpool2 = nn.ConvTranspose2d(in_channels=128, out_channels=128, kernel_size=2, stride=2)

        # self.dec2_1 = CBR2d(in_channels=128 + 128, out_channels=64)
        # self.unpool1 = nn.ConvTranspose2d(in_channels=64, out_channels=64, kernel_size=2, stride=2)

        # self.dec1_1 = CBR2d(in_channels=64 + 64, out_channels=64)
        # self.fc = nn.Conv2d(in_channels=64, out_channels=1, kernel_size=1, stride=1, padding=0, bias=True)
    
    # forward 함수 정의하기
    def forward(self, x):
        enc1_1 = self.enc1_1(x)
        enc1_2 = self.enc1_2(enc1_1)
        pool1 = self.pool1(enc1_2)

        enc2_1 = self.enc2_1(pool1)
        enc2_2 = self.enc2_2(enc2_1)
        pool2 = self.pool2(enc2_2)

        enc3_1 = self.enc3_1(pool2)
        enc3_2 = self.enc3_2(enc3_1)
        pool3 = self.pool3(enc3_2)

        enc4_1 = self.enc4_1(pool3)
        enc4_2 = self.enc4_2(enc4_1)
        pool4 = self.pool4(enc4_2)

        enc5_1 = self.enc5_1(pool4)

        dec5_1 = self.dec5_1(enc5_1)

        unpool4 = self.unpool4(dec5_1)
        cat4 = torch.cat((unpool4, enc4_2), dim=1)
        dec4_2 = self.dec4_2(cat4)
        dec4_1 = self.dec4_1(dec4_2)

        unpool3 = self.unpool3(dec4_1)
        cat3 = torch.cat((unpool3, enc3_2), dim=1)
        dec3_2 = self.dec3_2(cat3)
        dec3_1 = self.dec3_1(dec3_2)

        unpool2 = self.unpool2(dec3_1)
        cat2 = torch.cat((unpool2, enc2_2), dim=1)
        dec2_2 = self.dec2_2(cat2)
        dec2_1 = self.dec2_1(dec2_2)

        unpool1 = self.unpool1(dec2_1)
        cat1 = torch.cat((unpool1, enc1_2), dim=1)
        dec1_2 = self.dec1_2(cat1)
        dec1_1 = self.dec1_1(dec1_2)

        x = self.fc(dec1_1)

        # enc1_1 = self.enc1_1(x)
        # pool1 = self.pool1(enc1_1)

        # enc2_1 = self.enc2_1(pool1)
        # pool2 = self.pool2(enc2_1)

        # enc3_1 = self.enc3_1(pool2)
        # pool3 = self.pool3(enc3_1)

        # enc4_1 = self.enc4_1(pool3)
        # pool4 = self.pool4(enc4_1)

        # enc5_1 = self.enc5_1(pool4)

        # dec5_1 = self.dec5_1(enc5_1)

        # unpool4 = self.unpool4(dec5_1)
        # cat4 = torch.cat((unpool4, enc4_1), dim=1)
        # dec4_1 = self.dec4_1(cat4)

        # unpool3 = self.unpool3(dec4_1)
        # cat3 = torch.cat((unpool3, enc3_1), dim=1)
        # dec3_1 = self.dec3_1(cat3)

        # unpool2 = self.unpool2(dec3_1)
        # cat2 = torch.cat((unpool2, enc2_1), dim=1)
        # dec2_1 = self.dec2_1(cat2)

        # unpool1 = self.unpool1(dec2_1)
        # cat1 = torch.cat((unpool1, enc1_1), dim=1)
        # dec1_1 = self.dec1_1(cat1)

        # x = self.fc(dec1_1)

        return x

## 네트워크 저장하기
def save(ckpt_dir, net, optim, epoch):
    if not os.path.exists(ckpt_dir):
        os.makedirs(ckpt_dir)

    torch.save({'net': net.state_dict(), 'optim': optim.state_dict()},
               "%s/model_epoch%d.pth" % (ckpt_dir, epoch))

## 네트워크 불러오기
def load(ckpt_dir, net, optim):
    if not os.path.exists(ckpt_dir):
        epoch = 0
        return net, optim, epoch

    ckpt_lst = os.listdir(ckpt_dir)
    ckpt_lst.sort(key=lambda f: int(''.join(filter(str.isdigit, f))))

    dict_model = torch.load('%s/%s' % (ckpt_dir, ckpt_lst[-1]))

    net.load_state_dict(dict_model['net'])
    optim.load_state_dict(dict_model['optim'])
    epoch = int(ckpt_lst[-1].split('epoch')[1].split('.pth')[0])

    return net, optim, epoch

torch.cuda.empty_cache()
os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:128'

dir_img = 'C:/Users/pinb/Desktop/imgs'
dir_mask = 'C:/Users/pinb/Desktop/masks'
transform = transforms.Compose([Normalization(mean=0.5, std=0.5), RandomFlip(), Rotate(angle_range=(-90, 90)), ToTensor()])
train_set, val_set, test_set = create_datasets(img_dir=dir_img, mask_dir=dir_mask, transform=transform)

# 훈련 파라미터 설정하기
lr = 1e-3
batch_size = 4
num_epoch = 10

base_dir = './2nd_Battery/unet'
ckpt_dir = os.path.join(base_dir, "checkpoint")
log_dir = os.path.join(base_dir, "log")

# 네트워크 생성하기
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
net = UNet().to(device)

# 손실함수 정의하기
fn_loss = nn.BCEWithLogitsLoss().to(device)

# Optimizer 설정하기
optim = torch.optim.Adam(net.parameters(), lr=lr)

# 그 밖에 부수적인 functions 설정하기
fn_tonumpy = lambda x: x.to('cpu').detach().numpy().transpose(0, 2, 3, 1)
fn_denorm = lambda x, mean, std: (x * std) + mean
fn_class = lambda x: 1.0 * (x > 0.5)

# 훈련을 위한 Transform과 DataLoader
loader_train = DataLoader(train_set, batch_size=batch_size, shuffle=True, num_workers=0)
loader_val = DataLoader(val_set, batch_size=batch_size, shuffle=False, num_workers=0)

# 그밖에 부수적인 variables 설정하기
num_data_train = len(train_set)
num_data_val = len(val_set)

num_batch_train = np.ceil(num_data_train / batch_size)
num_batch_val = np.ceil(num_data_val / batch_size)


# Tensorboard 를 사용하기 위한 SummaryWriter 설정
writer_train = SummaryWriter(log_dir=os.path.join(log_dir, 'train'))
writer_val = SummaryWriter(log_dir=os.path.join(log_dir, 'val'))

# 네트워크 학습시키기
st_epoch = 0
# 학습한 모델이 있을 경우 모델 로드하기
net, optim, st_epoch = load(ckpt_dir=ckpt_dir, net=net, optim=optim) 

for epoch in range(st_epoch + 1, num_epoch + 1):
        net.train()
        loss_arr = []

        for batch, data in enumerate(loader_train, 1):
            # forward pass
            label = data['label'].to(device)
            input = data['input'].to(device)

            output = net(input)

            # backward pass
            optim.zero_grad()

            loss = fn_loss(output, label)
            loss.backward()

            optim.step()

            # 손실함수 계산
            loss_arr += [loss.item()]

            print("TRAIN: EPOCH %04d / %04d | BATCH %04d / %04d | LOSS %.4f" %
                  (epoch, num_epoch, batch, num_batch_train, np.mean(loss_arr)))

            # Tensorboard 저장하기
            label = fn_tonumpy(label)
            input = fn_tonumpy(fn_denorm(input, mean=0.5, std=0.5))
            output = fn_tonumpy(fn_class(output))

            writer_train.add_image('label', label, num_batch_train * (epoch - 1) + batch, dataformats='NHWC')
            writer_train.add_image('input', input, num_batch_train * (epoch - 1) + batch, dataformats='NHWC')
            writer_train.add_image('output', output, num_batch_train * (epoch - 1) + batch, dataformats='NHWC')

        writer_train.add_scalar('loss', np.mean(loss_arr), epoch)

        with torch.no_grad():
            net.eval()
            loss_arr = []

            for batch, data in enumerate(loader_val, 1):
                # forward pass
                label = data['label'].to(device)
                input = data['input'].to(device)

                output = net(input)

                # 손실함수 계산하기
                loss = fn_loss(output, label)

                loss_arr += [loss.item()]

                print("VALID: EPOCH %04d / %04d | BATCH %04d / %04d | LOSS %.4f" %
                      (epoch, num_epoch, batch, num_batch_val, np.mean(loss_arr)))

                # Tensorboard 저장하기
                label = fn_tonumpy(label)
                input = fn_tonumpy(fn_denorm(input, mean=0.5, std=0.5))
                output = fn_tonumpy(fn_class(output))

                writer_val.add_image('label', label, num_batch_val * (epoch - 1) + batch, dataformats='NHWC')
                writer_val.add_image('input', input, num_batch_val * (epoch - 1) + batch, dataformats='NHWC')
                writer_val.add_image('output', output, num_batch_val * (epoch - 1) + batch, dataformats='NHWC')

        writer_val.add_scalar('loss', np.mean(loss_arr), epoch)

        # epoch 5마다 모델 저장하기
        if epoch % 1 == 0:
            save(ckpt_dir=ckpt_dir, net=net, optim=optim, epoch=epoch)

        writer_train.close()
        writer_val.close()

print('train set: ' + str(len(train_set)))
print('val set: ' + str(len(val_set)))
print('test set: ' + str(len(test_set)))
print('total: ' + str(len(train_set)+ len(val_set)+ len(test_set)))

loader_test = DataLoader(test_set, batch_size=batch_size, shuffle=False, num_workers=0)

# 그밖에 부수적인 variables 설정하기
num_data_test = len(test_set)
num_batch_test = np.ceil(num_data_test / batch_size)

# 결과 디렉토리 생성하기
result_dir = os.path.join(base_dir, 'result')
if not os.path.exists(result_dir):
    os.makedirs(os.path.join(result_dir, 'png'))
    os.makedirs(os.path.join(result_dir, 'numpy'))


net, optim, st_epoch = load(ckpt_dir=ckpt_dir, net=net, optim=optim)

with torch.no_grad():
      net.eval()
      loss_arr = []

      for batch, data in enumerate(loader_test, 1):
          # forward pass
          label = data['label'].to(device)
          input = data['input'].to(device)

          output = net(input)

          # 손실함수 계산하기
          loss = fn_loss(output, label)

          loss_arr += [loss.item()]

          print("TEST: BATCH %04d / %04d | LOSS %.4f" %
                (batch, num_batch_test, np.mean(loss_arr)))

          # Tensorboard 저장하기
          label = fn_tonumpy(label)
          input = fn_tonumpy(fn_denorm(input, mean=0.5, std=0.5))
          output = fn_tonumpy(fn_class(output))

          # 테스트 결과 저장하기
          for j in range(label.shape[0]):
              id = num_batch_test * (batch - 1) + j

              plt.imsave(os.path.join(result_dir, 'png', 'label_%04d.png' % id), label[j].squeeze(), cmap='gray')
              plt.imsave(os.path.join(result_dir, 'png', 'input_%04d.png' % id), input[j].squeeze(), cmap='gray')
              plt.imsave(os.path.join(result_dir, 'png', 'output_%04d.png' % id), output[j].squeeze(), cmap='gray')

              np.save(os.path.join(result_dir, 'numpy', 'label_%04d.npy' % id), label[j].squeeze())
              np.save(os.path.join(result_dir, 'numpy', 'input_%04d.npy' % id), input[j].squeeze())
              np.save(os.path.join(result_dir, 'numpy', 'output_%04d.npy' % id), output[j].squeeze())

print("AVERAGE TEST: BATCH %04d / %04d | LOSS %.4f" %
        (batch, num_batch_test, np.mean(loss_arr)))

##
lst_data = os.listdir(os.path.join(result_dir, 'numpy'))

lst_label = [f for f in lst_data if f.startswith('label')]
lst_input = [f for f in lst_data if f.startswith('input')]
lst_output = [f for f in lst_data if f.startswith('output')]

lst_label.sort()
lst_input.sort()
lst_output.sort()

##
id = 0

label = np.load(os.path.join(result_dir,"numpy", lst_label[id]))
input = np.load(os.path.join(result_dir,"numpy", lst_input[id]))
output = np.load(os.path.join(result_dir,"numpy", lst_output[id]))

## 플롯 그리기
plt.figure(figsize=(8,6))
plt.subplot(131)
plt.imshow(input, cmap='gray')
plt.title('input')

plt.subplot(132)
plt.imshow(label, cmap='gray')
plt.title('label')

plt.subplot(133)
plt.imshow(output, cmap='gray')
plt.title('output')

plt.show()