Upload 6 files

app.py

@@ -0,0 +1,185 @@
+import gradio as gr
+from PIL import ImageColor
+
+import onnxruntime
+import cv2
+import numpy as np
+
+# The common resume photo size is 35mmx45mm
+RESUME_PHOTO_W = 350
+RESUME_PHOTO_H = 450
+
+
+# modified from https://github.com/opencv/opencv_zoo/blob/main/models/face_detection_yunet/yunet.py
+class YuNet:
+    def __init__(
+        self,
+        modelPath,
+        inputSize=[320, 320],
+        confThreshold=0.6,
+        nmsThreshold=0.3,
+        topK=5000,
+        backendId=0,
+        targetId=0,
+    ):
+        self._modelPath = modelPath
+        self._inputSize = tuple(inputSize)  # [w, h]
+        self._confThreshold = confThreshold
+        self._nmsThreshold = nmsThreshold
+        self._topK = topK
+        self._backendId = backendId
+        self._targetId = targetId
+
+        self._model = cv2.FaceDetectorYN.create(
+            model=self._modelPath,
+            config="",
+            input_size=self._inputSize,
+            score_threshold=self._confThreshold,
+            nms_threshold=self._nmsThreshold,
+            top_k=self._topK,
+            backend_id=self._backendId,
+            target_id=self._targetId,
+        )
+
+    @property
+    def name(self):
+        return self.__class__.__name__
+
+    def setBackendAndTarget(self, backendId, targetId):
+        self._backendId = backendId
+        self._targetId = targetId
+        self._model = cv2.FaceDetectorYN.create(
+            model=self._modelPath,
+            config="",
+            input_size=self._inputSize,
+            score_threshold=self._confThreshold,
+            nms_threshold=self._nmsThreshold,
+            top_k=self._topK,
+            backend_id=self._backendId,
+            target_id=self._targetId,
+        )
+
+    def setInputSize(self, input_size):
+        self._model.setInputSize(tuple(input_size))
+
+    def infer(self, image):
+        # Forward
+        faces = self._model.detect(image)
+        return faces[1]
+
+
+class ONNXModel:
+    def __init__(self, model_path, input_w, input_h):
+        self.model = onnxruntime.InferenceSession(model_path)
+        self.input_w = input_w
+        self.input_h = input_h
+
+    def preprocess(self, rgb, mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)):
+        # convert the input data into the float32 input
+        img_data = (
+            np.array(cv2.resize(rgb, (self.input_w, self.input_h)))
+            .transpose(2, 0, 1)
+            .astype("float32")
+        )
+
+        # normalize
+        norm_img_data = np.zeros(img_data.shape).astype("float32")
+
+        for i in range(img_data.shape[0]):
+            norm_img_data[i, :, :] = img_data[i, :, :] / 255
+            norm_img_data[i, :, :] = (norm_img_data[i, :, :] - mean[i]) / std[i]
+
+        # add batch channel
+        norm_img_data = norm_img_data.reshape(1, 3, self.input_h, self.input_w).astype(
+            "float32"
+        )
+        return norm_img_data
+
+    def forward(self, image):
+        input_data = self.preprocess(image)
+        output_data = self.model.run(["argmax_0.tmp_0"], {"x": input_data})
+
+        return output_data
+
+
+def make_resume_photo(rgb, background_color):
+    h, w, _ = rgb.shape
+    bgr = cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR)
+
+    # Initialize models
+    face_detector = YuNet("models/face_detection_yunet_2023mar.onnx")
+    face_detector.setInputSize([w, h])
+    human_segmentor = ONNXModel(
+        "models/human_pp_humansegv2_lite_192x192_inference_model.onnx", 192, 192
+    )
+
+    # yunet uses opencv bgr image format
+    detections = face_detector.infer(bgr)
+
+    results = []
+    for idx, det in enumerate(detections):
+        # bounding box
+        pt1 = np.array((det[0], det[1]))
+        pt2 = np.array((det[0] + det[2], det[1] + det[3]))
+
+        # face landmarks
+        landmarks = det[4:14].reshape((5, 2))
+        right_eye = landmarks[0]
+        left_eye = landmarks[1]
+
+        angle = np.arctan2(right_eye[1] - left_eye[1], (right_eye[0] - left_eye[0]))
+        rmat = cv2.getRotationMatrix2D((0, 0), -angle, 1)
+
+        # apply rotation
+        rotated_bgr = cv2.warpAffine(bgr, rmat, (bgr.shape[1], bgr.shape[0]))
+        rotated_pt1 = rmat[:, :-1] @ pt1
+        rotated_pt2 = rmat[:, :-1] @ pt2
+
+        face_w, face_h = rotated_pt2 - rotated_pt1
+        up_length = int(face_h / 4)
+        down_length = int(face_h / 3)
+        crop_h = face_h + up_length + down_length
+        crop_w = int(crop_h * (RESUME_PHOTO_W / RESUME_PHOTO_H))
+
+        pt1 = np.array(
+            (rotated_pt1[0] - (crop_w - face_w) / 2, rotated_pt1[1] - up_length)
+        ).astype(np.int32)
+        pt2 = np.array((pt1[0] + crop_w, pt1[1] + crop_h)).astype(np.int32)
+
+        resume_photo = rotated_bgr[pt1[1] : pt2[1], pt1[0] : pt2[0], :]
+
+        rgb = cv2.cvtColor(resume_photo, cv2.COLOR_BGR2RGB)
+        mask = human_segmentor.forward(rgb)
+        mask = mask[0].transpose(1, 2, 0)
+        mask = cv2.resize(
+            mask.astype(np.uint8), (resume_photo.shape[1], resume_photo.shape[0])
+        )
+
+        resume_photo = cv2.cvtColor(resume_photo, cv2.COLOR_BGR2RGB)
+        resume_photo[mask == 0] = ImageColor.getcolor(background_color, "RGB")
+        resume_photo = cv2.resize(resume_photo, (RESUME_PHOTO_W, RESUME_PHOTO_H))
+        results.append(resume_photo)
+
+    return results
+
+
+title = "Resume Photo Maker"
+
+demo = gr.Interface(
+    fn=make_resume_photo,
+    inputs=[
+        gr.Image(type="numpy", label="input"),
+        gr.ColorPicker(label="background color"),
+    ],
+    outputs=gr.Gallery(label="output"),
+    examples=[
+        ["images/elon.jpg", "#FFFFFF"],
+        ["images/9_Press_Conference_Press_Conference_9_45.jpg", "#FFFFFF"],
+    ],
+    title=title,
+    allow_flagging="never",
+    article="<p style='text-align: center;'><a href='https://github.com/bot66/resume-photo-maker' target='_blank'>Github Repo</a></p>",
+)
+
+if __name__ == "__main__":
+    demo.launch()

