wip

app.py

@@ -1,4 +1,4 @@
-# Part of the source code is from: fashn-ai/sapiens-body-part-segmentation
+# Part of the code is from: fashn-ai/sapiens-body-part-segmentation
 import os
 
 import gradio as gr
@@ -54,12 +54,11 @@ def segment(image: Image.Image, model_name: str) -> Image.Image:
     return blended_image
 
 
-# ----------------- GRADIO UI ----------------- #
 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:
+with gr.Blocks() as demo:
     gr.Markdown("# Sapiens Arena 🤸🏽‍♂️ - WIP devmode- Not yet available")
     with gr.Tabs():
         with gr.TabItem('Image'):

config.py

@@ -25,26 +25,26 @@ SAPIENS_LITE_MODELS_URL = {
 
 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"
+        "sapiens_0.3b": "checkpoints/depth/sapiens_0.3b_torchscript.pt2",
+        "sapiens_0.6b": "checkpoints/depth/sapiens_0.6b_torchscript.pt2",
+        "sapiens_1b": "checkpoints/depth/sapiens_1b_torchscript.pt2",
+        "sapiens_2b": "checkpoints/depth/sapiens_2b_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"
+        "sapiens_0.3b": "checkpoints/normal/sapiens_0.3b_torchscript.pt2",
+        "sapiens_0.6b": "checkpoints/normal/sapiens_0.6b_torchscript.pt2",
+        "sapiens_1b": "checkpoints/normal/sapiens_1b_torchscript.pt2",
+        "sapiens_2b": "checkpoints/normal/sapiens_2b_torchscript.pt2"
     },
     "pose": {
-        "sapiens_1b": "checkpoints/sapiens_1b/sapiens_1b_goliath_best_goliath_AP_640_torchscript.pt2"
+        "sapiens_1b": "checkpoints/pose/sapiens_1b_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"
+        "sapiens_0.3b": "checkpoints/seg/sapiens_0.3b_torchscript.pt2",
+        "sapiens_0.6b": "checkpoints/seg/sapiens_0.6b_torchscript.pt2",
+        "sapiens_1b": "checkpoints/seg/sapiens_1b_torchscript.pt2",
+        "sapiens_2b": "checkpoints/seg/sapiens_2b_torchscript.pt2"
     }
 }
 

utils/vis_utils.py

@@ -39,4 +39,44 @@ def visualize_mask_with_overlay(img: Image.Image, mask: Image.Image, labels_to_i
 
     blended = Image.fromarray(np.uint8(img_np * (1 - alpha) + overlay * alpha))
 
-    return blended
+    return blended
+
+def resize_image(pil_image, target_size):
+    """
+    Resize a PIL image while maintaining its aspect ratio.
+    
+    Args:
+    pil_image (PIL.Image): The input image.
+    target_size (tuple): The target size as (width, height).
+    
+    Returns:
+    PIL.Image: The resized image.
+    """
+    original_width, original_height = pil_image.size
+    target_width, target_height = target_size
+    
+    # Calculate aspect ratios
+    aspect_ratio = original_width / original_height
+    target_aspect = target_width / target_height
+    
+    if aspect_ratio > target_aspect:
+        # Image is wider than target, scale based on width
+        new_width = target_width
+        new_height = int(new_width / aspect_ratio)
+    else:
+        # Image is taller than target, scale based on height
+        new_height = target_height
+        new_width = int(new_height * aspect_ratio)
+    
+    # Resize the image
+    resized_image = pil_image.resize((new_width, new_height), Image.LANCZOS)
+    
+    # Create a new image with the target size and paste the resized image
+    new_image = Image.new('RGB', target_size, (0, 0, 0))
+    paste_x = (target_width - new_width) // 2
+    paste_y = (target_height - new_height) // 2
+    new_image.paste(resized_image, (paste_x, paste_y))
+    
+    return new_image
+
+