Update app.py

app.py

@@ -1,173 +1,3 @@
-
-# import gradio as gr
-# from load_image import load_img
-# import spaces
-# from transformers import AutoModelForImageSegmentation
-# import torch
-# from torchvision import transforms
-# from PIL import Image
-# import os
-# import numpy as np
-
-# torch.set_float32_matmul_precision(["high", "highest"][0])
-
-# # load 2 models
-
-# birefnet = AutoModelForImageSegmentation.from_pretrained(
-#     "ZhengPeng7/BiRefNet", trust_remote_code=True
-# )
-
-
-# # RMBG2 = AutoModelForImageSegmentation.from_pretrained(
-# #     "briaai/RMBG-2.0", trust_remote_code=True
-# # )
-
-# # Keep them in a dict to switch easily
-# models_dict = {
-#     "BiRefNet": birefnet
-#     # "RMBG-2.0": RMBG2,
-# }
-
-# # Transform
-
-# transform_image = transforms.Compose(
-#     [
-#         transforms.Resize((1024, 1024)),
-#         transforms.ToTensor(),
-#         transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
-#     ]
-# )
-
-# @spaces.GPU
-# def process(image: Image.Image, model_choice: str):
-#     """
-#     Runs inference to remove the background (adds alpha) 
-#     with the chosen segmentation model.
-#     """
-#     # Select the model
-#     current_model = models_dict[model_choice]
-
-#     # Prepare image
-#     image_size = image.size
-#     input_images = transform_image(image).unsqueeze(0)
-
-#     # Inference
-#     with torch.no_grad():
-#         # Each model returns a list of preds in its forward, 
-#         # so we take the last element, apply sigmoid, and move to CPU
-#         preds = current_model(input_images)[-1].sigmoid().cpu()
-
-#     # Convert single-channel pred to a PIL mask
-#     pred = preds[0].squeeze()
-#     pred_pil = transforms.ToPILImage()(pred)
-
-#     # Resize the mask back to original image size
-#     mask = pred_pil.resize(image_size)
-
-#     # Add alpha channel to the original
-#     image.putalpha(mask)
-#     return image
-
-# def fn(source: str, model_choice: str):
-#     """
-#     Used by Tab 1 & Tab 2 to produce a processed image with alpha.
-#     - 'source' is either a file path (type="filepath") or 
-#       a URL string (textbox).
-#     - 'model_choice' is the user's selection from the radio.
-#     """
-#     # Load from local path or URL
-#     im = load_img(source, output_type="pil")
-#     im = im.convert("RGB")
-
-#     # Process
-#     processed_image = process(im, model_choice)
-#     return processed_image
-
-# def process_file(file_path: str, model_choice: str):
-#     """
-#     For Tab 3 (file output).
-#     - Accepts a local path, returns path to a new .png with alpha channel.
-#     - 'model_choice' is also passed in for selecting the model.
-#     """
-#     name_path = file_path.rsplit(".", 1)[0] + ".png"
-#     im = load_img(file_path, output_type="pil")
-#     im = im.convert("RGB")
-
-#     # Run the chosen model
-#     transparent = process(im, model_choice)
-#     transparent.save(name_path)
-#     return name_path
-
-
-# # GRadio UI
-
-# model_selector_1 = gr.Radio(
-#     choices=["BiRefNet"], # , "RMBG-2.0"
-#     value="BiRefNet",
-#     label="Select Model"
-# )
-# model_selector_2 = gr.Radio(
-#     choices=["BiRefNet"],# "RMBG-2.0"],
-#     value="BiRefNet",
-#     label="Select Model"
-# )
-# model_selector_3 = gr.Radio(
-#     choices=["BiRefNet"],# "RMBG-2.0"],
-#     value="BiRefNet",
-#     label="Select Model"
-# )
-
-# # Outputs for tabs 1 & 2: single processed image
-# processed_img_upload = gr.Image(label="Processed Image (Upload)", type="pil")
-# processed_img_url = gr.Image(label="Processed Image (URL)", type="pil")
-
-# # For uploading local files
-# image_upload = gr.Image(label="Upload an image", type="filepath")
-# image_file_upload = gr.Image(label="Upload an image", type="filepath")
-
-# # For Tab 2 (URL input)
-# url_input = gr.Textbox(label="Paste an image URL")
-
-# # For Tab 3 (file output)
-# output_file = gr.File(label="Output PNG File")
-
-# # Tab 1: local image -> processed image
-# tab1 = gr.Interface(
-#     fn=fn,
-#     inputs=[image_upload, model_selector_1],
-#     outputs=processed_img_upload,
-#     api_name="image",
-#     description="Upload an image and choose your background removal model."
-# )
-
-# # Tab 2: URL input -> processed image
-# tab2 = gr.Interface(
-#     fn=fn,
-#     inputs=[url_input, model_selector_2],
-#     outputs=processed_img_url,
-#     api_name="text",
-#     description="Paste an image URL and choose your background removal model."
-# )
-
-# # Tab 3: file output -> returns path to .png
-# tab3 = gr.Interface(
-#     fn=process_file,
-#     inputs=[image_file_upload, model_selector_3],
-#     outputs=output_file,
-#     api_name="png",
-#     description="Upload an image, choose a model, and get a transparent PNG."
-# )
-
-# # Combine all tabs
-# demo = gr.TabbedInterface(
-#     [tab1, tab2, tab3],
-#     ["Image Upload", "URL Input", "File Output"],
-#     title="Background Removal Tool"
-# )
-
-# if __name__ == "__main__":
-#     demo.launch(show_error=True, share=True)
-
 import gradio as gr
 from load_image import load_img
 import spaces
