Refactor app.py for modularity and error handling, and clean up requirements.txt

app.py

@@ -1,71 +1,87 @@
-import gradio as gr
-from loadimg import load_img
-import spaces
-from transformers import AutoModelForImageSegmentation
-import torch
-from torchvision import transforms
-
-torch.set_float32_matmul_precision(["high", "highest"][0])
-
-birefnet = AutoModelForImageSegmentation.from_pretrained(
-    "ZhengPeng7/BiRefNet", trust_remote_code=True
-)
-birefnet.to("cuda")
-
-transform_image = transforms.Compose(
-    [
-        transforms.Resize((1024, 1024)),
-        transforms.ToTensor(),
-        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
-    ]
-)
-
-def fn(image):
-    im = load_img(image, output_type="pil")
-    im = im.convert("RGB")
-    origin = im.copy()
-    processed_image = process(im)
-    return (processed_image, origin)
-
-@spaces.GPU
-def process(image):
-    image_size = image.size
-    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
-
-def process_file(f):
-    name_path = f.rsplit(".", 1)[0] + ".png"
-    im = load_img(f, output_type="pil")
-    im = im.convert("RGB")
-    transparent = process(im)
-    transparent.save(name_path)
-    return name_path
-
-slider1 = gr.ImageSlider(label="Processed Image", type="pil", format="png")
-slider2 = gr.ImageSlider(label="Processed Image from URL", type="pil", format="png")
-image_upload = gr.Image(label="Upload an image")
-image_file_upload = gr.Image(label="Upload an image", type="filepath")
-url_input = gr.Textbox(label="Paste an image URL")
-output_file = gr.File(label="Output PNG File")
-
-# Example images
-chameleon = load_img("butterfly.jpg", output_type="pil")
-url_example = "https://hips.hearstapps.com/hmg-prod/images/gettyimages-1229892983-square.jpg"
-
-tab1 = gr.Interface(fn, inputs=image_upload, outputs=slider1, examples=[chameleon], api_name="image")
-tab2 = gr.Interface(fn, inputs=url_input, outputs=slider2, examples=[url_example], api_name="text")
-tab3 = gr.Interface(process_file, inputs=image_file_upload, outputs=output_file, examples=["butterfly.jpg"], api_name="png")
-
-demo = gr.TabbedInterface(
-    [tab1, tab2, tab3], ["Image Upload", "URL Input", "File Output"], title="Background Removal Tool"
-)
-
-if __name__ == "__main__":
-    demo.launch(show_error=True)
+
+import gradio as gr
+import spaces
+import torch
+from image_loader import load_image_from_url, load_image_from_file
+from image_processor import process_image
+import logging
+
+# Configure logging
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
+
+torch.set_float32_matmul_precision(["high", "highest"][0])
+
+try:
+    birefnet = AutoModelForImageSegmentation.from_pretrained(
+        "ZhengPeng7/BiRefNet", trust_remote_code=True
+    )
+    birefnet.to("cuda")
+    logging.info("BiRefNet model loaded successfully.")
+except Exception as e:
+    logging.error(f"Error loading BiRefNet model: {e}")
+    raise Exception(f"Error loading BiRefNet model: {e}")
+
+def fn(image_input):
+    try:
+        if isinstance(image_input, str):  # URL input
+            img = load_image_from_url(image_input)
+        else:  # File upload
+            img = load_image_from_file(image_input)
+
+        img = img.convert("RGB")
+        origin = img.copy()
+        processed_image = process(img)
+        return (processed_image, origin)
+    except Exception as e:
+        logging.error(f"Error in fn function: {e}")
+        return None, None  # Return None or a placeholder image
+
+@spaces.GPU
+def process(image):
+    try:
+        processed_image = process_image(image, birefnet)
+        return processed_image
+    except Exception as e:
+        logging.error(f"Error in process function: {e}")
+        raise gr.Error(f"Error processing image: {e}")
+
+
+def process_file(file_path):
+    try:
+        name_path = file_path.rsplit(".", 1)[0] + ".png"
+        img = load_image_from_file(file_path)
+        img = img.convert("RGB")
+        transparent = process(img)
+        transparent.save(name_path)
+        logging.info(f"Processed image saved to: {name_path}")
+        return name_path
+    except Exception as e:
+        logging.error(f"Error in process_file function: {e}")
+        raise gr.Error(f"Error processing file: {e}")
+
+slider1 = gr.ImageSlider(label="Processed Image", type="pil", format="png")
+slider2 = gr.ImageSlider(label="Processed Image from URL", type="pil", format="png")
+image_upload = gr.Image(label="Upload an image")
+image_file_upload = gr.Image(label="Upload an image", type="filepath")
+url_input = gr.Textbox(label="Paste an image URL")
+output_file = gr.File(label="Output PNG File")
+
+# Example images
+try:
+    chameleon = load_image_from_file("butterfly.jpg")
+except Exception as e:
+    logging.error(f"Error loading example image: {e}")
+    chameleon = None  # Or a placeholder image
+
+url_example = "https://hips.hearstapps.com/hmg-prod/images/gettyimages-1229892983-square.jpg"
+
+tab1 = gr.Interface(fn, inputs=image_upload, outputs=slider1, examples=[chameleon], api_name="image")
+tab2 = gr.Interface(fn, inputs=url_input, outputs=slider2, examples=[url_example], api_name="text")
+tab3 = gr.Interface(process_file, inputs=image_file_upload, outputs=output_file, examples=["butterfly.jpg"], api_name="png")
+
+demo = gr.TabbedInterface(
+    [tab1, tab2, tab3], ["Image Upload", "URL Input", "File Output"], title="Background Removal Tool"
+)
+
+if __name__ == "__main__":
+    demo.launch(show_error=True)

