Create predict.py

predict.py

@@ -0,0 +1,229 @@
+# Prediction interface for Cog ⚙️
+# https://github.com/replicate/cog/blob/main/docs/python.md
+import shutil
+import gradio as gr
+
+# from cog import BasePredictor, Input, Path
+
+import insightface
+import onnxruntime
+from insightface.app import FaceAnalysis
+import cv2
+import gfpgan
+import tempfile
+import time
+import uuid
+from typing import Any, Union
+from loggers import logger, request_id as _request_id
+import ssl
+from datetime import datetime
+import traceback
+import torch
+import os
+import requests
+import subprocess
+import sys
+from PIL import Image
+import numpy as np
+
+ssl._create_default_https_context = ssl._create_unverified_context
+
+if sys.platform == 'darwin':
+    cache_file_dir = '/tmp/file'
+else:
+    cache_file_dir = '/src/file'
+os.makedirs(cache_file_dir, exist_ok=True)
+
+
+def img_url_to_local_path(img_url, file_path=None):
+    filename = img_url.split('/')[-1]
+    max_count = 3
+    count = 0
+    if file_path is None:
+        temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=filename)
+        temp_file_name = temp_file.name
+    else:
+        temp_file_name = file_path
+    while True:
+        count += 1
+        try:
+            res = requests.get(img_url, timeout=60)
+            res.raise_for_status()
+            with open(temp_file_name, "wb") as f:
+                f.write(res.content)
+            return temp_file_name
+        except Exception as e:
+            logger.error(e)
+        if count >= max_count:
+            msg = f'request {max_count} time url: {img_url} failed, please check'
+            logger.error(msg)
+            raise Exception(msg)
+
+
+def delete_files_day_ago(cache_days=10):
+    command = f"find {cache_file_dir} -type f -ctime +{cache_days} -exec rm {{}} \;"
+    result = subprocess.run(command, shell=True, capture_output=True, text=True)
+    logger.info(result.stdout)
+
+
+def image_format_by_path(image_path):
+    image = Image.open(image_path)
+    image_format = image.format
+    if not image_format:
+        image_format = 'jpg'
+    elif image_format == "JPEG":
+        image_format = 'jpg'
+    else:
+        image_format = image_format.lower()
+    return image_format
+
+
+def local_file_for_url(url, cache_days=10):
+    filename = url.split('/')[-1]
+    _, ext = filename.split('.')
+    file_path = f'{cache_file_dir}/{filename}'
+    if not os.path.exists(file_path):
+        img_url_to_local_path(url, file_path)
+        logger.info(f'download file to {file_path}')
+        delete_files_day_ago(cache_days)
+    else:
+        logger.info(f'cache file {file_path}')
+    return file_path
+
+
+class Predictor:
+    def __init__(self):
+        self.det_thresh = 0.1
+
+    def setup(self):
+        self.face_swapper = insightface.model_zoo.get_model('cache/inswapper_128.onnx', providers=onnxruntime.get_available_providers())
+        self.face_enhancer = gfpgan.GFPGANer(model_path='cache/GFPGANv1.4.pth', upscale=1)
+        self.face_analyser = FaceAnalysis(name='buffalo_l')
+
+    def get_face(self, img_data, image_type='target'):
+        try:
+            logger.info(self.det_thresh)
+            self.face_analyser.prepare(ctx_id=0, det_thresh=0.5)
+            if image_type == 'source':
+                self.face_analyser.prepare(ctx_id=0, det_thresh=self.det_thresh)
+            analysed = self.face_analyser.get(img_data)
+            logger.info(f'face num: {len(analysed)}')
+            if len(analysed) == 0:
+                msg = 'no face'
+                logger.error(msg)
+                raise Exception(msg)
+            largest = max(analysed, key=lambda x: (x.bbox[2] - x.bbox[0]) * (x.bbox[3] - x.bbox[1]))
+            return largest
+        except Exception as e:
+            logger.error(str(e))
+            raise Exception(str(e))
+
+    def enhance_face(self, target_face, target_frame, weight=0.5):
+        start_x, start_y, end_x, end_y = map(int, target_face['bbox'])
+        padding_x = int((end_x - start_x) * 0.5)
+        padding_y = int((end_y - start_y) * 0.5)
+        start_x = max(0, start_x - padding_x)
+        start_y = max(0, start_y - padding_y)
+        end_x = max(0, end_x + padding_x)
+        end_y = max(0, end_y + padding_y)
+        temp_face = target_frame[start_y:end_y, start_x:end_x]
+        if temp_face.size:
+            _, _, temp_face = self.face_enhancer.enhance(
+                temp_face,
+                paste_back=True,
+                weight=weight
+            )
+            target_frame[start_y:end_y, start_x:end_x] = temp_face
+        return target_frame
+
+    def predict(
+            self,
+            source_image_path,
+            target_image_path,
+            enhance_face,
+            # request_id: str = Input(description="request_id", default=""),
+            # det_thresh: float = Input(description="det_thresh default 0.1", default=0.1),
+            # local_target: Path = Input(description="local target image", default=None),
+            # local_source: Path = Input(description="local source image", default=None),
+            # cache_days: int = Input(description="cache days default 10", default=10),
+            # weight: float = Input(description="weight default 0.5", default=0.5)
+
+    ) -> Any:
+        """Run a single prediction on the model"""
+        request_id = None
+        det_thresh = 0.1
+        cache_days = 10
+        weight = 0.5
+
+        device = 'cuda' if torch.cuda.is_available() else 'mps'
+        logger.info(f'device: {device}, det_thresh:{det_thresh}')
+
+        try:
+            self.det_thresh = det_thresh
+            start_time = time.time()
+            if not request_id:
+                request_id = str(uuid.uuid4())
+            _request_id.set(request_id)
+            frame = cv2.imread(str(target_image_path))
+            source_frame = cv2.imread(str(source_image_path))
+            source_face = self.get_face(source_frame, image_type='source')
+            target_face = self.get_face(frame)
+            try:
+                logger.info(f'{frame.shape}, {target_face.shape}, {source_face.shape}')
+            except Exception as e:
+                logger.error(f"printing shapes failed,  error:{str(e)}")
+                raise Exception(str(e))
+            ext = image_format_by_path(target_image_path)
+            size = os.path.getsize(target_image_path)
+            logger.info(f'origin {size/1024}k')
+            result = self.face_swapper.get(frame, target_face, source_face, paste_back=True)
+            if enhance_face:
+                result = self.enhance_face(target_face, result, weight)
+            # _, _, result = self.face_enhancer.enhance(
+            #     result,
+            #     paste_back=True
+            # )
+            out_path = f"{tempfile.mkdtemp()}/{uuid.uuid4()}.{ext}"
+            cv2.imwrite(str(out_path), result)
+            return Image.open(out_path) 
+
+            size = os.path.getsize(out_path)
+            logger.info(f'result {size / 1024}k')
+            cost_time = time.time() - start_time
+            logger.info(f'total time: {cost_time * 1000} ms')
+            data = {'code': 200, 'msg': 'succeed', 'image': out_path, 'status': 'succeed'}
+            return data
+        except Exception as e:
+            logger.error(traceback.format_exc())
+            data = {'code': 500, 'msg': str(e), 'image': '', 'status': 'failed'}
+            logger.error(f"{str(e)}")
+            return data
+
+def swap_faces(source_image_path, target_image_path, enhance_face):
+    predictor = Predictor()
+    predictor.setup()
+    return predictor.predict(
+        source_image_path,
+        target_image_path,
+        enhance_face
+    )
+
+if __name__ == "__main__":
+    demo = gr.Interface(
+      fn=swap_faces,
+      inputs=[
+          gr.Image(type="filepath"),
+          gr.Image(type="filepath"),
+          gr.Checkbox(label="Enhance Face", value=True),
+        #   gr.Checkbox(label="Enhance Frame", value=True),
+      ],
+      outputs=[
+          gr.Image(
+            type="pil",
+            show_download_button=True,
+          )
+      ],
+      title="Swap Faces",
+      allow_flagging="never"
+    )
+    demo.launch()