Upload folder using huggingface_hub

.gitattributes

@@ -33,7 +33,6 @@ 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
-example1.png filter=lfs diff=lfs merge=lfs -text
-example2.png filter=lfs diff=lfs merge=lfs -text
-example3.png filter=lfs diff=lfs merge=lfs -text
-examples.jpg filter=lfs diff=lfs merge=lfs -text
+example1.jpeg filter=lfs diff=lfs merge=lfs -text
+example2.jpeg filter=lfs diff=lfs merge=lfs -text
+example3.jpeg filter=lfs diff=lfs merge=lfs -text

app.py

@@ -1,4 +1,3 @@
-import spaces
 import numpy as np
 import torch
 import torch.nn.functional as F
@@ -6,27 +5,31 @@ import gradio as gr
 from ormbg import ORMBG
 from PIL import Image
 
+
 model_path = "ormbg.pth"
 
-# Load the model globally but don't send to device yet
 net = ORMBG()
-net.load_state_dict(torch.load(model_path, map_location="cpu"))
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+net.to(device)
+
+if torch.cuda.is_available():
+    net.load_state_dict(torch.load(model_path))
+    net = net.cuda()
+else:
+    net.load_state_dict(torch.load(model_path, map_location="cpu"))
 net.eval()
 
+
 def resize_image(image):
     image = image.convert("RGB")
     model_input_size = (1024, 1024)
     image = image.resize(model_input_size, Image.BILINEAR)
     return image
 
-@spaces.GPU
-@torch.inference_mode()
+
 def inference(image):
-    # Check for CUDA and set the device inside inference
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    net.to(device)
 
-    # Prepare input
+    # prepare input
     orig_image = Image.fromarray(image)
     w, h = orig_image.size
     image = resize_image(orig_image)
@@ -34,41 +37,50 @@ def inference(image):
     im_tensor = torch.tensor(im_np, dtype=torch.float32).permute(2, 0, 1)
     im_tensor = torch.unsqueeze(im_tensor, 0)
     im_tensor = torch.divide(im_tensor, 255.0)
-
     if torch.cuda.is_available():
-        im_tensor = im_tensor.to(device)
+        im_tensor = im_tensor.cuda()
 
-    # Inference
+    # inference
     result = net(im_tensor)
-    # Post process
+    # post process
     result = torch.squeeze(F.interpolate(result[0][0], size=(h, w), mode="bilinear"), 0)
     ma = torch.max(result)
     mi = torch.min(result)
     result = (result - mi) / (ma - mi)
-    # Image to PIL
+    # image to pil
     im_array = (result * 255).cpu().data.numpy().astype(np.uint8)
     pil_im = Image.fromarray(np.squeeze(im_array))
-    # Paste the mask on the original image
+    # paste the mask on the original image
     new_im = Image.new("RGBA", pil_im.size, (0, 0, 0, 0))
     new_im.paste(orig_image, mask=pil_im)
 
     return new_im
 
-# Gradio interface setup
+
+gr.Markdown("## Open Remove Background Model (ormbg)")
+gr.HTML(
+    """
+  <p style="margin-bottom: 10px; font-size: 94%">
+    This is a demo for Open Remove Background Model (ormbg) that using
+    <a href="https://huggingface.co/schirrmacher/ormbg" target="_blank">Open Remove Background Model (ormbg) model</a> as backbone.
+  </p>
+"""
+)
 title = "Open Remove Background Model (ormbg)"
 description = r"""
 This model is a <strong>fully open-source background remover</strong> optimized for images with humans.
+
 It is based on [Highly Accurate Dichotomous Image Segmentation research](https://github.com/xuebinqin/DIS).
 The model was trained with the synthetic [Human Segmentation Dataset](https://huggingface.co/datasets/schirrmacher/humans).
 
 This is the first iteration of the model, so there will be improvements!
-If you identify cases where the model fails, <a href='https://huggingface.co/schirrmacher/ormbg/discussions' target='_blank'>upload your examples</a>!
+If you identify cases were the model fails, <a href='https://huggingface.co/schirrmacher/ormbg/discussions' target='_blank'>upload your examples</a>!
 
 - <a href='https://huggingface.co/schirrmacher/ormbg' target='_blank'>Model card</a>: find inference code, training information, tutorials
 - <a href='https://huggingface.co/schirrmacher/ormbg' target='_blank'>Dataset</a>: see training images, segmentation data, backgrounds
 - <a href='https://huggingface.co/schirrmacher/ormbg\#research' target='_blank'>Research</a>: see current approach for improvements
-"""
 
+"""
 examples = ["./example1.png", "./example2.png", "./example3.png"]
 
 demo = gr.Interface(
@@ -77,7 +89,7 @@ demo = gr.Interface(
     outputs="image",
     examples=examples,
     title=title,
-    description=description
+    description=description,
 )
 
 if __name__ == "__main__":