Upload folder using huggingface_hub

inference.py

@@ -293,7 +293,7 @@ def pose(task: Task, s3_outkey: str = "_pose", poses: Optional[list] = None):
     lora_patcher = lora_style.get_patcher(controlnet.pipe2, task.get_style())
     lora_patcher.patch()
 
-    try:
+    if task.get_pose_coordinates():
         infered_pose = pose_detector.transform(
             image=task.get_imageUrl(),
             client_coordinates=task.get_pose_coordinates(),
@@ -301,8 +301,7 @@ def pose(task: Task, s3_outkey: str = "_pose", poses: Optional[list] = None):
             height=task.get_height(),
         )
         poses = [infered_pose] * num_return_sequences
-    except Exception as e:
-        print("Failed to detect pose, using Open Pose detector", e)
+    else:
         poses = [controlnet.detect_pose(task.get_imageUrl())] * num_return_sequences
 
     images, has_nsfw = controlnet.process_pose(

internals/pipelines/pose_detector.py

@@ -10,7 +10,6 @@ from models.pose.body import Body
 
 
 class PoseDetector:
-    # __det_model = "https://comic-assets.s3.ap-south-1.amazonaws.com/models/faster_rcnn_r50_fpn_1x_coco_20200130-047c8118.pth"
     __pose_model = (
         "https://comic-assets.s3.ap-south-1.amazonaws.com/models/body_pose_model.pth"
     )
@@ -41,11 +40,17 @@ class PoseDetector:
             image = download_image(image)
 
         infer_coordinates = self.infer(image, width, height)
+        candidate_list = self.make_pose_from_subset(
+            infer_coordinates["candidate"], infer_coordinates["subset"]
+        )
+        # hard check only one person
+        infer_coordinates["candidate"] = candidate_list[0]
+
         if client_coordinates and client_coordinates["candidate"]:
             client_coordinates = self.resize_coordinates(
                 client_coordinates, 384, 384, width, height
             )
-            infer_coordinates = self.map_head_to_body(
+            infer_coordinates = self.map_coordinates(
                 client_coordinates, infer_coordinates
             )
 
@@ -90,9 +95,10 @@ class PoseDetector:
                 )
 
         for i, point in enumerate(points):
-            x = point[0]
-            y = point[1]
-            draw.ellipse((x - 3, y - 3, x + 3, y + 3), fill=self.__points_color[i])
+            x = safe_index(point, 0)
+            y = safe_index(point, 1)
+            if x and y:
+                draw.ellipse((x - 3, y - 3, x + 3, y + 3), fill=self.__points_color[i])
 
         return image
 
@@ -111,9 +117,9 @@ class PoseDetector:
 
         candidate = [item[:2] for item in candidate]
 
-        return {"candidate": candidate[:18], "subset": subset[:18]}
+        return {"candidate": candidate, "subset": subset}
 
-    def map_head_to_body(
+    def map_coordinates(
         self, client_coordinates: dict, infer_coordinates: dict
     ) -> dict:
         client_points = client_coordinates["candidate"]
@@ -125,12 +131,34 @@ class PoseDetector:
         dx = i_neck[0] - c_neck[0]
         dy = i_neck[1] - c_neck[1]
 
-        for i in range(2, 15):
+        # Considering client coordinates truthy and translate it to the position of infered coordinates
+        for i in range(len(client_points)):
             point = client_points[i - 1]
             infer_points[i - 1] = [point[0] + dx, point[1] + dy]
 
         return {"candidate": infer_points, "subset": infer_coordinates["subset"]}
 
+    def make_pose_from_subset(self, candidate, subset):
+        "Maps pose coordinates for subset"
+
+        def make_pose_from_subset_item(candidate, subset_item):
+            pose = []
+            for j in range(18):
+                i = int(subset_item[j])
+                pose.append(
+                    None
+                    if i < 0 or not safe_index(candidate, i)
+                    else list(map(lambda x: x, candidate[i]))
+                )
+            return pose
+
+        return list(
+            map(
+                lambda subset_item: make_pose_from_subset_item(candidate, subset_item),
+                subset,
+            )
+        )
+
     def __convert_keypoints(self, keypoints):
         return [keypoints[i] for i in self.__kim]
 

internals/pipelines/upscaler.py

@@ -7,16 +7,21 @@ import cv2
 import numpy as np
 from basicsr.archs.rrdbnet_arch import RRDBNet
 from basicsr.utils.download_util import load_file_from_url
+from gfpgan import GFPGANer
 from PIL import Image
 from realesrgan import RealESRGANer
 
 import internals.util.image as ImageUtil
 from internals.util.commons import download_image
+from internals.util.config import get_root_dir
 
 
 class Upscaler:
     __model_esrgan_url = "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth"
     __model_esrgan_anime_url = "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.4/RealESRGAN_x4plus_anime_6B.pth"
+    __model_gfpgan_url = (
+        "https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.3.pth"
+    )
 
     __loaded = False
 
@@ -31,6 +36,9 @@ class Upscaler:
         self.__model_path_anime = self.__preload_model(
             self.__model_esrgan_anime_url, download_dir
         )
+        self.__model_path_gfpgan = self.__preload_model(
+            self.__model_gfpgan_url, download_dir
+        )
         self.__loaded = True
 
     def upscale(self, image: Union[str, Image.Image], resize_dimension: int) -> bytes:
@@ -88,13 +96,26 @@ class Upscaler:
         if isinstance(image, Image.Image):
             image = ImageUtil.to_bytes(image)
 
-        upsampler = RealESRGANer(
-            scale=4, model_path=model_path, model=rrbdnet, half="fp16", gpu_id="0"
-        )
         image_array = np.frombuffer(image, dtype=np.uint8)
         input_image = cv2.imdecode(image_array, cv2.IMREAD_COLOR)
         dimension = min(input_image.shape[0], input_image.shape[1])
         scale = max(math.floor(resize_dimension / dimension), 2)
-        output, _ = upsampler.enhance(input_image, outscale=scale)
+
+        os.chdir(str(Path.home() / ".cache"))
+        upsampler = RealESRGANer(
+            scale=4, model_path=model_path, model=rrbdnet, half="fp16", gpu_id="0"
+        )
+        face_enhancer = GFPGANer(
+            model_path=self.__model_path_gfpgan,
+            upscale=scale,
+            arch="clean",
+            channel_multiplier=2,
+            bg_upsampler=upsampler,
+        )
+
+        _, _, output = face_enhancer.enhance(
+            input_image, has_aligned=False, only_center_face=True, paste_back=True
+        )
+        os.chdir(get_root_dir())
         out_bytes = cv2.imencode(".png", output)[1].tobytes()
         return out_bytes

internals/util/commons.py

@@ -5,7 +5,7 @@ import random
 import re
 from io import BytesIO
 from pathlib import Path
-from typing import Optional, Union
+from typing import Any, Optional, Union
 
 import boto3
 import requests
@@ -191,7 +191,9 @@ def construct_default_s3_url(key):
     return "https://comic-assets.s3.ap-south-1.amazonaws.com/" + key
 
 
-def safe_index(array, index) -> Optional:
+def safe_index(array, index) -> Optional[Any]:
+    if not array:
+        return None
     if index < 0:
         return None
     if index >= len(array):

requirements.txt

@@ -10,6 +10,7 @@ rembg==2.0.30
 gfpgan==1.3.8
 rembg==2.0.30
 controlnet-aux==0.0.5
+gfpgan>=1.3.4
 realesrgan==0.3.0
 compel==1.0.4
 scikit-image>=0.19.3