@@ -176,66 +6,140 @@ import torch
 from torchvision import transforms
 from PIL import Image
 import os
+import numpy as np
 
-# precision tweak
-torch.set_float32_matmul_precision("high")
+torch.set_float32_matmul_precision(["high", "highest"][0])
 
-# ─── Model & Transforms ─────────────────────────────────────────────────────────
+# load model
 birefnet = AutoModelForImageSegmentation.from_pretrained(
     "ZhengPeng7/BiRefNet", trust_remote_code=True
 )
 
-transform_image = transforms.Compose([
-    transforms.Resize((1024, 1024)),
-    transforms.ToTensor(),
-    transforms.Normalize([0.485, 0.456, 0.406],
-                         [0.229, 0.224, 0.225]),
-])
+# Keep model in a dict for easy switching
+models_dict = {
+    "BiRefNet": birefnet
+}
+
+# Transform
+transform_image = transforms.Compose(
+    [
+        transforms.Resize((1024, 1024)),
+        transforms.ToTensor(),
+        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
+    ]
+)
 
-# ─── Inference fn ────────────────────────────────────────────────────────────────
 @spaces.GPU
-def remove_bg(image: Image.Image):
-    orig_size = image.size
-    x = transform_image(image).unsqueeze(0)
-    out = birefnet(x)
-    logits = out.logits
-    mask = logits.sigmoid().cpu().squeeze(0)
-    mask_pil = transforms.ToPILImage()(mask).resize(orig_size)
-    image.putalpha(mask_pil)
+def process(image: Image.Image, model_choice: str):
+    """
+    Runs inference to remove the background (adds alpha) 
+    with the chosen segmentation model.
+    """
+    current_model = models_dict[model_choice]
+
+    # Prepare image
+    image_size = image.size
+    input_images = transform_image(image).unsqueeze(0)
+
+    # Inference
+    with torch.no_grad():
+        preds = current_model(input_images)[-1].sigmoid().cpu()
+
+    # Convert single-channel pred to a PIL mask
+    pred = preds[0].squeeze()
+    pred_pil = transforms.ToPILImage()(pred)
+
+    # Resize the mask back to original image size
+    mask = pred_pil.resize(image_size)
+
+    # Add alpha channel to the original
+    image.putalpha(mask)
     return image
 
