Update app.py

app.py

@@ -61,7 +61,7 @@ def translate_prompt_to_english(prompt):
         logger.exception("번역 중 오류 발생:")
         return prompt
 
-def preprocess_prompt(prompt, image1, image2, image3):
+def preprocess_prompt(prompt, image1, image2, image3, background_options=None):
     """
     프롬프트를 처리하고 기능 명령을 해석
     """
@@ -103,8 +103,17 @@ def preprocess_prompt(prompt, image1, image2, image3):
     elif "4. 옷바꾸기" in prompt:
         prompt = "첫 번째 이미지의 인물 의상을 두 번째 이미지의 의상으로 변경해주세요. 의상의 스타일과 색상은 두 번째 이미지를 따르되, 신체 비율과 포즈는 첫 번째 이미지를 유지해주세요."
 
-    elif "5. 배경바꾸기" in prompt:
-        prompt = "첫 번째 이미지의 배경을 두 번째 이미지의 배경으로 변경해주세요. 첫 번째 이미지의 주요 피사체는 유지하고, 두 번째 이미지의 배경과 조화롭게 합성해주세요."
+    elif "5. 배경바꾸기" in prompt or "배경 변경" in prompt:
+        # 배경 옵션이 있으면 그것을 기반으로 프롬프트 생성
+        if background_options and all(background_options.values()):
+            product_category = background_options.get("product_category", "일반 상품")
+            background_type = background_options.get("background_type", "자연 배경")
+            background_style = background_options.get("background_style", "숲")
+            
+            # 상품 카테고리별 특화된 프롬프트 생성
+            prompt = f"첫 번째 이미지의 {product_category} 상품은 그대로 유지하면서, 배경을 {background_type} 중 {background_style}으로 자연스럽게 변경해주세요. 상품의 색상, 질감, 형태는 완벽하게 보존하고 고품질 상품 사진 스타일로 만들어주세요."
+        else:
+            prompt = "첫 번째 이미지의 배경을 두 번째 이미지의 배경으로 변경해주세요. 첫 번째 이미지의 주요 피사체는 유지하고, 두 번째 이미지의 배경과 조화롭게 합성해주세요."
 
     elif "상품합성" in prompt:
         prompt = "#1 상품 이미지에 #2 이미지 또는 #3 이미지를 [모든 이미지의 주요 요소를 포함하고, 특히 상품이 돋보이도록 조화롭게 통합하여]자연스럽게 합성하라."
@@ -164,7 +173,7 @@ def generate_with_images(prompt, images):
         logger.exception("이미지 생성 중 오류 발생:")
         return None, f"오류 발생: {str(e)}"
 