image_loader.py

@@ -0,0 +1,74 @@
+
+from PIL import Image
+import requests
+from io import BytesIO
+import logging
+
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
+
+def load_image_from_url(url):
+    """Loads an image from a URL.
+
+    Args:
+        url (str): The URL of the image.
+
+    Returns:
+        PIL.Image.Image: The loaded image.
+
+    Raises:
+        Exception: If the image cannot be loaded from the URL.
+    """
+    try:
+        response = requests.get(url, stream=True)
+        response.raise_for_status()  # Raise HTTPError for bad responses (4xx or 5xx)
+        image = Image.open(BytesIO(response.content))
+        logging.info(f"Image loaded successfully from URL: {url}")
+        return image
+    except requests.exceptions.RequestException as e:
+        logging.error(f"Error loading image from URL: {url} - {e}")
+        raise Exception(f"Error loading image from URL: {url} - {e}")
+    except Exception as e:
+        logging.error(f"Error opening image from URL: {url} - {e}")
+        raise Exception(f"Error opening image from URL: {url} - {e}")
+
+
+def load_image_from_file(file_path):
+    """Loads an image from a file.
+
+    Args:
+        file_path (str): The path to the image file.
+
+    Returns:
+        PIL.Image.Image: The loaded image.
+
+    Raises:
+        Exception: If the image cannot be loaded from the file.
+    """
+    try:
+        image = Image.open(file_path)
+        logging.info(f"Image loaded successfully from file: {file_path}")
+        return image
+    except FileNotFoundError:
+        logging.error(f"File not found: {file_path}")
+        raise Exception(f"File not found: {file_path}")
+    except Exception as e:
+        logging.error(f"Error loading image from file: {file_path} - {e}")
+        raise Exception(f"Error loading image from file: {file_path} - {e}")
+
+if __name__ == '__main__':
+    # Example Usage
+    try:
+        image_url = "https://hips.hearstapps.com/hmg-prod/images/gettyimages-1229892983-square.jpg"
+        image_from_url = load_image_from_url(image_url)
+        print("Image loaded from URL successfully!")
+        # image_from_url.show() # Display the image (optional)
+    except Exception as e:
+        print(e)
+
+    try:
+        image_path = "butterfly.jpg"
+        image_from_file = load_image_from_file(image_path)
+        print("Image loaded from file successfully!")
+        # image_from_file.show() # Display the image (optional)
+    except Exception as e:
+        print(e)

image_processor.py

@@ -0,0 +1,43 @@
+
+import torch
+from torchvision import transforms
+from PIL import Image
+import logging
+
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
+
+transform_image = transforms.Compose(
+    [
+        transforms.Resize((1024, 1024)),
+        transforms.ToTensor(),
+        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
+    ]
+)
+
+def process_image(image, birefnet, device="cuda"):
+    """Processes the input image to remove the background.
+
+    Args:
+        image (PIL.Image.Image): The image to process.
+        birefnet (torch.nn.Module): The BiRefNet model.
+        device (str): The device to run the model on (default: "cuda").
+
+    Returns:
+        PIL.Image.Image: The processed image with background removed.
+    """
+    try:
+        image_size = image.size
+        input_images = transform_image(image).unsqueeze(0).to(device)
+
+        # 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)
+        logging.info("Image processed successfully.")
+        return image
+    except Exception as e:
+        logging.error(f"Error processing image: {e}")
+        raise Exception(f"Error processing image: {e}")

requirements.txt

@@ -1,17 +1,18 @@
-torch
-accelerate
-opencv-python
-spaces
-pillow
-numpy
-timm
-kornia
-prettytable
-typing
-scikit-image
-huggingface_hub
-transformers>=4.39.1
-gradio
-gradio_imageslider
-loadimg>=0.1.1
-einops
+
+accelerate==0.27.2
+einops==0.7.0
+gradio==4.16.0
+gradio_imageslider==0.2.0
+huggingface_hub==0.20.3
+kornia==0.7.1
+loadimg==0.1.1
+numpy==1.26.4
+opencv-python==4.9.0.54
+pillow==10.2.0
+prettytable==4.0.0
+scikit-image==0.23.0
+spaces==0.35.0
+timm==0.9.12
+torch==2.2.0
+transformers==4.39.1
+typing==3.7.4.3