Update app.py

app.py

@@ -5,11 +5,29 @@ 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", "여기에_API_토큰을_입력하세요")
+os.environ["REPLICATE_API_TOKEN"] = REPLICATE_API_TOKEN
+
+# 세션별 이미지 저장용 딕셔너리
+session_images = {}
+
 def adjust_brightness(image, value):
     """이미지 밝기 조절"""
     value = float(value - 1) * 100  # 0-2 범위를 -100에서 +100으로 변환
@@ -50,96 +68,221 @@ def adjust_temperature(image, value):
     b = np.clip(b, 0, 255)
     return cv2.merge([b, g, r])
 
-def adjust_tint(image, value):
-    """이미지 색조 조절"""
-    hsv_image = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
-    h, s, v = cv2.split(hsv_image)
-    h = cv2.add(h, int(value))
-    h = np.clip(h, 0, 179)  # Hue 값은 0-179 범위
-    final_hsv = cv2.merge((h, s, v))
-    return cv2.cvtColor(final_hsv, cv2.COLOR_HSV2BGR)
-
-def adjust_exposure(image, value):
-    """이미지 노출 조절"""
-    enhancer = ImageEnhance.Brightness(Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB)))
-    img_enhanced = enhancer.enhance(1 + float(value) / 5.0)
-    return cv2.cvtColor(np.array(img_enhanced), cv2.COLOR_RGB2BGR)
-
-def adjust_vibrance(image, value):
-    """이미지 활기 조절"""
-    img = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
-    converter = ImageEnhance.Color(img)
-    factor = 1 + (float(value) / 100.0)
-    img = converter.enhance(factor)
-    return cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)
-
-def adjust_color_mixer_blues(image, value):
-    """이미지 컬러 믹서 (블루) 조절"""
-    b, g, r = cv2.split(image)
-    b = cv2.add(b, float(value))
-    b = np.clip(b, 0, 255)
-    return cv2.merge([b, g, r])
+def remove_background(image, session_id):
+    """배경 제거 기능 - Replicate API 사용"""
+    logger.info(f"세션 {session_id}: 배경 제거 시작")
+    
+    # 임시 파일로 이미지 저장
+    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로 배경 제거 요청
+        output = replicate.run(
+            "851-labs/background-remover:a029dff38972b5fda4ec5d75d7d1cd25aeff621d2cf4946a41055d7db66b80bc",
+            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 adjust_shadows(image, value):
-    """이미지 그림자 조절"""
-    pil_image = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
-    enhancer = ImageEnhance.Brightness(pil_image)
-    factor = 1 + (float(value) / 100.0)
-    pil_image = enhancer.enhance(factor)
-    return cv2.cvtColor(np.array(pil_image), cv2.COLOR_BGR2RGB)
+def get_korean_timestamp():
+    """한국 시간 타임스탬프 생성"""
+    korea_time = datetime.utcnow() + timedelta(hours=9)
+    return korea_time.strftime('%Y%m%d_%H%M%S')
 
-def process_image(image, brightness, contrast, saturation, temperature, tint, exposure, vibrance, color_mixer_blues, shadows):
-    """모든 조정 사항을 이미지에 적용"""
+def handle_upload(image, session_id=None):
+    """��미지 업로드 처리 및 배경 제거"""
     if image is None:
         return None
-        
-    # PIL 이미지를 OpenCV 형식으로 변환
-    image = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
     
-    # 조정 사항 순차 적용
-    image = adjust_brightness(image, brightness)
-    image = adjust_contrast(image, contrast)
-    image = adjust_saturation(image, saturation)
-    image = adjust_temperature(image, temperature)
-    image = adjust_tint(image, tint)
-    image = adjust_exposure(image, exposure)
-    image = adjust_vibrance(image, vibrance)
-    image = adjust_color_mixer_blues(image, color_mixer_blues)
-    image = adjust_shadows(image, shadows)
+    # 세션 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
     
-    # PIL 이미지로 다시 변환
-    return Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
+    # 배경 제거 실행
+    bg_removed_image = remove_background(original_image, session_id)
+    
+    # 필터 적용 전 기본 이미지 반환 (필터 없이 원본 이미지)
+    return original_image, session_id
 
-def download_image(image, input_image_name):
-    """이미지를 JPG 형식으로 저장하고 경로 반환"""
-    if image is None:
+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))
 