-def process_images_with_prompt(image1, image2, image3, prompt):
+def process_images_with_prompt(image1, image2, image3, prompt, background_options=None):
     """
     3개의 이미지와 프롬프트를 처리하여 최종 영어 프롬프트(final_prompt)를 생성한 후,
     API를 호출하여 결과 이미지를 반환합니다.
@@ -176,7 +185,7 @@ def process_images_with_prompt(image1, image2, image3, prompt):
             return None, "적어도 하나의 이미지를 업로드해주세요.", ""
 
         if prompt and prompt.strip():
-            processed_prompt = preprocess_prompt(prompt, image1, image2, image3)
+            processed_prompt = preprocess_prompt(prompt, image1, image2, image3, background_options)
             if re.search("[가-힣]", processed_prompt):
                 final_prompt = translate_prompt_to_english(processed_prompt)
             else:
@@ -198,13 +207,46 @@ def process_images_with_prompt(image1, image2, image3, prompt):
         logger.exception("이미지 처리 중 오류 발생:")
         return None, f"오류 발생: {str(e)}", prompt
 
-def process_and_show_prompt(image1, image2, image3, prompt):
+def process_and_show_prompt(image1, image2, image3, prompt, product_category=None, background_type=None, background_style=None):
     try:
-        result_img, status, final_prompt = process_images_with_prompt(image1, image2, image3, prompt)
-        return result_img, status, final_prompt
+        # 배경 옵션이 있을 경우 딕셔너리로 전달
+        background_options = None
+        if all([product_category, background_type, background_style]):
+            background_options = {
+                "product_category": product_category,
+                "background_type": background_type,
+                "background_style": background_style
+            }
+            
+        result_img, status, final_prompt = process_images_with_prompt(image1, image2, image3, prompt, background_options)
+        
+        # 선택된 옵션을 문자열로 구성
+        options_summary = ""
+        if background_options:
+            options_summary = f"제품군: {product_category}\n배경 유형: {background_type}\n세부 스타일: {background_style}"
+        
+        return result_img, status, final_prompt, options_summary
     except Exception as e:
         logger.exception("처리 중 오류 발생:")
-        return None, f"오류 발생: {str(e)}", prompt
+        return None, f"오류 발생: {str(e)}", prompt, ""
+
+# 선택 옵션 표시/���김 처리 함수
+def toggle_background_options(selection):
+    if selection == "배경바꾸기":
+        return gr.update(visible=True), gr.update(visible=True), gr.update(visible=True), gr.update(visible=True)
+    else:
+        return gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
+
+# 배경 유형에 따른 세부 스타일 옵션 업데이트
+def update_background_style_options(background_type):
+    if background_type == "자연 배경":
+        return gr.update(choices=["숲", "해변", "산", "들판", "계곡", "열대 정원", "사막", "눈 풍경"], value="숲")
+    elif background_type == "실내 배경":
+        return gr.update(choices=["거실", "주방", "서재", "화장대", "호텔 로비", "카페", "미니멀 공간"], value="거실")
+    elif background_type == "특수 배경":
+        return gr.update(choices=["도시 풍경", "루프탑", "해안가", "모던 스튜디오", "빈티지 공간", "미니어처 공간"], value="도시 풍경")
+    else:
+        return gr.update(choices=["선택하세요"], value="선택하세요")
 
 # Gradio 인터페이스
 with gr.Blocks() as demo:
@@ -217,6 +259,8 @@ with gr.Blocks() as demo:
         """
     )
     
+    current_selection = gr.State("")
+    
     with gr.Row():
         with gr.Column():
             with gr.Row():
@@ -228,50 +272,133 @@ with gr.Blocks() as demo:
                 placeholder="프롬프트를 입력하거나 비워두면 자동 합성됩니다.",
                 label="프롬프트 (선택 사항)"
             )
+            
             with gr.Row():
                 image_change_btn = gr.Button("이미지 변경")
                 text_remove_btn = gr.Button("글자지우기")
                 clothes_change_btn = gr.Button("옷바꾸기")
                 background_change_btn = gr.Button("배경바꾸기")
-                # 버튼 라벨과 클릭 시 반환되는 프롬프트 수정
                 composite_product_btn = gr.Button("상품합성")
+            
+            # 배경 변경 관련 옵션들 (기본적으로 숨김 상태)
+            with gr.Row(visible=False) as background_options_row:
+                gr.HTML("<h3>배경 변경 옵션</h3>")
+            
+            # 제품군 카테고리 드롭다운
+            product_category = gr.Dropdown(
+                choices=["화장품", "가방", "액세서리", "보석류", "홈 데코", "신발", "의류", "주방용품", "전자기기", "기타"],
+                label="제품군 선택",
+                value="기타",
+                visible=False
+            )
+            
+            # 배경 유형 드롭다운
+            background_type = gr.Dropdown(
+                choices=["자연 배경", "실내 배경", "특수 배경"],
+                label="배경 유형 선택",
+                value="자연 배경",
+                visible=False
+            )
+            
+            # 배경 세부 스타일 드롭다운 (배경 유형에 따라 동적 변경)
+            background_style = gr.Dropdown(
+                choices=["숲", "해변", "산", "들판", "계곡", "열대 정원", "사막", "눈 풍경"],
+                label="배경 세부 스타일 선택",
+                value="숲",
+                visible=False
+            )
+            
             submit_btn = gr.Button("이미지 생성", variant="primary")
+        
         with gr.Column():
             output_image = gr.Image(label="생성된 이미지")
             output_text = gr.Textbox(label="상태 메시지")
             prompt_display = gr.Textbox(label="사용된 프롬프트 (영어)", visible=True)
