Update app.py

app.py

@@ -6,7 +6,12 @@ import gfpgan
 import tempfile
 import time
 import gradio as gr
+import sys
+from torchvision.transforms import functional
+from PIL import Image
 
+# 참조코드에서 사용된 모듈 임포트 수정
+sys.modules["torchvision.transforms.functional_tensor"] = functional
 
 class Predictor:
     def __init__(self):
@@ -23,63 +28,164 @@ class Predictor:
             os.system(
                 'wget https://huggingface.co/ashleykleynhans/inswapper/resolve/main/inswapper_128.onnx'
             )
-        os.chdir('..')
+        os.chdir('..')  # 디렉토리 변경 완료
 
-        """Load the model into memory to make running multiple predictions efficient"""
+        """💎 Load the model into memory to make running multiple predictions efficient"""
         self.face_swapper = insightface.model_zoo.get_model('models/inswapper_128.onnx',
                                                             providers=onnxruntime.get_available_providers())
+        # self.face_swapper.prepare(ctx_id=0, det_size=(640, 640))  # 이 줄을 제거합니다.
+
         self.face_enhancer = gfpgan.GFPGANer(model_path='models/GFPGANv1.4.pth', upscale=1)
         self.face_analyser = insightface.app.FaceAnalysis(name='buffalo_l')
         self.face_analyser.prepare(ctx_id=0, det_size=(640, 640))
 
-    def get_face(self, img_data):
-        analysed = self.face_analyser.get(img_data)
-        try:
-            largest = max(analysed, key=lambda x: (x.bbox[2] - x.bbox[0]) * (x.bbox[3] - x.bbox[1]))
-            return largest
-        except:
-            print("No face found")
-            return None
+    def get_face_image(self, img_data, face):
+        # 얼굴 영역을 잘라내기
+        x1, y1, x2, y2 = [int(coord) for coord in face.bbox]
+        face_img = img_data[y1:y2, x1:x2]
+        return face_img
 
-    def predict(self, input_image, swap_image):
-        """Run a single prediction on the model"""
+    def predict(self, input_image_path, swap_image_path):
+        """🧶 Run a single prediction on the model"""
         try:
-            frame = cv2.imread(input_image.name)
-            face = self.get_face(frame)
-            source_face = self.get_face(cv2.imread(swap_image.name))
-            try:
-                print(frame.shape, face.shape, source_face.shape)
-            except:
-                print("printing shapes failed.")
-            result = self.face_swapper.get(frame, face, source_face, paste_back=True)
+            frame = cv2.imread(input_image_path)
+            if frame is None:
+                print("❌ Target image could not be read.")
+                return None
+            analysed = self.face_analyser.get(frame)
+            if not analysed:
+                print("❌ No face found in target image.")
+                return None
+            face = max(analysed, key=lambda x: (x.bbox[2] - x.bbox[0]) * (x.bbox[3] - x.bbox[1]))
+            target_face_img = self.get_face_image(frame, face)
+
+            swap_frame = cv2.imread(swap_image_path)
+            if swap_frame is None:
+                print("❌ Swap image could not be read.")
+                return None
+            swap_analysed = self.face_analyser.get(swap_frame)
+            if not swap_analysed:
+                print("❌ No face found in swap image.")
+                return None
+            swap_face = max(swap_analysed, key=lambda x: (x.bbox[2] - x.bbox[0]) * (x.bbox[3] - x.bbox[1]))
+            swap_face_img = self.get_face_image(swap_frame, swap_face)
+
+            # 얼굴 교체 수행
+            result = self.face_swapper.get(frame, face, swap_face, paste_back=True)
 
+            # 얼굴 향상 수행
             _, _, result = self.face_enhancer.enhance(
                 result,
                 paste_back=True
             )
-            out_path = tempfile.mkdtemp() + f"/{str(int(time.time()))}.jpg"
+            out_path = os.path.join(tempfile.mkdtemp(), f"{str(int(time.time()))}.jpg")
             cv2.imwrite(out_path, result)
             return out_path
         except Exception as e:
             print(f"{e}")
             return None
 
