Update app.py

app.py

@@ -3,8 +3,9 @@ import cv2
 import numpy as np
 from PIL import Image
 from io import BytesIO
+import os
 
-def process_image(image, 
+def process_image(file, 
                   convert_bw, 
                   denoise, 
                   sharpen, 
@@ -15,56 +16,78 @@ def process_image(image,
     """
     이미지 처리 함수
     """
-    # 원본 이미지 복사
+    if file is None:
+        return None, None, None, None
+
+    # Extract original filename
+    original_filename = os.path.basename(file.name)
+
+    # Open the image
+    image = Image.open(file).convert("RGB")
     original_image = image.copy()
-    
+
     # Convert PIL Image to OpenCV format
     img = np.array(image)
     img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
-    
+
     # 노이즈 제거
     if denoise:
         img = cv2.fastNlMeansDenoisingColored(img, None, 10, 10, 7, 21)
-    
+
     # 샤프닝
     if sharpen:
         kernel = np.array([[0, -1, 0],
                            [-1, 5,-1],
                            [0, -1, 0]])
         img = cv2.filter2D(img, -1, kernel)
-    
+
     # 감마 보정
     if gamma != 1.0:
         invGamma = 1.0 / gamma
         table = np.array([((i / 255.0) ** invGamma) * 255
                           for i in np.arange(256)]).astype("uint8")
         img = cv2.LUT(img, table)
-    
+
     # 밝기 조정
     if brightness != 0:
         img = cv2.convertScaleAbs(img, alpha=1, beta=brightness)
-    
+
     # 대비 조정
     if contrast != 1.0:
         img = cv2.convertScaleAbs(img, alpha=contrast, beta=0)
-    
+
     # 채도 조정
     if saturation != 1.0:
         hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV).astype(np.float32)
         hsv[...,1] = hsv[...,1] * saturation
         hsv[...,1] = np.clip(hsv[...,1], 0, 255)
         img = cv2.cvtColor(hsv.astype(np.uint8), cv2.COLOR_HSV2BGR)
-    
+
     # 흑백 변환
     if convert_bw:
         img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
         img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
-    
+
     # 변환된 이미지를 PIL 형식으로 변환
     transformed_image = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
     transformed_image = Image.fromarray(transformed_image)
-    
-    return transformed_image, original_image, transformed_image
+
+    # Prepare the download file
+    buf = BytesIO()
+    transformed_image.save(buf, format="JPEG")
+    buf.seek(0)
+
+    # 새로운 파일명 설정 (예: original_transformed.jpg)
+    if '.' in original_filename:
+        base, ext = original_filename.rsplit('.', 1)
+    else:
+        base, ext = original_filename, 'jpg'
+    new_filename = f"{base}_transformed.{ext}"
+
+    # gr.File을 위한 파일 이름 설정
+    buf.name = new_filename
+
+    return transformed_image, original_image, transformed_image, buf
 
 def blend_images(original, transformed, alpha):
     """
@@ -75,17 +98,6 @@ def blend_images(original, transformed, alpha):
     blended = Image.blend(original, transformed, alpha)
     return blended
 
-def download_image(transformed):
-    """
-    변환된 이미지를 다운로드할 수 있도록 하는 함수
-    """
-    if transformed is None:
-        return None
-    buf = BytesIO()
-    transformed.save(buf, format="JPEG")
-    buf.seek(0)
-    return buf
-
 # Gradio 인터페이스 구성
 def create_interface():
     with gr.Blocks() as demo:
@@ -93,7 +105,7 @@ def create_interface():
         
         with gr.Row():
             with gr.Column(scale=1):
-                input_image = gr.Image(type="pil", label="원본 이미지 업로드")
+                input_image = gr.File(type="file", label="원본 이미지 업로드")
                 convert_bw = gr.Checkbox(label="흑백으로 변환", value=True)
                 denoise = gr.Checkbox(label="노이즈 제거", value=False)
                 sharpen = gr.Checkbox(label="샤프닝", value=False)
@@ -109,20 +121,13 @@ def create_interface():
             
             with gr.Column(scale=1):
                 transformed_image = gr.Image(type="pil", label="변환된 이미지")
-                download_btn = gr.DownloadButton("JPG로 다운로드", file_name="transformed_image.jpg")
+                download_btn = gr.File(label="JPG로 다운로드")
         
-        # 이미지 처리 결과를 State에 저장하고 변환된 이미지 표시
+        # 이미지 처리 결과를 State에 저장하고 변환된 이미지 및 다운로드 파일 표시
         submit.click(
             fn=process_image, 
             inputs=[input_image, convert_bw, denoise, sharpen, gamma, brightness, contrast, saturation],
-            outputs=[transformed_image, original_state, transformed_state]
-        )
-        
-        # 다운로드 버튼에 변환된 이미지 연결
-        download_btn.click(
-            fn=download_image,
-            inputs=transformed_state,
-            outputs=download_btn
+            outputs=[transformed_image, original_state, transformed_state, download_btn]
         )
         
         gr.Markdown("### 이미지 비교")