from rembg import remove, new_session
from PIL import Image
import os

import imageio
import numpy as np

import tensorflow as tf
from tensorflow import keras
from skimage import color
from skimage.transform import resize

#이미지 배경을 제거하고 결과를 저장
def process_image(upload_path, result_path, model_name):
    input_image = Image.open(upload_path)
    session = new_session(model_name) # 모델 세션 생성
    output_image = remove(input_image, session=session) # 모델 세션 전달
    output_image.save(result_path)

#결과 파일 이름을 생성
def generate_result_filename(filename):
    base_filename, _ = os.path.splitext(filename)
    return f"result_{base_filename}.png"

#결과 이미지 경로를 생성
def get_result_image_path(result_folder, result_filename):
    return os.path.join(result_folder, result_filename)

#업로드된 파일을 저장
def save_uploaded_file(upload_folder, file):
    upload_path = os.path.join(upload_folder, file.filename)
    file.save(upload_path)
    return upload_path

#이미지를 흑백으로 바꾸고 결과를 저장
def process_gray(upload_path, result_path):
    input_image = Image.open(upload_path)
    image = np.array(input_image)
    gray_image = np.mean(image, axis=-1).reshape(image.shape[:2] + (1,)).astype(np.float64)
    gray_image = Image.fromarray(gray_image.squeeze())
    gray_image = gray_image.convert("L") # L모드로 변환
    gray_image.save(result_path)

#이미지를 컬러로 바꾸고 결과를 저장
def process_color(upload_path, result_path):
    input_image = Image.open(upload_path)
    model_path = os.path.join(os.path.dirname(__file__), 'color.keras')
    model = tf.keras.models.load_model(model_path) # 모델 로드
    original_size = input_image.size # 원본 이미지 사이즈 저장

    resized = resize(np.array(input_image), (256, 256, 3)) # 이미지 크기 리사이징
    resized = np.mean(resized, axis=-1, keepdims=True) # 흑백으로 변경
    harry_gray_reshaped = resized.reshape((1, 256, 256, 1)) # 배치 추가

    test = model.predict(harry_gray_reshaped) # 모델 예측
    test_resized = resize(test[0], (original_size[1], original_size[0])) # 너비와 높이를 바꿔서 전달
    test_resized = Image.fromarray(np.uint8(test_resized * 255)) # 0~1 사이 값을 0~255로 조정

    test_resized.save(result_path)


# 이미지를 업스케일링
def process_sizer(upload_path, result_path):
    input_image = Image.open(upload_path)
    model_path = os.path.join(os.path.dirname(__file__), 'sizer.keras')
    model = keras.models.load_model(model_path)
    input_image, original_size = preprocess_image(input_image)
    if len(input_image.shape) == 3:
        input_image = np.expand_dims(input_image, axis=0)
        output_image = model.predict(input_image)[0]
    else:
        output_image = process_large_image(input_image, original_size, model)

    output_image = postprocess_image(output_image, original_size)
    output_image = Image.fromarray(output_image.astype(np.uint8))
    output_image.save(result_path)

# 크기에 따른 전처리
def preprocess_image(img, input_size=(128, 128)):
    img = img.convert('RGB')
    original_width, original_height = img.size
    img_array = np.array(img) / 255.0

    resized_img = resize(img_array, input_size)
    return resized_img, (original_width, original_height)

# 모델 출력 이미지 후처리
def postprocess_image(output_image, original_size):
    target_size = (original_size[1] * 2, original_size[0] * 2)

    output_image = resize(output_image, target_size)
    output_image = np.clip(output_image * 255, 0, 255).astype(np.uint8)

    return output_image

# 큰 이미지 분할 처리 병합
def process_large_image(img_array, original_size, model):
    tile_size = 128
    overlap = 16
    stride = tile_size - overlap

    original_height, original_width, _ = img_array.shape
    output_height, output_width = original_height * 2, original_width * 2

    output_img = np.zeros((output_height, output_width, 3), dtype=np.uint8)

    for y in range(0, original_height - tile_size + 1, stride):
        for x in range(0, original_width - tile_size + 1, stride):
            tile = img_array[y:y + tile_size, x:x + tile_size]
            tile = np.expand_dims(tile, axis=0)
            output_tile = model.predict(tile)[0]
            output_tile = np.clip(output_tile * 255, 0, 255).astype(np.uint8)

            output_y, output_x = y * 2, x * 2
            output_img[output_y:output_y + tile_size * 2, output_x:output_x + tile_size * 2] = output_tile

    return output_img