wip

.gitattributes

@@ -20,6 +20,7 @@
 *.pickle filter=lfs diff=lfs merge=lfs -text
 *.pkl filter=lfs diff=lfs merge=lfs -text
 *.pt filter=lfs diff=lfs merge=lfs -text
+*.pt2 filter=lfs diff=lfs merge=lfs -text
 *.pth filter=lfs diff=lfs merge=lfs -text
 *.rar filter=lfs diff=lfs merge=lfs -text
 *.safetensors filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1 @@
+__pycache__

app.py

@@ -1,7 +1,121 @@
+# Part of the source code is in: fashn-ai/sapiens-body-part-segmentation
+import os
+
 import gradio as gr
+import numpy as np
+import spaces
+import torch
+from gradio.themes.utils import sizes
+from PIL import Image
+from torchvision import transforms
+from utils.vis_utils import get_palette, visualize_mask_with_overlay
+
+if torch.cuda.is_available() and torch.cuda.get_device_properties(0).major >= 8:
+    torch.backends.cuda.matmul.allow_tf32 = True
+    torch.backends.cudnn.allow_tf32 = True
+
+ASSETS_DIR = os.path.join(os.path.dirname(__file__), "assets")
+
+
+CHECKPOINTS_DIR = os.path.join(ASSETS_DIR, "checkpoints")
+
+CHECKPOINTS = {
+    "0.3B": "sapiens_0.3b_goliath_best_goliath_mIoU_7673_epoch_194_torchscript.pt2",
+    "0.6B": "sapiens_0.6b_goliath_best_goliath_mIoU_7777_epoch_178_torchscript.pt2",
+    "1B": "sapiens_1b_goliath_best_goliath_mIoU_7994_epoch_151_torchscript.pt2",
+    "2B": "sapiens_2b_goliath_best_goliath_mIoU_8179_epoch_181_torchscript.pt2",
+}
+
+
+def load_model(checkpoint_name: str):
+    checkpoint_path = os.path.join(CHECKPOINTS_DIR, CHECKPOINTS[checkpoint_name])
+    model = torch.jit.load(checkpoint_path)
+    model.eval()
+    model.to("cuda")
+    return model
+
+
+MODELS = {name: load_model(name) for name in CHECKPOINTS.keys()}
+
+
+@torch.inference_mode()
+def run_model(model, input_tensor, height, width):
+    output = model(input_tensor)
+    output = torch.nn.functional.interpolate(output, size=(height, width), mode="bilinear", align_corners=False)
+    _, preds = torch.max(output, 1)
+    return preds
+
+
+transform_fn = transforms.Compose(
+    [
+        transforms.Resize((1024, 768)),
+        transforms.ToTensor(),
+        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+    ]
+)
+# ----------------- CORE FUNCTION ----------------- #
+
+
+@spaces.GPU
+def segment(image: Image.Image, model_name: str) -> Image.Image:
+    input_tensor = transform_fn(image).unsqueeze(0).to("cuda")
+    model = MODELS[model_name]
+    preds = run_model(model, input_tensor, height=image.height, width=image.width)
+    mask = preds.squeeze(0).cpu().numpy()
+    mask_image = Image.fromarray(mask.astype("uint8"))
+    blended_image = visualize_mask_with_overlay(image, mask_image, LABELS_TO_IDS, alpha=0.5)
+    return blended_image
+
+
+# ----------------- GRADIO UI ----------------- #
+
+
+with open("banner.html", "r") as file:
+    banner = file.read()
+with open("tips.html", "r") as file:
+    tips = file.read()
+
+CUSTOM_CSS = """
+.image-container  img {
+    max-width: 512px;
+    max-height: 512px;
+    margin: 0 auto;
+    border-radius: 0px;
+.gradio-container {background-color: #fafafa}
+"""
+
+with gr.Blocks(css=CUSTOM_CSS, theme=gr.themes.Monochrome(radius_size=sizes.radius_md)) as demo:
+    gr.HTML(banner)
+    gr.HTML(tips)
+    with gr.Row():
+        with gr.Column():
+            input_image = gr.Image(label="Input Image", type="pil", format="png")
+            model_name = gr.Dropdown(
+                label="Model Version",
+                choices=list(CHECKPOINTS.keys()),
+                value="0.3B",
+            )
+
+            example_model = gr.Examples(
+                inputs=input_image,
+                examples_per_page=10,
+                examples=[
+                    os.path.join(ASSETS_DIR, "examples", img)
+                    for img in os.listdir(os.path.join(ASSETS_DIR, "examples"))
+                ],
+            )
+        with gr.Column():
+            result_image = gr.Image(label="Segmentation Result", format="png")
+            run_button = gr.Button("Run")
+
+            gr.Image(os.path.join(ASSETS_DIR, "legend.png"), label="Legend", type="filepath")
+
+    run_button.click(
+        fn=segment,
+        inputs=[input_image, model_name],
+        outputs=[result_image],
+    )
 
