Update app.py

app.py

@@ -1,6 +1,5 @@
 import colorsys
 import os
-
 import gradio as gr
 import matplotlib.colors as mcolors
 import numpy as np
@@ -16,7 +15,6 @@ from torchvision import transforms
 ASSETS_DIR = os.path.join(os.path.dirname(__file__), "assets")
 os.makedirs(ASSETS_DIR, exist_ok=True)
 
-
 LABELS_TO_IDS = {
     "Background": 0,
     "Apparel": 1,
@@ -48,7 +46,6 @@ LABELS_TO_IDS = {
     "Tongue": 27,
 }
 
-
 def get_palette(num_cls):
     palette = [0] * (256 * 3)
     palette[0:3] = [0, 0, 0]
@@ -63,12 +60,10 @@ def get_palette(num_cls):
 
     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)
@@ -86,7 +81,6 @@ def visualize_mask_with_overlay(img: Image.Image, mask: Image.Image, labels_to_i
 
     return blended
 
-
 def create_legend_image(labels_to_ids: dict[str, int], filename="legend.png"):
     num_cls = len(labels_to_ids)
     palette = get_palette(num_cls)
@@ -128,30 +122,28 @@ def create_legend_image(labels_to_ids: dict[str, int], filename="legend.png"):
     plt.savefig(filename, dpi=300, bbox_inches="tight")
     plt.close()
 
-
-# create_legend_image(LABELS_TO_IDS, filename=os.path.join(ASSETS_DIR, "legend.png"))
-
-
 # ----------------- MODEL ----------------- #
 
-URL = "https://huggingface.co/facebook/sapiens/blob/main/sapiens_lite_host/torchscript/normal/checkpoints/sapiens_2b/sapiens_2b_normal_render_people_epoch_70_torchscript.pt2?download=true"
+URL = "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"
 CHECKPOINTS_DIR = os.path.join(ASSETS_DIR, "checkpoints")
 os.makedirs(CHECKPOINTS_DIR, exist_ok=True)
 
 model_path = os.path.join(CHECKPOINTS_DIR, "sapiens_2b_normal_render_people_epoch_70_torchscript.pt2")
 
-if not os.path.exists(model_path):
-    os.makedirs(CHECKPOINTS_DIR, exist_ok=True)
+if not os.path.exists(model_path) or os.path.getsize(model_path) == 0:
+    print("Downloading model...")
     import requests
 
     response = requests.get(URL)
-    with open(model_path, "wb") as file:
-        file.write(response.content)
+    if response.status_code == 200:
+        with open(model_path, "wb") as file:
+            file.write(response.content)
+    else:
+        raise Exception("Failed to download the model. Please check the URL.")
 
 model = torch.jit.load(model_path)
 model.eval()
 
-
 @torch.no_grad()
 def run_model(input_tensor, height, width):
     output = model(input_tensor)
@@ -159,7 +151,6 @@ def run_model(input_tensor, height, width):
     _, preds = torch.max(output, 1)
     return preds
 
-
 transform_fn = transforms.Compose(
     [
         transforms.Resize((1024, 768)),
@@ -167,8 +158,8 @@ transform_fn = transforms.Compose(
         transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
     ]
 )
-# ----------------- CORE FUNCTION ----------------- #
 
+# ----------------- CORE FUNCTION ----------------- #
 
 def segment(image: Image.Image) -> Image.Image:
     input_tensor = transform_fn(image).unsqueeze(0)
@@ -178,10 +169,8 @@ def segment(image: Image.Image) -> Image.Image:
     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:
@@ -202,27 +191,15 @@ with gr.Blocks(css=CUSTOM_CSS, theme=gr.themes.Monochrome(radius_size=sizes.radi
     with gr.Row():
         with gr.Column():
             input_image = gr.Image(label="Input Image", type="pil", format="png")
-
-            '''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],
         outputs=[result_image],
     )
 
-
 if __name__ == "__main__":
     demo.launch(share=False)