added remove background

Makefile

@@ -21,4 +21,9 @@ dev:
 
 hf:
 	chmod 777 hf.sh
-	./hf.sh
+	./hf.sh
+
+requirements:
+	uv pip compile --no-annotate pyproject.toml --no-deps --no-strip-extras --no-header \
+	| sed -E 's/([a-zA-Z0-9_-]+(\[[a-zA-Z0-9_,-]+\])?)[=><~!].*/\1/g' \
+	> requirements.txt

src/app.py

@@ -6,6 +6,7 @@ from tools import (
     change_color_objects_hsv,
     change_color_objects_lab,
     privacy_preserve_image,
+    remove_background,
 )
 
 gr.set_static_paths(paths=[Path.cwd().absolute() / "assets"])
@@ -30,7 +31,7 @@ hsv_df_input = gr.Dataframe(
 lab_df_input = gr.Dataframe(
     headers=["Object", "New A", "New B"],
     datatype=["str", "number", "number"],
-    col_count=(3,"fixed"),
+    col_count=(3, "fixed"),
     label="Target Objects and New Settings",
     type="array",
 )
@@ -119,13 +120,37 @@ privacy_preserve_tool = gr.Interface(
     ],
 )
 
+remove_background_tool = gr.Interface(
+    fn=remove_background,
+    inputs=[
+        gr.Image(label="Input Image", type="pil"),
+    ],
+    outputs=gr.Image(label="Output Image"),
+    title="Remove Image Background Tool",
+    description="Upload an image remove the background.",
+    examples=[
+        [
+            "https://raw.githubusercontent.com/mahan-ym/ImageAlfred/main/src/assets/examples/test_5.jpg",
+        ],
+        [
+            "https://raw.githubusercontent.com/mahan-ym/ImageAlfred/main/src/assets/examples/test_6.jpg",
+        ],
+    ],
+)
+
 demo = gr.TabbedInterface(
     [
         change_color_objects_hsv_tool,
         change_color_objects_lab_tool,
         privacy_preserve_tool,
+        remove_background_tool,
+    ],
+    [
+        "Change Color Objects HSV",
+        "Change Color Objects LAB",
+        "Privacy Preserving Tool",
+        "Remove Background Tool",
     ],
-    ["Change Color Objects HSV", "Change Color Objects LAB", "Privacy Preserving Tool"],
     title=title,
     theme=gr.themes.Default(
         primary_hue="blue",

src/modal_app.py

@@ -48,6 +48,10 @@ image = (
         "git+https://github.com/luca-medeiros/lang-segment-anything.git",
         gpu="A10G",
     )
+    .pip_install(
+        "git+https://github.com/PramaLLC/BEN2.git#egg=ben2",
+        gpu="A10G",
+    )
 )
 
 
@@ -79,11 +83,14 @@ def lang_sam_segment(
     if len(langsam_results[0]["labels"]) == 0:
         print("No masks found for the given prompt.")
         return None
-    
+
     print(f"found {len(langsam_results[0]['labels'])} masks for prompt: {prompt}")
     print("labels:", langsam_results[0]["labels"])
     print("scores:", langsam_results[0]["scores"])
-    print("masks scores:", langsam_results[0].get("mask_scores", "No mask scores available"))  # noqa: E501
+    print(
+        "masks scores:",
+        langsam_results[0].get("mask_scores", "No mask scores available"),
+    )  # noqa: E501
 
     return langsam_results
 
@@ -284,7 +291,7 @@ def preserve_privacy(
 
     for result in langsam_results:
         print(f"result: {result}")
-        
+
         for i, mask in enumerate(result["masks"]):
             if "mask_scores" in result:
                 if (
@@ -310,3 +317,27 @@ def preserve_privacy(
     output_image_pil = Image.fromarray(img_array)
 
     return output_image_pil
+
+
+@app.function(
+    gpu="A10G",
+    image=image,
+    volumes={volume_path: volume},
+    timeout=60 * 2,
+)
+def remove_background(image_pil: Image.Image) -> Image.Image:
+    from ben2 import BEN_Base
+    import torch
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"Using device: {device}")
+    print("type of image_pil:", type(image_pil))
+    model = BEN_Base.from_pretrained("PramaLLC/BEN2")
+    model.to(device).eval()
+
+    output_image = model.inference(
+        image_pil,
+        refine_foreground=True,
+    )
+    print(f"output type: {type(output_image)}")
+    return output_image

src/tools.py

@@ -9,6 +9,35 @@ from PIL import Image
 modal_app_name = "ImageAlfred"
 
 
+def remove_background(
+    input_img,
+) -> np.ndarray | Image.Image | str | Path | None:
+    """
+    Remove the background of the image.
+
+    Args:
+        input_img: Input image or can be URL string of the image or base64 string. Cannot be None.
+    Returns:
+        bytes: Binary image data of the modified image.
+    """  # noqa: E501
+    if not input_img:
+        raise gr.Error("Input image cannot be None or empty.")
+
+    func = modal.Function.from_name("ImageAlfred", "remove_background")
+    output_pil = func.remote(
+        image_pil=input_img,
+    )
+
+    if output_pil is None:
+        raise gr.Error("Received None from server.")
+    if not isinstance(output_pil, Image.Image):
+        raise gr.Error(
+            f"Expected Image.Image from server function, got {type(output_pil)}"
+        )
+
+    return output_pil
+
+
 def privacy_preserve_image(
     input_img,
     input_prompt,
@@ -99,7 +128,7 @@ def change_color_objects_hsv(
         )
     if not input_img:
         raise gr.Error("input img cannot be None or empty.")
-    
+
     print("before processing input:", user_input)
     valid_pattern = re.compile(r"^[a-zA-Z\s]+$")
     for item in user_input:
@@ -198,7 +227,7 @@ def change_color_objects_lab(
         raise gr.Error("input img cannot be None or empty.")
     valid_pattern = re.compile(r"^[a-zA-Z\s]+$")
     print("before processing input:", user_input)
-    
+
     for item in user_input:
         if len(item) != 3:
             raise gr.Error(