-def greet(name):
-    return "Hello " + name + "!!"
 
-demo = gr.Interface(fn=greet, inputs="text", outputs="text")
-demo.launch()
+if __name__ == "__main__":
+    demo.launch(share=False)

checkpoints/depth/sapiens_0.3b_torchscript.pt2

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

checkpoints/depth/sapiens_0.6b_torchscript.pt2

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

checkpoints/depth/sapiens_1b_torchscript.pt2

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

checkpoints/depth/sapiens_2b_torchscript.pt2

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

checkpoints/normal/sapiens_0.3b_torchscript.pt2

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

checkpoints/normal/sapiens_0.6b_torchscript.pt2

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

checkpoints/normal/sapiens_1b_torchscript.pt2

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

checkpoints/normal/sapiens_2b_torchscript.pt2

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

checkpoints/pose/sapiens_1b_torchscript.pt2

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

checkpoints/seg/sapiens_0.3b_torchscript.pt2

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

checkpoints/seg/sapiens_0.6b_torchscript.pt2

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

checkpoints/seg/sapiens_1b_torchscript.pt2

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

checkpoints/seg/sapiens_2b_torchscript.pt2

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

config.py

@@ -0,0 +1,55 @@
+SAPIENS_LITE_MODELS = {
+    "depth": {
+        "sapiens_0.3b": "https://huggingface.co/facebook/sapiens/resolve/main/sapiens_lite_host/torchscript/depth/checkpoints/sapiens_0.3b/sapiens_0.3b_render_people_epoch_100_torchscript.pt2?download=true",
+        "sapiens_0.6b": "https://huggingface.co/facebook/sapiens/resolve/main/sapiens_lite_host/torchscript/depth/checkpoints/sapiens_0.6b/sapiens_0.6b_render_people_epoch_70_torchscript.pt2?download=true",
+        "sapiens_1b": "https://huggingface.co/facebook/sapiens/resolve/main/sapiens_lite_host/torchscript/depth/checkpoints/sapiens_1b/sapiens_1b_render_people_epoch_88_torchscript.pt2?download=true",
+        "sapiens_2b": "https://huggingface.co/facebook/sapiens/resolve/main/sapiens_lite_host/torchscript/depth/checkpoints/sapiens_2b/sapiens_2b_render_people_epoch_25_torchscript.pt2?download=true"
+    },
+    "detector": {},
+    "normal": {
+        "sapiens_0.3b": "https://huggingface.co/facebook/sapiens/resolve/main/sapiens_lite_host/torchscript/normal/checkpoints/sapiens_0.3b/sapiens_0.3b_normal_render_people_epoch_66_torchscript.pt2?download=true",
+        "sapiens_0.6b": "https://huggingface.co/facebook/sapiens/resolve/main/sapiens_lite_host/torchscript/normal/checkpoints/sapiens_0.6b/sapiens_0.6b_normal_render_people_epoch_200_torchscript.pt2?download=true",
+        "sapiens_1b": "https://huggingface.co/facebook/sapiens/resolve/main/sapiens_lite_host/torchscript/normal/checkpoints/sapiens_1b/sapiens_1b_normal_render_people_epoch_115_torchscript.pt2?download=true",
+        "sapiens_2b": "https://huggingface.co/facebook/sapiens/resolve/main/sapiens_lite_host/torchscript/normal/checkpoints/sapiens_2b/sapiens_2b_normal_render_people_epoch_70_torchscript.pt2?download=true"
+    },
+    "pose": {
+        "sapiens_1b": "https://huggingface.co/facebook/sapiens/resolve/main/sapiens_lite_host/torchscript/pose/checkpoints/sapiens_1b/sapiens_1b_goliath_best_goliath_AP_640_torchscript.pt2?download=true"
+    },
+    "seg": {
+        "sapiens_0.3b": "https://huggingface.co/facebook/sapiens/resolve/main/sapiens_lite_host/torchscript/seg/checkpoints/sapiens_0.3b/sapiens_0.3b_goliath_best_goliath_mIoU_7673_epoch_194_torchscript.pt2?download=true",
+        "sapiens_0.6b": "https://huggingface.co/facebook/sapiens/resolve/main/sapiens_lite_host/torchscript/seg/checkpoints/sapiens_0.6b/sapiens_0.6b_goliath_best_goliath_mIoU_7777_epoch_178_torchscript.pt2?download=true",
+        "sapiens_1b": "https://huggingface.co/facebook/sapiens/resolve/main/sapiens_lite_host/torchscript/seg/checkpoints/sapiens_1b/sapiens_1b_goliath_best_goliath_mIoU_7994_epoch_151_torchscript.pt2?download=true",
+        "sapiens_2b": "https://huggingface.co/facebook/sapiens/resolve/main/sapiens_lite_host/torchscript/seg/checkpoints/sapiens_2b/sapiens_2b_goliath_best_goliath_mIoU_8179_epoch_181_torchscript.pt2?download=true"
+    }
+}
+
+LABELS_TO_IDS = {
+    "Background": 0,
+    "Apparel": 1,
+    "Face Neck": 2,
+    "Hair": 3,
+    "Left Foot": 4,
+    "Left Hand": 5,
+    "Left Lower Arm": 6,
+    "Left Lower Leg": 7,
+    "Left Shoe": 8,
+    "Left Sock": 9,
+    "Left Upper Arm": 10,
+    "Left Upper Leg": 11,
+    "Lower Clothing": 12,
+    "Right Foot": 13,
+    "Right Hand": 14,
+    "Right Lower Arm": 15,
+    "Right Lower Leg": 16,
+    "Right Shoe": 17,
+    "Right Sock": 18,
+    "Right Upper Arm": 19,
+    "Right Upper Leg": 20,
+    "Torso": 21,
+    "Upper Clothing": 22,
+    "Lower Lip": 23,
+    "Upper Lip": 24,
+    "Lower Teeth": 25,
+    "Upper Teeth": 26,
+    "Tongue": 27,
+}