models/face_detection_yunet_2023mar.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8f2383e4dd3cfbb4553ea8718107fc0423210dc964f9f4280604804ed2552fa4
+size 232589

models/human_pp_humansegv2_lite_192x192_inference_model.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:34edc335d7833f5a96bb2dadafb1d9da24bac072a26b447c18dd021ea8f29215
+size 12219997

requirements.txt

@@ -0,0 +1,73 @@
+aiofiles==23.2.1
+altair==5.1.2
+annotated-types==0.6.0
+anyio==3.7.1
+attrs==23.1.0
+certifi==2023.7.22
+charset-normalizer==3.3.2
+click==8.1.7
+colorama==0.4.6
+coloredlogs==15.0.1
+contourpy==1.2.0
+cycler==0.12.1
+exceptiongroup==1.1.3
+fastapi==0.104.1
+ffmpy==0.3.1
+filelock==3.13.1
+flatbuffers==23.5.26
+fonttools==4.44.0
+fsspec==2023.10.0
+gradio==4.3.0
+gradio_client==0.7.0
+h11==0.14.0
+httpcore==1.0.1
+httpx==0.25.1
+huggingface-hub==0.19.2
+humanfriendly==10.0
+idna==3.4
+importlib-resources==6.1.1
+Jinja2==3.1.2
+jsonschema==4.19.2
+jsonschema-specifications==2023.7.1
+kiwisolver==1.4.5
+markdown-it-py==3.0.0
+MarkupSafe==2.1.3
+matplotlib==3.8.1
+mdurl==0.1.2
+mpmath==1.3.0
+numpy==1.26.1
+onnxruntime==1.16.1
+opencv-python==4.8.1.78
+orjson==3.9.10
+packaging==23.2
+pandas==2.1.2
+Pillow==10.1.0
+protobuf==4.25.0
+pydantic==2.4.2
+pydantic_core==2.10.1
+pydub==0.25.1
+Pygments==2.16.1
+pyparsing==3.1.1
+python-dateutil==2.8.2
+python-multipart==0.0.6
+pytz==2023.3.post1
+PyYAML==6.0.1
+referencing==0.30.2
+requests==2.31.0
+rich==13.6.0
+rpds-py==0.12.0
+semantic-version==2.10.0
+shellingham==1.5.4
+six==1.16.0
+sniffio==1.3.0
+starlette==0.27.0
+sympy==1.12
+tomlkit==0.12.0
+toolz==0.12.0
+tqdm==4.66.1
+typer==0.9.0
+typing_extensions==4.8.0
+tzdata==2023.3
+urllib3==2.0.7
+uvicorn==0.24.0.post1
+websockets==11.0.3