Update app.py

app.py

@@ -339,12 +339,13 @@ def on_change_bbox(prompts: dict[str, Any] | None):
 def on_change_prompt(img: Image.Image | None, prompt: str | None, bg_prompt: str | None = None):
     return gr.update(interactive=bool(img and prompt))
 
-
-        
+# process_prompt 함수 수정
 def process_prompt(img: Image.Image, prompt: str, bg_prompt: str | None = None, 
                   aspect_ratio: str = "1:1", position: str = "bottom-center", 
-                  scale_percent: float = 100) -> tuple[Image.Image, Image.Image]:
+                  scale_percent: float = 100, text_params: dict | None = None) -> tuple[Image.Image, Image.Image]:
     try:
+        
+
         if img is None or prompt.strip() == "":
             raise gr.Error("Please provide both image and prompt")
         
@@ -374,12 +375,18 @@ def process_prompt(img: Image.Image, prompt: str, bg_prompt: str | None = None,
                 return results[1], results[2]
         
         return results[1], results[2]
+
+        if text_params and text_params.get('text'):
+            combined = add_text_to_image(combined, text_params)
+        
+        return combined, results[2]
     except Exception as e:
         print(f"Error in process_prompt: {str(e)}")
         raise gr.Error(str(e))
     finally:
         clear_memory()
-        
+
+
 def process_bbox(img: Image.Image, box_input: str) -> tuple[Image.Image, Image.Image]:
     try:
         if img is None or box_input.strip() == "":
@@ -482,9 +489,63 @@ button.primary:hover {
 }
 """
 
-# UI 구성
-# UI 구성 부분에서 process_btn을 위로 이동하고 position_grid.click 부분 제거
+def add_text_with_stroke(draw, text, x, y, font, text_color, stroke_width):
+    """텍스트에 외곽선을 추가하는 헬퍼 함수"""
+    for adj_x in range(-stroke_width, stroke_width + 1):
+        for adj_y in range(-stroke_width, stroke_width + 1):
+            draw.text((x + adj_x, y + adj_y), text, font=font, fill=text_color)
 
+def add_text_to_image(image, text_params):
+    """이미지에 텍스트를 추가하는 함수"""
+    if not text_params.get('text'):
+        return image
+    
+    if image.mode != 'RGBA':
+        image = image.convert('RGBA')
+    
+    txt_overlay = Image.new('RGBA', image.size, (255, 255, 255, 0))
+    draw = ImageDraw.Draw(txt_overlay)
+    
+    try:
+        font = ImageFont.truetype("DejaVuSans.ttf", text_params['font_size'])
+    except:
+        try:
+            font = ImageFont.truetype("arial.ttf", text_params['font_size'])
+        except:
+            font = ImageFont.load_default()
+    
+    color_map = {
+        'White': (255, 255, 255),
+        'Black': (0, 0, 0),
+        'Red': (255, 0, 0),
+        'Green': (0, 255, 0),
+        'Blue': (0, 0, 255),
+        'Yellow': (255, 255, 0),
+        'Purple': (128, 0, 128)
+    }
+    
+    rgb_color = color_map.get(text_params['color'], (255, 255, 255))
+    text_color = (*rgb_color, text_params['opacity'])
+    
+    text_bbox = draw.textbbox((0, 0), text_params['text'], font=font)
+    text_width = text_bbox[2] - text_bbox[0]
+    text_height = text_bbox[3] - text_bbox[1]
+    
+    x = int((image.width - text_width) * (text_params['x_position'] / 100))
+    y = int((image.height - text_height) * (text_params['y_position'] / 100))
+    
+    add_text_with_stroke(
+        draw,
+        text_params['text'],
+        x,
+        y,
+        font,
+        text_color,
+        text_params['thickness']
+    )
+    
+    return Image.alpha_composite(image, txt_overlay)
+    
 # UI 구성
 with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
     gr.HTML("""
@@ -582,6 +643,53 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
                     height=256
                 )
 
+    
+    with gr.Row():
+        with gr.Column(scale=1):
+            # 텍스트 입력 섹션 추가
+            with gr.Group(visible=True) as text_group:
+                text_input = gr.Textbox(label="Text to Add", placeholder="Enter text...")
+                font_size = gr.Slider(minimum=10, maximum=800, value=400, step=10, 
+                                    label="Font Size")
+                thickness = gr.Slider(minimum=0, maximum=20, value=0, step=1,
+                                    label="Text Thickness")
+                color_dropdown = gr.Dropdown(
+                    choices=["White", "Black", "Red", "Green", "Blue", "Yellow", "Purple"],
+                    value="White",
+                    label="Text Color"
+                )
+                opacity_slider = gr.Slider(minimum=0, maximum=255, value=255, step=1, 
+                                         label="Opacity")
+                text_x_position = gr.Slider(minimum=0, maximum=100, value=50, step=1, 
+                                          label="Text X Position (%)")
+                text_y_position = gr.Slider(minimum=0, maximum=100, value=50, step=1, 
+                                          label="Text Y Position (%)")
+    
+    # 이벤트 바인딩 수정
+    process_btn.click(
+        fn=process_prompt,
+        inputs=[
+            input_image,
+            text_prompt,
+            bg_prompt,
+            aspect_ratio,
+            position,
+            scale_slider,
+            # 텍스트 파라미터 추가
+            gr.State(lambda: {
+                'text': text_input.value,
+                'font_size': font_size.value,
+                'thickness': thickness.value,
+                'color': color_dropdown.value,
+                'opacity': opacity_slider.value,
+                'x_position': text_x_position.value,
+                'y_position': text_y_position.value
+            })
+        ],
+        outputs=[combined_image, extracted_image]
+    )
+
+    
     # Event bindings
     input_image.change(
         fn=update_process_button,