-    # 한국 시간 타임스탬프 생성 함수
-    from datetime import datetime, timedelta
-    def get_korean_timestamp():
-        korea_time = datetime.utcnow() + timedelta(hours=9)
-        return korea_time.strftime('%Y%m%d_%H%M%S')
+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()
-    if input_image_name and hasattr(input_image_name, 'name'):
-        base_name = input_image_name.name.split('.')[0]  # 파일 객체에서 이름 추출
-    else:
-        base_name = "이미지"
+    
+    # 원본 파일명 가져오기 (있는 경우)
+    original_name = "이미지"
+    if 'original_name' in session_images[session_id]:
+        original_name = session_images[session_id]['original_name']
 
-    file_name = f"[끝장AI]끝장필터_{base_name}_{timestamp}.jpg"
+    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="JPEG")
+    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 {
+    .download-button, .download-output, .reset-button {
         width: 100%;
     }
     .download-container {
@@ -149,110 +292,216 @@ def create_interface():
         width: 100%;
     }
     #gradio-app {
-        margin: 0 !important; /* 모든 방향 여백 제거 */
-        text-align: left !important; /* 왼쪽 정렬 강제 */
-        padding: 20px !important; /* 패딩 추가 */
+        margin: 0 !important;
+        text-align: left !important;
+        padding: 0 !important;
     }
     .gradio-container {
-        max-width: 100% !important; /* 가로 폭 전체 사용 */
-        margin-left: 0 !important; /* 왼쪽 정렬 */
-        padding: 20px !important; /* 패딩 추가 */
+        max-width: 100% !important;
+        margin: 0 !important;
+        padding: 0 !important;
     }
-    .download-button {
+    .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",  # 가장 밝은 주황
+            c50="#FFF7ED",
             c100="#FFEDD5",
             c200="#FED7AA",
             c300="#FDBA74",
             c400="#FB923C",
-            c500="#F97316",  # 기본 주황
+            c500="#F97316",
             c600="#EA580C",
             c700="#C2410C",
             c800="#9A3412",
-            c900="#7C2D12",  # 가장 어두운 주황
+            c900="#7C2D12",
             c950="#431407",
         ),
-        secondary_hue="zinc",  # 모던한 느낌의 회색 계열
+        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():
-            # 왼쪽 열: 비율 3
-            with gr.Column(scale=3):
+            # 왼쪽 열: 입력 패널
+            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")
 
-                # 조정 슬라이더
-                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="채도 조절")
-                temperature_slider = gr.Slider(-1.0, 1.0, value=0.0, step=0.1, label="색온도 조절")
-                tint_slider = gr.Slider(-100, 100, value=0, step=1, label="색조 조절")
-                exposure_slider = gr.Slider(-5.0, 5.0, value=0.0, step=0.1, label="노출 조절")
-                vibrance_slider = gr.Slider(-100.0, 100.0, value=0.0, step=1.0, label="활기 조절")
-                color_mixer_blues_slider = gr.Slider(-100.0, 100.0, value=0.0, step=1.0, label="컬러 믹서 (블루)")
-                shadows_slider = gr.Slider(-100.0, 100.0, value=0.0, step=1.0, label="그림자 조절")
+            # 오른쪽 열: 출력 패널
+            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="변환된 이미지 다운로드")
 
-            # 오른쪽 열: 비율 7
-            with gr.Column(scale=7):
-                # 처리된 이미지 출력
-                output_image = gr.Image(type="pil", label="처리된 이미지")
-
-                # 변환된 이미지 다운로드 버튼
-                with gr.Row(elem_classes="download-container"):
-                    download_button = gr.Button("JPG로 변환하기", elem_classes="download-button")
-                with gr.Row(elem_classes="download-container"):
-                    download_output = gr.File(label="JPG 이미지 다운로드", elem_classes="download-output")
-
-        # 이미지 처리 함수 연결
-        inputs = [
-            input_image,
-            brightness_slider,
-            contrast_slider,
-            saturation_slider,
+        # 이미지 업로드 처리 및 초기 출력 설정
+        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,
-            tint_slider,
-            exposure_slider,
-            vibrance_slider,
-            color_mixer_blues_slider,
-            shadows_slider
-        ]
-
-        input_components = [
             brightness_slider,
             contrast_slider,
             saturation_slider,
-            temperature_slider,
-            tint_slider,
-            exposure_slider,
-            vibrance_slider,
-            color_mixer_blues_slider,
-            shadows_slider
+            filter_mode
         ]
-
-        for input_component in input_components:
+        
+        # 슬라이더 변경 시 처리
+        for input_component in [temperature_slider, brightness_slider, contrast_slider, saturation_slider]:
             input_component.change(
-                fn=process_image,
-                inputs=inputs,
-                outputs=output_image
+                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=[output_image, input_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
 
@@ -261,4 +510,4 @@ if __name__ == "__main__":
     logger.info("애플리케이션 시작")
     interface = create_interface()
     interface.queue()
-    interface.launch()
+    interface.launch()