-
-# Instantiate the Predictor class
+# Predictor 클래스 인스턴스 생성
 predictor = Predictor()
-title = "Swap Faces Using Our Model!!!"
-
-# Create Gradio Interface
-iface = gr.Interface(
-    fn=predictor.predict,
-    inputs=[
-        gr.Image(type="filepath", label="Target Image"),
-        gr.Image(type="filepath", label="Swap Image")
-    ],
-    outputs=gr.Image(type="filepath", label="Result"),
-    title=title,
-    examples=[["input.jpg", "swap img.jpg"]])
-
-
-# Launch the Gradio Interface
-iface.launch()
+
+# CSS 및 테마 설정
+css = """
+/* "Swap Faces" 버튼 스타일 */
+button#swap-button {
+    background-color: #FB923C !important; /* 주황색 배경 */
+    color: white !important; /* 흰색 글씨 */
+}
+/* "이미지 다운로드 (JPG)" 버튼 스타일 */
+button#download-button {
+    background-color: #FB923C !important; /* 주황색 배경 */
+    color: white !important; /* 흰색 글씨 */
+}
+/* 필요에 따라 추가적인 스타일을 여기에 작성할 수 있습니다 */
+"""
+
+demo_theme = gr.themes.Soft(
+    primary_hue=gr.themes.Color(
+        c50="#FFF7ED",
+        c100="#FFEDD5",
+        c200="#FED7AA",
+        c300="#FDBA74",
+        c400="#FB923C",
+        c500="#F97316",
+        c600="#EA580C",
+        c700="#C2410C",
+        c800="#9A3412",
+        c900="#7C2D12",
+        c950="#431407",
+    ),
+    secondary_hue="zinc",
+    neutral_hue="zinc",
+    font=("Pretendard", "sans-serif")
+)
+
+# JPG 다운로드 기능 구현
+def save_as_jpg(file_path):
+    try:
+        if file_path is None:
+            return None
+        # 파일 경로를 받아 PIL 이미지로 변환
+        img = Image.open(file_path)
+        with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as tmp:
+            img.save(tmp, format="JPEG")
+            tmp_path = tmp.name
+        return tmp_path  # 파일 경로 반환
+    except Exception as error:
+        print(f"Error saving as JPG: {error}")
+        return None
+
+# Clear 함수: 입력 및 출력 초기화
+def clear_all():
+    return [None, None, None]
+
+# Gradio Interface 구성
+with gr.Blocks(theme=demo_theme, css=css) as demo:
+    with gr.Row():
+        # 왼쪽 섹션: 입력
+        with gr.Column(scale=1):
+            target_image = gr.Image(
+                type="filepath",
+                label="얼굴을 변경할 이미지"
+            )
+            swap_image = gr.Image(
+                type="filepath",
+                label="교체할 얼굴"
+            )
+            swap_button = gr.Button("얼굴 교체", elem_id="swap-button")
+            clear_button = gr.Button("리셋하기")
+        
+        # 오른쪽 섹션: 출력
+        with gr.Column(scale=1):
+            result_image = gr.Image(
+                type="filepath",
+                label="결과 이미지"
+            )
+            download_jpg_button = gr.Button("JPG로 변환하기", elem_id="download-button")
+            download_file = gr.File(label="JPG 이미지 다운받기")
+    
+    # 버튼 클릭 시 예측 함수 호출
+    swap_button.click(
+        fn=predictor.predict,
+        inputs=[target_image, swap_image],
+        outputs=result_image
+    )
+
+    # 리셋하기 버튼 클릭 시 입력 및 출력 이미지 초기화
+    clear_button.click(
+        fn=clear_all,
+        inputs=None,
+        outputs=[target_image, swap_image, result_image]
+    )
+
+    # JPG 다운로드 버튼 클릭 시 파일 생성
+    download_jpg_button.click(
+        fn=save_as_jpg,
+        inputs=result_image,
+        outputs=download_file
+    )
+
+# Gradio Interface 실행
+demo.launch()