cleanup

app.py

@@ -11,8 +11,6 @@ from base_utils import (
     parse_url,
 )
 
-# from background_removal import remove_bg
-
 pdf_to_img = gr.Interface(
     convert_pdf_to_image, gr.File(), gr.Gallery(), api_name="pdf_to_img"
 )
@@ -69,12 +67,6 @@ url_parser = gr.Interface(
     api_name="url_to_text",
 )
 
-# rmbg = gr.Interface(
-#     remove_bg,
-#     inputs=["image"],
-#     outputs=["image"],
-#     api_name="rmbg",
-# )
 
 demo = gr.TabbedInterface(
     [
@@ -97,7 +89,6 @@ demo = gr.TabbedInterface(
         "Extract PPTX Text",
         "Extract text from URL",
         "Extract Json",
-        # "Remove Background",
     ],
 )
 

background_removal.py

@@ -1,29 +0,0 @@
-import spaces
-from loadimg import load_img
-import torch
-from torchvision import transforms
-# Load BiRefNet with weights
-from transformers import AutoModelForImageSegmentation
-birefnet = AutoModelForImageSegmentation.from_pretrained('ZhengPeng7/BiRefNet', trust_remote_code=True)
-
-@spaces.GPU
-def remove_bg(imagepath):
-    # Data settings
-    image_size = (1024, 1024)
-    transform_image = transforms.Compose([
-        transforms.Resize(image_size),
-        transforms.ToTensor(),
-        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
-    ])
-
-    image = load_img(imagepath).convert("RGB")
-    input_images = transform_image(image).unsqueeze(0).to('cuda')
-
-    # Prediction
-    with torch.no_grad():
-        preds = birefnet(input_images)[-1].sigmoid().cpu()
-    pred = preds[0].squeeze()
-    pred_pil = transforms.ToPILImage()(pred)
-    mask = pred_pil.resize(image.size)
-    image.putalpha(mask)
-    return image