Upload 6 files

.gitattributes

@@ -33,3 +33,16 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+
+### Python ###
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/

app.py

@@ -0,0 +1,108 @@
+from rtmlib import YOLOX, RTMPose, draw_bbox, draw_skeleton
+import functools
+from typing import Callable
+from pathlib import Path
+import gradio as gr
+import numpy as np
+import PIL.Image
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from PIL import Image
+import torchvision.transforms as T
+
+
+TITLE = 'Face Parsing'
+
+def get_palette(num_cls):
+    """ Returns the color map for visualizing the segmentation mask.
+    Args:
+        num_cls: Number of classes
+    Returns:
+        The color map
+    """
+
+    n = num_cls
+    palette = [0] * (n * 3)
+    for j in range(0, n):
+        lab = j
+        palette[j * 3 + 0] = 0
+        palette[j * 3 + 1] = 0
+        palette[j * 3 + 2] = 0
+        i = 0
+        while lab:
+            palette[j * 3 + 0] |= (((lab >> 0) & 1) << (7 - i))
+            palette[j * 3 + 1] |= (((lab >> 1) & 1) << (7 - i))
+            palette[j * 3 + 2] |= (((lab >> 2) & 1) << (7 - i))
+            i += 1
+            lab >>= 3
+    return palette
+
+@torch.inference_mode()
+def predict(image: PIL.Image.Image, model, transform: Callable,
+            device: torch.device,palette) -> np.ndarray:
+    img_show = np.array(image.copy())
+    bboxes = model[1](np. array(image))
+    img_show = draw_bbox(img_show, bboxes)
+    keypoints,scores = model[2](np. array(image),bboxes=bboxes)
+    img_show = draw_skeleton(img_show,keypoints,scores)
+
+    data = transform(image)
+    data = data.unsqueeze(0).to(device)
+    out = model[0](data)
+    out =  F.interpolate(out, [image.size[1],image.size[0]], mode="bilinear")
+    output = out[0].permute(1,2,0)
+    parsing = torch.argmax(output,dim=2).cpu().numpy()
+
+    output_im = Image.fromarray(np.asarray(parsing, dtype=np.uint8))
+    image = Image.fromarray(np.asarray(img_show, dtype=np.uint8))
+    output_im.putpalette(palette)
+    output_im = output_im.convert('RGB')
+    # output_im.save('output.png')
+
+    res = Image.blend(image.convert('RGB'), output_im, 0.5)
+    return output_im, res
+
+
+def load_parsing_model():
+    model = torch.jit.load(Path("models/faceparsing_512_512.pt"))
+    model.eval()
+    return model
+
+
+def main():
+    device = torch.device('cpu')
+    model_ls =[]
+    model = load_parsing_model()
+
+    transform = T.Compose([
+        T.Resize((512, 512), interpolation=PIL.Image.NEAREST),
+        T.ToTensor(),
+        T.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
+    ])
+    palette = get_palette(20)
+    det_model = YOLOX('models/det.onnx',model_input_size=(640,640),backend='onnxruntime', device='cpu')
+    pose_model = RTMPose('models/pose.onnx', model_input_size=(192, 256),to_openpose=False, backend='onnxruntime', device='cpu')
+
+    model_ls.append(model)
+    model_ls.append(det_model)
+    model_ls.append(pose_model)
+
+    func = functools.partial(predict,
+                             model=model_ls,
+                             transform=transform,
+                             device=device,palette=palette)
+
+
+    gr.Interface(
+        fn=func,
+        inputs=gr.Image(label='Input', type='pil'),
+        outputs=[
+            gr.Image(label='Predicted Labels', type='pil'),
+            gr.Image(label='Masked', type='pil'),
+        ],
+        title=TITLE,
+    ).queue().launch(show_api=False)
+
+if __name__ == "__main__":
+    main()

models/det.onnx

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

models/faceparsing_512_512.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8267ce549f91c0e87f575e80d339b6cebf743c05fb23064cc82e2ac180bfc5e8
+size 257745243

models/pose.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7bca282009dd5e5b8a94bb27e8015f610630643659b862818803febf9107e2e5
+size 368041127