app.py

import cv2
import numpy as np
from PIL import Image, ImageEnhance, ImageFilter
import gradio as gr
from io import BytesIO
import tempfile
import logging
import replicate
import requests
import os
from dotenv import load_dotenv
import base64
from datetime import datetime, timedelta
import uuid
import time

# .env 파일에서 환경 변수 로드 (없으면 무시)
load_dotenv()

# 로깅 설정 - INFO 레벨로 변경
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Replicate API 토큰 설정
REPLICATE_API_TOKEN = os.getenv("REPLICATE_API_TOKEN", "")
os.environ["REPLICATE_API_TOKEN"] = REPLICATE_API_TOKEN

# 배경 제거 모델 ID를 환경 변수로 숨김 처리
BACKGROUND_REMOVER_MODEL = os.getenv("BACKGROUND_REMOVER_MODEL", "")

# 세션별 이미지 저장용 딕셔너리
session_images = {}

def adjust_brightness(image, value):
    """이미지 밝기 조절"""
    value = float(value - 1) * 100  # 0-2 범위를 -100에서 +100으로 변환
    hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
    h, s, v = cv2.split(hsv)
    v = cv2.add(v, value)
    v = np.clip(v, 0, 255)
    final_hsv = cv2.merge((h, s, v))
    return cv2.cvtColor(final_hsv, cv2.COLOR_HSV2BGR)

def adjust_contrast(image, value):
    """이미지 대비 조절"""
    value = float(value)
    return np.clip(image * value, 0, 255).astype(np.uint8)

def adjust_saturation(image, value):
    """이미지 채도 조절"""
    value = float(value - 1) * 100  # 0-2 범위를 -100에서 +100으로 변환
    hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
    h, s, v = cv2.split(hsv)
    s = cv2.add(s, value)
    s = np.clip(s, 0, 255)
    final_hsv = cv2.merge((h, s, v))
    return cv2.cvtColor(final_hsv, cv2.COLOR_HSV2BGR)

def adjust_temperature(image, value):
    """이미지 색온도 조절 (색상 밸런스)"""
    value = float(value) * 30  # 효과 스케일 조절
    b, g, r = cv2.split(image)
    if value > 0:  # 따뜻하게
        r = cv2.add(r, value)
        b = cv2.subtract(b, value)
    else:  # 차갑게
        r = cv2.add(r, value)
        b = cv2.subtract(b, value)
    
    r = np.clip(r, 0, 255)
    b = np.clip(b, 0, 255)
    return cv2.merge([b, g, r])

def get_korean_timestamp():
    """한국 시간 타임스탬프 생성"""
    korea_time = datetime.utcnow() + timedelta(hours=9)
    return korea_time.strftime('%Y%m%d_%H%M%S')

def remove_background(image, session_id):
    """배경 제거 기능 - Replicate API 사용"""
    logger.info(f"세션 {session_id}: 배경 제거 시작")
    
    # 모델 ID가 설정되지 않은 경우 처리
    if not BACKGROUND_REMOVER_MODEL:
        logger.error("배경 제거 모델 ID가 설정되지 않았습니다.")
        return image
    
    # 임시 파일로 이미지 저장
    temp_file = tempfile.NamedTemporaryFile(suffix='.png', delete=False)
    temp_file_path = temp_file.name
    temp_file.close()
    
    # PIL 이미지를 임시 파일로 저장
    image.save(temp_file_path, format="PNG")
    
    try:
        # Replicate API로 배경 제거 요청 - 환경 변수에서 모델 ID 가져오기
        output = replicate.run(
            BACKGROUND_REMOVER_MODEL,  # 환경 변수로 변경
            input={
                "image": open(temp_file_path, "rb"),
                "format": "png",
                "reverse": False,
                "threshold": 0,
                "background_type": "rgba"
            }
        )
        
        # URL에서 이미지 다운로드
        response = requests.get(output)
        if response.status_code == 200:
            # 바이트 데이터를 PIL 이미지로 변환
            img = Image.open(BytesIO(response.content))
            
            # PNG 파일로 배경 제거된 이미지 저장 (마스크용)
            timestamp = get_korean_timestamp()
            bg_removed_path = tempfile.gettempdir() + f"/bg_removed_{session_id}_{timestamp}.png"
            img.save(bg_removed_path, format="PNG")
            
            logger.info(f"세션 {session_id}: 배경 제거 완료: {bg_removed_path}")
            
            # 세션 정보 저장
            session_images[session_id]['bg_removed_image'] = img
            session_images[session_id]['bg_removed_path'] = bg_removed_path
            
            return img
        else:
            logger.error(f"세션 {session_id}: API 응답 오류: {response.status_code}")
            return image
            
    except Exception as e:
        logger.error(f"세션 {session_id}: 배경 제거 오류: {e}")
        return image
    finally:
        # 임시 파일 삭제
        if os.path.exists(temp_file_path):
            os.unlink(temp_file_path)

