#모델 도출까지는 구글 코랩 이용
#dataset은 AISegment 중 무작위 3000장 이용

#1 기본 셋팅

#%%
#사용할 라이브러리 로드
import tensorflow as tf
from tensorflow import keras
from keras.layers import Dense
from keras.models import Sequential

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
import warnings 
from IPython.display import Image

from skimage import color

from keras.layers import Dense, Input, Conv2D, Conv2DTranspose, Flatten, Reshape, Cropping2D
from keras.models import Model
from keras.layers import BatchNormalization, Dropout, Activation, MaxPool2D, concatenate

warnings.filterwarnings('ignore')
%matplotlib inline

SEED = 34

#%%
#GPU와 파일 로딩시 메모리 설정
# GPU 사용 설정
gpus = tf.config.list_physical_devices('GPU')
if gpus:
    try:
        # GPU 메모리 사용 제한 설정 (필요에 따라 조절)
        for gpu in gpus:
            tf.config.experimental.set_memory_growth(gpu, True)
        logical_gpus = tf.config.list_logical_devices('GPU')
        print(len(gpus), "Physical GPUs,", len(logical_gpus), "Logical GPUs")
    except RuntimeError as e:
        print(e)

# 데이터 로드 및 전처리 (데이터 제너레이터 사용)
class DataGenerator(keras.utils.Sequence):
    def __init__(self, images, masks, batch_size=128):
        self.images = images
        self.masks = masks
        self.batch_size = batch_size

    def __len__(self):
        return int(np.ceil(len(self.images) / float(self.batch_size)))

    def __getitem__(self, idx):
        batch_images = self.images[idx * self.batch_size:(idx + 1) * self.batch_size]
        batch_masks = self.masks[idx * self.batch_size:(idx + 1) * self.batch_size]

        return np.array(batch_images), np.array(batch_masks)
#%%
# .npz 파일 경로
file_path = r"C:\Users\remil\바탕 화면\productive\last.npz"

# .npz 파일 로드
data = np.load(file_path)

#%%
#데이터 로드
train_generator = DataGenerator(data["train_images"], batch_size=128)
test_generator = DataGenerator(data["test_images"], batch_size=128)

#%%
#값 출력
tr_images, tr_masks = train_generator[5]
t_images, t_masks = test_generator[5]
print(tr_images.min(), tr_images.max())
print(tr_masks.min(), tr_masks.max())
print(t_images.min(), t_images.max())
print(t_masks.min(), t_masks.max())

#%%
#데이터 유형 확인
print(tr_images.dtype, tr_images.shape)
print(tr_masks.dtype, tr_masks.shape)
print(t_images.dtype, t_images.shape)
print(t_masks.dtype, t_masks.shape)

#%%
plt.imshow(train_generator[0,0])
plt.imshow(tr_masks[0,0])

#%%
#모델 작성
from keras.layers import Dense, Input, Conv2D, Conv2DTranspose, Flatten, Reshape, Cropping2D
from keras.models import Model
from keras.layers import BatchNormalization, Dropout, Activation, MaxPool2D, concatenate

def conv2d_block(x, channel):
  x = Conv2D(channel, 3, padding="same")(x)
  x = BatchNormalization()(x)
  x = Activation("relu")(x)

  x = Conv2D(channel, 3, padding="same")(x)
  x = BatchNormalization()(x)
  x = Activation("relu")(x)
  return x

def unet_like():
  inputs = Input((256, 256, 3))

  c1 = conv2d_block(inputs, 16)
  p1 = MaxPool2D((2,2))(c1)
  p1 = Dropout(0.1)(p1)

  c2 = conv2d_block(p1, 32)
  p2 = MaxPool2D((2,2))(c2)
  p2 = Dropout(0.1)(p2)

  c3 = conv2d_block(p2, 64)
  p3 = MaxPool2D((2,2))(c3)
  p3 = Dropout(0.1)(p3)

  c4 = conv2d_block(p3, 128)
  p4 = MaxPool2D((2,2))(c4)
  p4 = Dropout(0.1)(p4)

  c5 = conv2d_block(p4, 256)

  u6 = Conv2DTranspose(128, 2, 2)(c5) 
  u6 = concatenate([u6, c4])
  u6 = Dropout(0.1)(u6)
  c6 = conv2d_block(u6, 128)

  u7 = Conv2DTranspose(64, 2, 2)(c6)
  u7 = concatenate([u7, c3])
  u7 = Dropout(0.1)(u7)
  c7 = conv2d_block(u7, 64)

  u8 = Conv2DTranspose(32, 2, 2)(c7)
  u8 = concatenate([u8, c2])
  u8 = Dropout(0.1)(u8)
  c8 = conv2d_block(u8, 32)

  u9 = Conv2DTranspose(16, 2, 2)(c8) 
  u9 = concatenate([u9, c1])
  u9 = Dropout(0.1)(u9)
  c9 = conv2d_block(u9, 16)

  outputs = Conv2D(1, (1,1), activation="sigmoid")(c9)

  model = Model(inputs, outputs)
  return model