download_checkpoints.py

@@ -0,0 +1,42 @@
+import os
+import json
+import requests
+from tqdm import tqdm
+from config import SAPIENS_LITE_MODELS
+
+def download_file(url, filename):
+    response = requests.get(url, stream=True)
+    total_size = int(response.headers.get('content-length', 0))
+    
+    with open(filename, 'wb') as file, tqdm(
+        desc=filename,
+        total=total_size,
+        unit='iB',
+        unit_scale=True,
+        unit_divisor=1024,
+    ) as progress_bar:
+        for data in response.iter_content(chunk_size=1024):
+            size = file.write(data)
+            progress_bar.update(size)
+
+def main():
+    # Load the JSON file with model URLs
+    model_urls = SAPIENS_LITE_MODELS
+
+    for task, models in model_urls.items():
+        checkpoints_dir = os.path.join('checkpoints', task)
+        os.makedirs(checkpoints_dir, exist_ok=True)
+        
+        for model_name, url in models.items():
+            model_filename = f"{model_name}_torchscript.pt2"
+            model_path = os.path.join(checkpoints_dir, model_filename)
+            
+            if not os.path.exists(model_path):
+                print(f"Downloading {task} {model_name} model...")
+                download_file(url, model_path)
+                print(f"{task} {model_name} model downloaded successfully.")
+            else:
+                print(f"{task} {model_name} model already exists. Skipping download.")
+
+if __name__ == "__main__":
+    main()

requirements.txt

@@ -0,0 +1,7 @@
+gradio
+numpy
+torch
+torchvision
+matplotlib
+pillow
+spaces

sapiens

@@ -0,0 +1 @@
+Subproject commit 04bdc575d33ae93735f4c64887383e132951d8a4

utils/vis_utils.py

@@ -0,0 +1,42 @@
+# source: huggingface: fashn-ai/sapiens-body-part-segmentation
+import colorsys
+import matplotlib.colors as mcolors
+import numpy as np
+from PIL import Image
+
+def get_palette(num_cls):
+    palette = [0] * (256 * 3)
+    palette[0:3] = [0, 0, 0]
+
+    for j in range(1, num_cls):
+        hue = (j - 1) / (num_cls - 1)
+        saturation = 1.0
+        value = 1.0 if j % 2 == 0 else 0.5
+        rgb = colorsys.hsv_to_rgb(hue, saturation, value)
+        r, g, b = [int(x * 255) for x in rgb]
+        palette[j * 3 : j * 3 + 3] = [r, g, b]
+
+    return palette
+
+
+def create_colormap(palette):
+    colormap = np.array(palette).reshape(-1, 3) / 255.0
+    return mcolors.ListedColormap(colormap)
+
+
+def visualize_mask_with_overlay(img: Image.Image, mask: Image.Image, labels_to_ids: dict[str, int], alpha=0.5):
+    img_np = np.array(img.convert("RGB"))
+    mask_np = np.array(mask)
+
+    num_cls = len(labels_to_ids)
+    palette = get_palette(num_cls)
+    colormap = create_colormap(palette)
+
+    overlay = np.zeros((*mask_np.shape, 3), dtype=np.uint8)
+    for label, idx in labels_to_ids.items():
+        if idx != 0:
+            overlay[mask_np == idx] = np.array(colormap(idx)[:3]) * 255
+
+    blended = Image.fromarray(np.uint8(img_np * (1 - alpha) + overlay * alpha))
+
+    return blended