update blocks

app.py

@@ -1,4 +1,4 @@
-# Part of the source code is in: fashn-ai/sapiens-body-part-segmentation
+# Part of the source code is from: fashn-ai/sapiens-body-part-segmentation
 import os
 
 import gradio as gr
@@ -9,23 +9,13 @@ 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
+from config import SAPIENS_LITE_MODELS_PATH
 
 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",
-}
-
+CHECKPOINTS_DIR = "checkpoints"
 
 def load_model(checkpoint_name: str):
     checkpoint_path = os.path.join(CHECKPOINTS_DIR, CHECKPOINTS[checkpoint_name])
@@ -35,8 +25,7 @@ def load_model(checkpoint_name: str):
     return model
 
 
-MODELS = {name: load_model(name) for name in CHECKPOINTS.keys()}
-
+#MODELS = {name: load_model(name) for name in CHECKPOINTS.keys()}
 
 @torch.inference_mode()
 def run_model(model, input_tensor, height, width):
@@ -53,8 +42,6 @@ transform_fn = transforms.Compose(
         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:
@@ -68,47 +55,47 @@ def segment(image: Image.Image, model_name: str) -> Image.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")
+def update_model_choices(task):
+    model_choices = list(SAPIENS_LITE_MODELS_PATH[task.lower()].keys())
+    return gr.Dropdown(choices=model_choices, value=model_choices[0] if model_choices else None)
+
+with gr.Blocks(theme=gr.themes.Monochrome(radius_size=sizes.radius_md)) as demo:
+    gr.Markdown("# Sapiens Arena 🤸🏽‍♂️ - WIP devmode- Not yet available")
+    with gr.Tabs():
+        with gr.TabItem('Image'):
+            with gr.Row():
+                with gr.Column():
+                    input_image = gr.Image(label="Input Image", type="pil", format="png")
+                    select_task = gr.Radio(
+                        ["Seg", "Pose", "Depth", "Normal"], 
+                        label="Task", 
+                        info="Choose the task to perfom",
+                        choices=list(SAPIENS_LITE_MODELS_PATH.keys())
+                    )
+                    model_name = gr.Dropdown(
+                        label="Model Version",
+                        choices=list(SAPIENS_LITE_MODELS_PATH["seg"].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")
+        
+        with gr.TabItem('Video'):
+            gr.Markdown("In construction")
+
+    select_task.change(fn=update_model_choices, inputs=select_task, outputs=model_name)
 
     run_button.click(
         fn=segment,

config.py

@@ -1,4 +1,4 @@
-SAPIENS_LITE_MODELS = {
+SAPIENS_LITE_MODELS_URL = {
     "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",
@@ -23,6 +23,31 @@ SAPIENS_LITE_MODELS = {
     }
 }
 
+SAPIENS_LITE_MODELS_PATH = {
+    "depth": {
+        "sapiens_0.3b": "checkpoints/sapiens_0.3b/sapiens_0.3b_render_people_epoch_100_torchscript.pt2",
+        "sapiens_0.6b": "checkpoints/sapiens_0.6b/sapiens_0.6b_render_people_epoch_70_torchscript.pt2",
+        "sapiens_1b": "checkpoints/sapiens_1b/sapiens_1b_render_people_epoch_88_torchscript.pt2",
+        "sapiens_2b": "checkpoints/sapiens_2b/sapiens_2b_render_people_epoch_25_torchscript.pt2"
+    },
+    "detector": {},
+    "normal": {
+        "sapiens_0.3b": "checkpoints/sapiens_0.3b/sapiens_0.3b_normal_render_people_epoch_66_torchscript.pt2",
+        "sapiens_0.6b": "checkpoints/sapiens_0.6b/sapiens_0.6b_normal_render_people_epoch_200_torchscript.pt2",
+        "sapiens_1b": "checkpoints/sapiens_1b/sapiens_1b_normal_render_people_epoch_115_torchscript.pt2",
+        "sapiens_2b": "checkpoints/sapiens_2b/sapiens_2b_normal_render_people_epoch_70_torchscript.pt2"
+    },
+    "pose": {
+        "sapiens_1b": "checkpoints/sapiens_1b/sapiens_1b_goliath_best_goliath_AP_640_torchscript.pt2"
+    },
+    "seg": {
+        "sapiens_0.3b": "checkpoints/sapiens_0.3b/sapiens_0.3b_goliath_best_goliath_mIoU_7673_epoch_194_torchscript.pt2",
+        "sapiens_0.6b": "checkpoints/sapiens_0.6b/sapiens_0.6b_goliath_best_goliath_mIoU_7777_epoch_178_torchscript.pt2",
+        "sapiens_1b": "checkpoints/sapiens_1b/sapiens_1b_goliath_best_goliath_mIoU_7994_epoch_151_torchscript.pt2",
+        "sapiens_2b": "checkpoints/sapiens_2b/sapiens_2b_goliath_best_goliath_mIoU_8179_epoch_181_torchscript.pt2"
+    }
+}
+
 LABELS_TO_IDS = {
     "Background": 0,
     "Apparel": 1,

sapiens

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