import gradio as gr
import numpy as np
import os
from PIL import Image

def upscale_image(input_image, upscale_factor):
    try:
        # PIL Image로 변환
        if isinstance(input_image, np.ndarray):
            input_image = Image.fromarray(input_image)
        
        # 현재 이미지 크기
        current_width, current_height = input_image.size
        
        # 새로운 크기 계산
        new_width = int(current_width * upscale_factor)
        new_height = int(current_height * upscale_factor)
        
        # 이미지 리사이즈
        output_image = input_image.resize((new_width, new_height), Image.Resampling.LANCZOS)
        
        # 임시 파일 저장을 위한 디렉토리 확인 및 생성
        if not os.path.exists("temp"):
            os.makedirs("temp")
            
        # 결과 이미지 저장
        output_path = os.path.join("temp", "upscaled_image.png")
        output_image.save(output_path, "PNG", quality=100)
        
        return output_path
        
    except Exception as e:
        print(f"Error in upscale_image: {str(e)}")
        raise gr.Error("이미지 처리 중 오류가 발생했습니다. 다시 시도해주세요.")

# 인터페이스 설명
DESCRIPTION = """
# 이미지 업스케일러 (Image Upscaler)
이미지의 크기와 품질을 향상시켜 보세요!
You can increase the size and quality of your images!

## 사용 방법 (How to use):
1. 이미지를 업로드하세요 (Upload an image)
2. 업스케일 레벨을 선택하세요 (Select upscale level)
3. Submit 버튼을 클릭하세요 (Click submit)
"""

# Gradio 인터페이스 정의
demo = gr.Interface(
    fn=upscale_image,
    inputs=[
        gr.Image(label="입력 이미지 (Input Image)", type="pil"),
        gr.Radio(
            label="업스케일 레벨 (Upscale Level)", 
            choices=[2, 4, 6, 8],
            value=2,
            type="value"  # 'number' 에서 'value'로 변경
        )
    ],
    outputs=gr.File(label="업스케일된 이미지 다운로드 (Download Upscaled Image)"),
    title="Image Upscaler",
    description=DESCRIPTION,
    allow_flagging="never"
)

# 인터페이스 실행
demo.launch()