def handle_upload(image, session_id=None):
    """이미지 업로드 처리 및 배경 제거"""
    if image is None:
        return None
    
    # 세션 ID 생성 또는 확인
    if session_id is None or session_id not in session_images:
        session_id = str(uuid.uuid4())
        session_images[session_id] = {}
    
    # 원본 이미지 저장
    original_image = Image.fromarray(np.array(image))
    session_images[session_id]['original_image'] = original_image
    
    # 배경 제거 실행
    bg_removed_image = remove_background(original_image, session_id)
    
    # 필터 적용 전 기본 이미지 반환 (필터 없이 원본 이미지)
    return original_image, session_id

def process_image(session_id, temperature, brightness, contrast, saturation, filter_mode):
    """모드에 따른 이미지 처리"""
    if session_id is None or session_id not in session_images:
        logger.warning(f"세션 ID가 유효하지 않습니다: {session_id}")
        return None
    
    original_image = session_images[session_id].get('original_image')
    bg_removed_image = session_images[session_id].get('bg_removed_image')
    
    if original_image is None:
        return None
        
    if bg_removed_image is None:
        logger.warning(f"세션 {session_id}: 배경 제거된 이미지가 없습니다. 전체 필터 모드로 처리합니다.")
        # 전체 필터로 처리
        cv_image = cv2.cvtColor(np.array(original_image), cv2.COLOR_RGB2BGR)
        cv_image = adjust_temperature(cv_image, temperature)
        cv_image = adjust_brightness(cv_image, brightness)
        cv_image = adjust_contrast(cv_image, contrast)
        cv_image = adjust_saturation(cv_image, saturation)
        return Image.fromarray(cv2.cvtColor(cv_image, cv2.COLOR_BGR2RGB))
    
    # 필터 모드에 따라 다르게 처리
    if filter_mode == "전체 필터":
        # 원본 이미지에 필터 적용
        cv_image = cv2.cvtColor(np.array(original_image), cv2.COLOR_RGB2BGR)
        cv_image = adjust_temperature(cv_image, temperature)
        cv_image = adjust_brightness(cv_image, brightness)
        cv_image = adjust_contrast(cv_image, contrast)
        cv_image = adjust_saturation(cv_image, saturation)
        return Image.fromarray(cv2.cvtColor(cv_image, cv2.COLOR_BGR2RGB))
    
    elif filter_mode == "배경만 필터 적용":
        # 원본에 필터 적용 후 배경 제거된 이미지 합성
        cv_image = cv2.cvtColor(np.array(original_image), cv2.COLOR_RGB2BGR)
        cv_image = adjust_temperature(cv_image, temperature)
        cv_image = adjust_brightness(cv_image, brightness)
        cv_image = adjust_contrast(cv_image, contrast)
        cv_image = adjust_saturation(cv_image, saturation)
        filtered_original = Image.fromarray(cv2.cvtColor(cv_image, cv2.COLOR_BGR2RGB))
        
        # 배경 제거된 이미지가 RGBA 모드인지 확인
        if bg_removed_image.mode != 'RGBA':
            logger.warning(f"세션 {session_id}: 배경 제거된 이미지가 RGBA 모드가 아닙니다.")
            return filtered_original
        
        # 필터링된 이미지를 RGBA로 변환 (알파 채널 추가)
        if filtered_original.mode != 'RGBA':
            filtered_original = filtered_original.convert('RGBA')
        
        # 합성 (배경 제거된 이미지를 필터링된 원본 위에 배치)
        result = Image.new('RGBA', filtered_original.size, (0, 0, 0, 0))
        result.paste(filtered_original, (0, 0), None)
        result.paste(bg_removed_image, (0, 0), bg_removed_image)
        
        return result.convert('RGB')
    
    elif filter_mode == "상품만 필터 적용":
        # 배경 제거된 이미지만 필터 적용
        cv_image = cv2.cvtColor(np.array(bg_removed_image), cv2.COLOR_RGBA2BGRA)
        # 알파 채널 분리
        b, g, r, a = cv2.split(cv_image)
        # BGR 이미지로 합치기
        bgr = cv2.merge([b, g, r])
        # 필터 적용
        bgr = adjust_temperature(bgr, temperature)
        bgr = adjust_brightness(bgr, brightness)
        bgr = adjust_contrast(bgr, contrast)
        bgr = adjust_saturation(bgr, saturation)
        # 다시 알파 채널과 합치기
        filtered_bgra = cv2.merge([bgr[:,:,0], bgr[:,:,1], bgr[:,:,2], a])
        filtered_bg_removed = Image.fromarray(cv2.cvtColor(filtered_bgra, cv2.COLOR_BGRA2RGBA))
        
        # 배경 제거된 이미지가 RGBA 모드인지 확인
        if bg_removed_image.mode != 'RGBA':
            logger.warning(f"세션 {session_id}: 배경 제거된 이미지가 RGBA 모드가 아닙니다.")
            # 필터만 적용하고 반환
            return filtered_bg_removed
        
        # 원본 이미지를 RGBA로 변환
        original_rgba = original_image.convert('RGBA')
        
        # 원본 이미지와 필터링된 배경 제거 이미지를 합성
        result = Image.new('RGBA', original_rgba.size, (0, 0, 0, 0))
        result.paste(original_rgba, (0, 0), None)
        result.paste(filtered_bg_removed, (0, 0), bg_removed_image)
        
        return result.convert('RGB')
    
    # 기본값은 전체 필터
    cv_image = cv2.cvtColor(np.array(original_image), cv2.COLOR_RGB2BGR)
    cv_image = adjust_temperature(cv_image, temperature)
    cv_image = adjust_brightness(cv_image, brightness)
    cv_image = adjust_contrast(cv_image, contrast)
    cv_image = adjust_saturation(cv_image, saturation)
    return Image.fromarray(cv2.cvtColor(cv_image, cv2.COLOR_BGR2RGB))