#%%
#서머리 출력
model = unet_like()
model.summary()

#%%
#로스 옵티마이저 매트릭 설정
model.compile(loss="mse", optimizer="adam", metrics=["accuracy"])

# %%
#모델 학습
hist = model.fit(train_generator, validation_data=test_generator, epochs=10, verbose=1)

# %%
tf.keras.saving.save_model(model, 'my_model.keras')
#%%
plt.imshow(tr_masks[11])
# %%
#모델 로드
model45 = tf.keras.models.load_model("./nukki25645.h5")
model90 = tf.keras.models.load_model("./nukki12090.h5")
# %%
import imageio
from skimage.transform import resize
import matplotlib.pyplot as plt
import numpy as np

harry = imageio.imread('./harry.jpeg')
original_shape = harry.shape

# 이미지 크기를 (256, 256, 3)으로 조정
harry_resized = resize(harry, (256, 256))

# 예측
test = model.predict(harry_resized[None, ...])

# 예측 결과를 원본 크기로 다시 조정
test_resized = resize(test[0], original_shape[:2])

# 예측 결과를 0 또는 1로 이원화 (마스크 생성)
mask = (test_resized > 0.5).astype(np.uint8) 

# 마스크를 3채널로 복제하여 원본 이미지와 곱하기
mask_3d = np.repeat(mask[:, :, np.newaxis], 3, axis=2) 
cropped_harry = harry * mask_3d[:,:,0]

# 잘라낸 이미지 표시
plt.imshow(cropped_harry)
plt.show()



# %%
import imageio
from skimage.transform import resize

harry = imageio.imread('./harry.jpeg')
original_shape = harry.shape  # 원본 이미지 크기 저장

# 이미지 크기를 (256, 256, 3)으로 조정
harry_resized = resize(harry, (256, 256))

# 이제 크기가 조정된 이미지를 사용하여 예측
test = model.predict(harry_resized[None, ...])

# 예측 결과를 원본 크기로 다시 조정
test_resized = resize(test[0], original_shape[:2])  # 높이와 너비만 사용

# 예측 결과 표시 (원본 크기)
plt.imshow(test_resized, cmap='gray')
plt.show()



# %%
import imageio
from skimage.transform import resize
import matplotlib.pyplot as plt
import numpy as np

# 이미지 로드 및 크기 조정
harry = imageio.imread('./harry.jpeg')
original_shape = harry.shape
harry_resized = resize(harry, (256, 256))

# 모델 예측
test = model.predict(harry_resized[None, ...])
test_resized = resize(test[0], original_shape[:2])

# 마스크 생성
mask = (test_resized > 0.5).astype(np.uint8)

# 마스크를 3채널로 복제
mask_3d = np.repeat(mask[:, :, np.newaxis], 3, axis=2)

# 색상 적용 (예: 빨간색)
colored_mask = np.zeros_like(harry)
colored_mask[:, :, 0] = 255 * mask

# 원본 이미지와 색상 입힌 마스크 합성
colored_harry = harry * (1 - mask_3d) + colored_mask * mask_3d

# 결과 표시
plt.imshow(colored_harry)
plt.show()


# %%
harry = imageio.imread('./harry.jpeg')
original_shape = harry.shape
harry_resized = resize(harry, (256, 256))

# 모델 예측 (확률값)
test = model.predict(harry_resized[None, ...])[0]
test_resized = resize(test, original_shape[:2])

# 마스크 생성 (1% 이상 확률값)
mask = (test_resized > 0.20).astype(np.uint8)

# 마스크를 3채널로 복제
mask_3d = np.repeat(mask[:, :, np.newaxis], 3, axis=2)

mask_3d = mask_3d[:, :, :, 0] # or mask_3d = mask_3d[:, :, :, 0]

# 원본 이미지와 마스크 곱하기
cropped_harry = harry * mask_3d

# 결과 표시
plt.imshow(cropped_harry)
plt.show()



# %%
plt.imshow(harry)
# %%