Update app.py

app.py

@@ -9,12 +9,35 @@ import time
 from google import genai
 from google.genai import types
 from dotenv import load_dotenv
+from transformers import pipeline  # 새로 추가된 import
+
 load_dotenv()
 
 # 기존 코드 유지 (로깅, 함수 등 모든 기능 코드)
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 logger = logging.getLogger(__name__)
 
+# 배경 제거 기능 추가
+def remove_background(image):
+    if image is None:
+        return None, "이미지가 업로드되지 않았습니다."
+    
+    try:
+        logger.info("배경 제거 시작")
+        
+        # 모델을 처음 로드할 때는 다운로드할 수 있으므로 시간이 걸릴 수 있음
+        pipe = pipeline("image-segmentation", model="briaai/RMBG-1.4", trust_remote_code=True)
+        
+        # 배경이 제거된 이미지 가져오기
+        output_image = pipe(image)
+        
+        logger.info("배경 제거 완료")
+        return output_image, "배경이 성공적으로 제거되었습니다."
+    except Exception as e:
+        logger.exception("배경 제거 중 오류 발생:")
+        return None, f"오류 발생: {str(e)}"
+
+
 def save_binary_file(file_name, data):
     with open(file_name, "wb") as f:
         f.write(data)
@@ -837,6 +860,33 @@ with gr.Blocks(css=custom_css) as demo:
                             gfpgan_output = gr.Image(type="numpy", label="업스케일 결과", elem_classes="image-container")
                     
                     gfpgan_btn = gr.Button('✨ 업스케일 및 복원', elem_classes="custom-button primary")
+        
+        # 새로운 탭: 배경 지우기 기능
+        with gr.TabItem("✂️ 배경 지우기", elem_classes="tab-content"):
+            with gr.Row():
+                with gr.Column(elem_classes="panel"):
+                    gr.HTML('<div class="section-title"><i class="fas fa-cut"></i> 이미지 배경 제거</div>')
+                    gr.HTML("""
+                    <p style="margin-bottom: 1rem; font-size: 1.05rem;">
+                         이미지에서 배경을 자동으로 제거하여 투명한 PNG 이미지로 변환합니다.
+                    </p>
+                    <div class="tip-box">
+                        <i class="fas fa-lightbulb"></i> <strong>팁:</strong> 배경과 주요 피사체 간의 대비가 뚜렷할수록 더 나은 결과를 얻을 수 있습니다.
+                    </div>
+                    """)
+                    
+                    with gr.Row():
+                        with gr.Column():
+                            gr.HTML('<div class="section-title"><i class="fas fa-upload"></i> 입력 이미지</div>')
+                            bg_remove_input = gr.Image(type="pil", label="업로드", elem_classes="image-container")
+                        
+                        with gr.Column():
+                            gr.HTML('<div class="section-title"><i class="fas fa-download"></i> 결과 이미지</div>')
+                            bg_remove_output = gr.Image(type="pil", label="배경 제거 결과", elem_classes="image-container", format="png")
+                    
+                    with gr.Row():
+                        bg_remove_status = gr.Textbox(label="상태", lines=2, elem_classes="custom-input")
+                        bg_remove_btn = gr.Button('✂️ 배경 제거하기', elem_classes="custom-button primary")
     
     # 버튼 이벤트 연결
     image_change_btn1.click(
@@ -895,6 +945,13 @@ with gr.Blocks(css=custom_css) as demo:
         inputs=gfpgan_input, 
         outputs=gfpgan_output
     )
+    
+    # 배경 제거 버튼 이벤트 연결
+    bg_remove_btn.click(
+        fn=remove_background,
+        inputs=bg_remove_input,
+        outputs=[bg_remove_output, bg_remove_status]
+    )
 
 demo.queue()
 demo.launch()