def download_image(image, session_id, format_type):
    """이미지를 지정된 형식으로 저장하고 경로 반환"""
    if image is None or session_id is None or session_id not in session_images:
        return None

    timestamp = get_korean_timestamp()
    
    # 원본 파일명 가져오기 (있는 경우)
    original_name = "이미지"
    if 'original_name' in session_images[session_id]:
        original_name = session_images[session_id]['original_name']

    if format_type == "JPG":
        file_name = f"[끝장AI]끝장필터_{original_name}_{timestamp}.jpg"
        format_str = "JPEG"
    else:  # PNG
        file_name = f"[끝장AI]끝장필터_{original_name}_{timestamp}.png"
        format_str = "PNG"

    # 파일 저장
    temp_file_path = tempfile.gettempdir() + "/" + file_name
    image.save(temp_file_path, format=format_str)
    return temp_file_path

def reset_filters(session_id):
    """필터 설정 초기화"""
    if session_id is None or session_id not in session_images:
        return None, 0.0, 1.0, 1.0, 1.0
    
    original_image = session_images[session_id].get('original_image')
    return original_image, 0.0, 1.0, 1.0, 1.0

def create_interface():
    css = """
    footer {
        visibility: hidden;
    }
    .download-button, .download-output, .reset-button {
        width: 100%;
    }
    .download-container {
        display: flex;
        flex-direction: column;
        align-items: center;
        width: 100%;
    }
    #gradio-app {
        margin: 0 !important;
        text-align: left !important;
        padding: 0 !important;
    }
    .gradio-container {
        max-width: 100% !important;
        margin: 0 !important;
        padding: 0 !important;
    }
    .download-button, .reset-button {
        background-color: black !important;
        color: white !important;
        border: none !important;
        padding: 10px !important;
        font-size: 16px !important;
        border-radius: 4px !important;
    }
    .reset-button {
        background-color: #666 !important;
    }
    .filter-mode {
        margin-top: 20px;
        margin-bottom: 20px;
    }
    .input-panel {
        padding: 20px;
    }
    .output-panel {
        padding: 20px;
        height: 100%;
        display: flex;
        flex-direction: column;
    }
    .output-image {
        flex-grow: 1;
        min-height: 400px;
    }
    .filter-section {
        margin: 15px 0;
        padding: 15px;
        border: 1px solid #eee;
        border-radius: 8px;
    }
    .session-id {
        display: none;
    }
    """

    with gr.Blocks(theme=gr.themes.Soft(
        primary_hue=gr.themes.Color(
            c50="#FFF7ED",
            c100="#FFEDD5",
            c200="#FED7AA",
            c300="#FDBA74",
            c400="#FB923C",
            c500="#F97316",
            c600="#EA580C",
            c700="#C2410C",
            c800="#9A3412",
            c900="#7C2D12",
            c950="#431407",
        ),
        secondary_hue="zinc",
        neutral_hue="zinc",
        font=("Pretendard", "sans-serif")
    ), css=css) as interface:

        # 세션 ID (사용자에게 보이지 않음)
        session_id = gr.Textbox(visible=False, elem_classes="session-id")

        with gr.Row():
            # 왼쪽 열: 입력 패널
            with gr.Column(scale=1, elem_classes="input-panel"):
                # 이미지 업로드
                input_image = gr.Image(type="pil", label="이미지 업로드")
                
                # 필터 모드 선택
                with gr.Group(elem_classes="filter-section"):
                    filter_mode = gr.Radio(
                        ["전체 필터", "배경만 필터 적용", "상품만 필터 적용"],
                        label="필터 적용 모드",
                        value="전체 필터",
                        elem_classes="filter-mode"
                    )
                
                # 조정 슬라이더 그룹 (중요 슬라이더만 남김)
                with gr.Group(elem_classes="filter-section"):
                    gr.Markdown("### 이미지 필터 설정")
                    
                    # 색온도를 가장 위로 배치
                    temperature_slider = gr.Slider(-1.0, 1.0, value=0.0, step=0.1, label="색온도 조절")
                    brightness_slider = gr.Slider(0.0, 2.0, value=1.0, step=0.1, label="밝기 조절")
                    contrast_slider = gr.Slider(0.5, 1.5, value=1.0, step=0.1, label="대비 조절")
                    saturation_slider = gr.Slider(0.0, 2.0, value=1.0, step=0.1, label="채도 조절")
                    
                    # 리셋 버튼
                    reset_button = gr.Button("필터 초기화", elem_classes="reset-button")

            # 오른쪽 열: 출력 패널
            with gr.Column(scale=1, elem_classes="output-panel"):
                # 필터 적용된 이미지 출력
                filtered_output = gr.Image(type="pil", label="필터 적용된 이미지", elem_classes="output-image")
                
                # 이미지 저장 옵션 및 다운로드
                with gr.Row():
                    with gr.Column(scale=1):
                        format_select = gr.Radio(
                            ["JPG", "PNG"],
                            label="저장 형식",
                            value="JPG"
                        )
                    
                    with gr.Column(scale=2):
                        download_button = gr.Button("이미지 변환하기", elem_classes="download-button")
                
                # 다운로드 링크
                download_output = gr.File(label="변환된 이미지 다운로드")

        # 이미지 업로드 처리 및 초기 출력 설정
        def update_image_and_session(image):
            if image is not None:
                # 파일 이름 저장 (가능한 경우)
                original_name = "이미지"
                if hasattr(image, 'name'):
                    original_name = image.name.split('.')[0]
                
                # 이미지 처리 및 세션 생성
                result_image, new_session_id = handle_upload(image)
                
                # 원본 파일명 저장
                if new_session_id in session_images:
                    session_images[new_session_id]['original_name'] = original_name
                
                return result_image, new_session_id
            return None, None
        
        input_image.change(
            fn=update_image_and_session,
            inputs=[input_image],
            outputs=[filtered_output, session_id]
        )
        
        # 필터 적용 함수
        def apply_filters(session_id, temperature, brightness, contrast, saturation, filter_mode):
            return process_image(session_id, temperature, brightness, contrast, saturation, filter_mode)
        
        # 필터 입력 변경 시 처리
        filter_inputs = [
            session_id,
            temperature_slider,
            brightness_slider,
            contrast_slider,
            saturation_slider,
            filter_mode
        ]
        
        # 슬라이더 변경 시 처리
        for input_component in [temperature_slider, brightness_slider, contrast_slider, saturation_slider]:
            input_component.change(
                fn=apply_filters,
                inputs=filter_inputs,
                outputs=filtered_output
            )
        
        # 필터 모드 변경 시 처리
        filter_mode.change(
            fn=apply_filters,
            inputs=filter_inputs,
            outputs=filtered_output
        )
        
        # 리셋 버튼 기능
        reset_button.click(
            fn=reset_filters,
            inputs=[session_id],
            outputs=[filtered_output, temperature_slider, brightness_slider, contrast_slider, saturation_slider]
        )
        
        # 다운로드 버튼 기능
        download_button.click(
            fn=download_image,
            inputs=[filtered_output, session_id, format_select],
            outputs=download_output
        )
        
        # 세션 정리 타이머 (30분 이상 된 세션 제거)
        def cleanup_sessions():
            current_time = time.time()
            to_remove = []
            
            for sess_id in session_images:
                if 'created_at' in session_images[sess_id]:
                    if current_time - session_images[sess_id]['created_at'] > 1800:  # 30분
                        to_remove.append(sess_id)
            
            for sess_id in to_remove:
                # 임시 파일 제거
                if 'bg_removed_path' in session_images[sess_id]:
                    path = session_images[sess_id]['bg_removed_path']
                    if os.path.exists(path):
                        os.unlink(path)
                
                # 세션 정보 제거
                del session_images[sess_id]
                
            # 10분마다 실행
            gr.set_interval(cleanup_sessions, 600)
        
        # 세션 정리 시작
        interface.load(cleanup_sessions)

    return interface

# 인터페이스 생성 및 실행
if __name__ == "__main__":
    logger.info("애플리케이션 시작")
    interface = create_interface()
    interface.queue()
    interface.launch()