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

Refactored app.py to improve modularity by creating separate modules for image loading (image_loader.py) and processing (image_processor.py). Implemented error handling for image loading from URLs and files, displaying informative messages to the user. Added logging to track image loading, processing, and saving events. Cleaned up requirements.txt by sorting packages, pinning versions, and removing duplicates.

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