+            options_summary = gr.Textbox(label="선택된 옵션", visible=True)
     
     # 선택 옵션 버튼 클릭 시 프롬프트 입력란 업데이트 (한국어 문구)
     image_change_btn.click(
         fn=lambda: "#1 이미지의 [다른 모습]으로 변경하라.",
         inputs=[],
         outputs=prompt_input
+    ).then(
+        fn=lambda: "이미지 변경",
+        inputs=[],
+        outputs=current_selection
+    ).then(
+        fn=toggle_background_options,
+        inputs=[current_selection],
+        outputs=[background_options_row, product_category, background_type, background_style]
     )
+    
     text_remove_btn.click(
         fn=lambda: "#1 이미지의 [중국어를 모두]를 제거하라.",
         inputs=[],
         outputs=prompt_input
+    ).then(
+        fn=lambda: "글자지우기",
+        inputs=[],
+        outputs=current_selection
+    ).then(
+        fn=toggle_background_options,
+        inputs=[current_selection],
+        outputs=[background_options_row, product_category, background_type, background_style]
     )
+    
     clothes_change_btn.click(
         fn=lambda: "#1이미지에서 [신체 비율과 포즈는 유지한 체] 의상[스타일과 색상을]#2, #3으로 변경해주세요.",
         inputs=[],
         outputs=prompt_input
+    ).then(
+        fn=lambda: "옷바꾸기",
+        inputs=[],
+        outputs=current_selection
+    ).then(
+        fn=toggle_background_options,
+        inputs=[current_selection],
+        outputs=[background_options_row, product_category, background_type, background_style]
     )
+    
     background_change_btn.click(
-        fn=lambda: "#1 이미지의 배경을 #2 이미지의 배경으로 변경하라.",
+        fn=lambda: "배경 변경",
         inputs=[],
         outputs=prompt_input
+    ).then(
+        fn=lambda: "배경바꾸기",
+        inputs=[],
+        outputs=current_selection
+    ).then(
+        fn=toggle_background_options,
+        inputs=[current_selection],
+        outputs=[background_options_row, product_category, background_type, background_style]
     )
+    
     composite_product_btn.click(
         fn=lambda: "이미지 #1의 모델이 필요에 따라 자세나 표정을 유동적으로 변형한 후, 이미지 #2(또는 #3)의 상품을 최대한 원래 형태와 디테일을 유지하면서 자연스럽게 착용한 모습을 합성하라.",
         inputs=[],
         outputs=prompt_input
+    ).then(
+        fn=lambda: "상품합성",
+        inputs=[],
+        outputs=current_selection
+    ).then(
+        fn=toggle_background_options,
+        inputs=[current_selection],
+        outputs=[background_options_row, product_category, background_type, background_style]
+    )
+    
+    # 배경 유형에 따라 세부 스타일 옵션 업데이트
+    background_type.change(
+        fn=update_background_style_options,
+        inputs=[background_type],
+        outputs=[background_style]
     )
     
+    # 이미지 생성 버튼 클릭 이벤트
     submit_btn.click(
         fn=process_and_show_prompt,
-        inputs=[image1_input, image2_input, image3_input, prompt_input],
-        outputs=[output_image, output_text, prompt_display],
+        inputs=[image1_input, image2_input, image3_input, prompt_input, product_category, background_type, background_style],
+        outputs=[output_image, output_text, prompt_display, options_summary],
     )
     
     gr.Markdown(
@@ -281,10 +408,11 @@ with gr.Blocks() as demo:
         1. **자동 합성**: 이미지를 업로드하고 프롬프트를 비워두면 자동으로 합성됩니다.
         2. **이미지 참조**: #1, #2, #3으로 각 이미지를 참조할 수 있습니다.
         3. **선택 옵션**: 위의 버튼을 클릭하면 프롬프트 입력란에 한국어 문구로 입력하시면 됩니다.
+        4. **배경 변경**: 배경바꾸기 버튼을 클릭하면 추가 옵션이 나타납니다.
         
         > **팁**: 프롬프트를 직접 수정할 수도 있습니다.
         """
     )
 
 if __name__ == "__main__":
-    demo.launch(share=True)
+    demo.launch(share=True)