-# ─── URL wrapper ─────────────────────────────────────────────────────────────────
-def remove_bg_url(url: str):
-    img = load_img(url, output_type="pil").convert("RGB")
-    return remove_bg(img)
-
-# ─── File‐based version ──────────────────────────────────────────────────────────
-def remove_bg_file(path: str):
-    img = load_img(path, output_type="pil").convert("RGB")
-    out = remove_bg(img)
-    out_path = os.path.splitext(path)[0] + ".png"
-    out.save(out_path)
-    return out_path
-
-# ─── Gradio UI ───────────────────────────────────────────────────────────────────
-with gr.Blocks() as demo:
-    gr.Markdown("## Background Removal (BiRefNet only)")
-
-    with gr.Tab("Upload"):
-        inp = gr.Image(type="filepath", label="Upload Image")
-        out = gr.Image(type="pil", label="Result")
-        inp.change(remove_bg, inp, out)
-
-    with gr.Tab("URL"):
-        url_in = gr.Textbox(label="Image URL")
-        out_url = gr.Image(type="pil", label="Result")
-        url_in.submit(remove_bg_url, url_in, out_url)
-
-    with gr.Tab("Save PNG"):
-        inp2 = gr.File(label="Image File")
-        out2 = gr.File(label="Transparent PNG")
-        inp2.change(remove_bg_file, inp2, out2)
+
+def fn(source: str, model_choice: str):
+    """
+    Used by Tab 1 & Tab 2 to produce a processed image with alpha.
+    - 'source' is either a file path (type="filepath") or 
+      a URL string (textbox).
+    - 'model_choice' is the user's selection from the radio.
+    """
+    im = load_img(source, output_type="pil")
+    im = im.convert("RGB")
+
+    # Process
+    processed_image = process(im, model_choice)
+    return processed_image
+
+
+def process_file(file_path: str, model_choice: str):
+    """
+    For Tab 3 (file output).
+    - Accepts a local path, returns path to a new .png with alpha channel.
+    - 'model_choice' is also passed in for selecting the model.
+    """
+    name_path = file_path.rsplit(".", 1)[0] + ".png"
+    im = load_img(file_path, output_type="pil")
+    im = im.convert("RGB")
+
+    transparent = process(im, model_choice)
+    transparent.save(name_path)
+    return name_path
+
+
+# Gradio UI\ nmodel_selector = gr.Radio(
+    choices=["BiRefNet"],
+    value="BiRefNet",
+    label="Select Model"
+)
+
+# Outputs for tabs 1 & 2: single processed image
+processed_img_upload = gr.Image(label="Processed Image (Upload)", type="pil")
+processed_img_url = gr.Image(label="Processed Image (URL)", type="pil")
+
+# For uploading local files
+image_upload = gr.Image(label="Upload an image", type="filepath")
+image_file_upload = gr.Image(label="Upload an image", type="filepath")
+
+# For Tab 2 (URL input)
+url_input = gr.Textbox(label="Paste an image URL")
+
+# For Tab 3 (file output)
+output_file = gr.File(label="Output PNG File")
+
+# Tab 1: local image -> processed image
+tab1 = gr.Interface(
+    fn=fn,
+    inputs=[image_upload, model_selector],
+    outputs=processed_img_upload,
+    api_name="image",
+    description="Upload an image and choose your background removal model."
+)
+
+# Tab 2: URL input -> processed image
+tab2 = gr.Interface(
+    fn=fn,
+    inputs=[url_input, model_selector],
+    outputs=processed_img_url,
+    api_name="text",
+    description="Paste an image URL and choose your background removal model."
+)
+
+# Tab 3: file output -> returns path to .png
+tab3 = gr.Interface(
+    fn=process_file,
+    inputs=[image_file_upload, model_selector],
+    outputs=output_file,
+    api_name="png",
+    description="Upload an image, choose a model, and get a transparent PNG."
+)
+
+# Combine all tabs
+demo = gr.TabbedInterface(
+    [tab1, tab2, tab3],
+    ["Image Upload", "URL Input", "File Output"],
+    title="Background Removal Tool"
+)
 
 if __name__ == "__main__":
     demo.launch(show_error=True, share=True)
-