Update app.py

app.py

@@ -1,171 +1,241 @@
 
-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
+# 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
+# )
 
-torch.set_float32_matmul_precision(["high", "highest"][0])
 
-# load 2 models
+# # RMBG2 = AutoModelForImageSegmentation.from_pretrained(
+# #     "briaai/RMBG-2.0", trust_remote_code=True
+# # )
 
-birefnet = AutoModelForImageSegmentation.from_pretrained(
-    "ZhengPeng7/BiRefNet", trust_remote_code=True
-).to("cuda")
+# # Keep them in a dict to switch easily
+# models_dict = {
+#     "BiRefNet": birefnet
+#     # "RMBG-2.0": RMBG2,
+# }
 
+# # Transform
 
-# RMBG2 = AutoModelForImageSegmentation.from_pretrained(
-#     "briaai/RMBG-2.0", 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 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"
-)
+# @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"],
+#     choices=["BiRefNet"],# "RMBG-2.0"],
 #     value="BiRefNet",
 #     label="Select Model"
 # )
 # model_selector_3 = gr.Radio(
-#     choices=["BiRefNet", "RMBG-2.0"],
+#     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")
+# # 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 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 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")
+# # 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 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 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."
-)
+# # 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
+from transformers import AutoModelForImageSegmentation
+import torch
+from torchvision import transforms
+from PIL import Image
+import os
 
-# Combine all tabs
-demo = gr.TabbedInterface(
-    [tab1, tab2, tab3],
-    ["Image Upload", "URL Input", "File Output"],
-    title="Background Removal Tool"
+# precision tweak
+torch.set_float32_matmul_precision("high")
+
+# ─── Model & Transforms ─────────────────────────────────────────────────────────
+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]),
+])
+
+# ─── 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)
+    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)
+
 if __name__ == "__main__":
     demo.launch(show_error=True, share=True)
 
-