Update app.py

app.py

@@ -6,41 +6,59 @@ from PIL import Image
 import numpy as np
 import io
 import base64
+import sys
+
+# Configure error logging
+import logging
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
+logger = logging.getLogger("BackgroundRemover")
 
 # Initialize the segmentation model for background removal
-try:
-    # Using RMBG-1.4 which is specifically designed for background removal
-    # Added trust_remote_code=True to allow custom code execution
-    segmenter = pipeline(
-        "image-segmentation", 
-        model="briaai/RMBG-1.4",
-        device=0 if torch.cuda.is_available() else -1,
-        trust_remote_code=True  # Add this parameter to allow custom code execution
-    )
-    print("Successfully loaded RMBG-1.4 model")
-except Exception as e:
-    print(f"Error loading RMBG model: {e}")
-    # Fallback to a more standard segmentation model that doesn't require custom code
+segmenter = None
+
+def init_model():
+    global segmenter
     try:
+        # Using RMBG-1.4 which is specifically designed for background removal
+        logger.info("Loading RMBG-1.4 model...")
         segmenter = pipeline(
-            "image-segmentation",
-            model="facebook/detr-resnet-50-panoptic",
-            device=0 if torch.cuda.is_available() else -1
+            "image-segmentation", 
+            model="briaai/RMBG-1.4",
+            device=0 if torch.cuda.is_available() else -1,
+            trust_remote_code=True  # Allow custom code execution for the model
         )
-        print("Using fallback model: facebook/detr-resnet-50-panoptic")
-    except Exception as e2:
-        print(f"Error loading fallback model: {e2}")
-        segmenter = None
+        logger.info("Successfully loaded RMBG-1.4 model")
+    except Exception as e:
+        logger.error(f"Error loading RMBG model: {e}")
+        # Fallback to a more standard segmentation model that doesn't require custom code
+        try:
+            logger.info("Attempting to load fallback model...")
+            segmenter = pipeline(
+                "image-segmentation",
+                model="facebook/detr-resnet-50-panoptic",
+                device=0 if torch.cuda.is_available() else -1
+            )
+            logger.info("Using fallback model: facebook/detr-resnet-50-panoptic")
+        except Exception as e2:
+            logger.error(f"Error loading fallback model: {e2}")
+            segmenter = None
 
 def remove_background(input_image):
     """Remove background from an image using segmentation."""
-    if input_image is None:
-        return None
+    global segmenter
     
+    # Initialize model if not already done
     if segmenter is None:
-        print("No segmentation model available")
+        init_model()
+    
+    if segmenter is None:
+        logger.error("No segmentation model available")
         return input_image
     
+    if input_image is None:
+        logger.error("No input image provided")
+        return None
+    
     try:
         # Convert input image to numpy array if it's not already
         if isinstance(input_image, str):
@@ -51,32 +69,30 @@ def remove_background(input_image):
         
         # Check if image is valid
         if input_array.size == 0:
-            print("Empty input image")
+            logger.error("Empty input image")
             return input_image
             
-        print(f"Processing image of shape {input_array.shape}")
+        logger.info(f"Processing image of shape {input_array.shape}")
         
         # Run image segmentation
         result = segmenter(input_image)
-        print(f"Segmentation result type: {type(result)}")
+        logger.info(f"Segmentation result type: {type(result)}")
         
         # For the RMBG model, we directly get the mask
         if isinstance(result, dict) and 'mask' in result:
             # Direct mask from RMBG model
             mask_array = np.array(result['mask'])
             mask_array = mask_array / 255.0  # Normalize if needed
-            print("Using RMBG mask")
+            logger.info("Using RMBG mask")
         elif isinstance(result, list) and len(result) > 0:
             # Standard segmentation model output - try to create a foreground mask
-            # This will work differently depending on the model
-            # For DETR model, we need to identify person/object segments
             foreground_classes = ['person', 'animal', 'vehicle', 'object']
             
             # Initialize an empty mask
             if len(input_array.shape) == 3:
                 mask_array = np.zeros((input_array.shape[0], input_array.shape[1]), dtype=np.float32)
             else:
-                print("Invalid input image shape")
+                logger.error("Invalid input image shape")
                 return input_image
             
             # Combine all foreground segments
@@ -93,9 +109,9 @@ def remove_background(input_image):
                                 (mask_array.shape[1], mask_array.shape[0])))
                         # Add this segment to the foreground mask
                         mask_array = np.maximum(mask_array, segment_mask)
-            print("Created composite mask from segmentation model")
+            logger.info("Created composite mask from segmentation model")
         else:
-            print("Unexpected model output format")
+            logger.error("Unexpected model output format")
             return input_image
         
         # Create an RGBA image
@@ -111,19 +127,22 @@ def remove_background(input_image):
                 # For other models, we may need to invert the mask
                 rgba[:,:,3] = (mask_array * 255).astype(np.uint8)
             
-            print("Successfully created RGBA image")
+            logger.info("Successfully created RGBA image")
             return Image.fromarray(rgba)
         else:
-            print(f"Unexpected image format: shape {input_array.shape}")
+            logger.error(f"Unexpected image format: shape {input_array.shape}")
             return input_image
         
     except Exception as e:
-        print(f"Error in background removal: {e}")
+        logger.error(f"Error in background removal: {e}")
         # Return original image if processing failed
         return input_image
 
-# Create Gradio interface
-with gr.Blocks(css="footer {visibility: hidden}") as demo:
+# Initialize model on startup to avoid lazy loading during request
+init_model()
+
+# Create a simpler Gradio interface with minimal components to avoid internal errors
+with gr.Blocks(theme=gr.themes.Default(), css="footer {visibility: hidden}") as demo:
     gr.Markdown(
         """
         # Space BG Erase Studio
@@ -142,6 +161,7 @@ with gr.Blocks(css="footer {visibility: hidden}") as demo:
         with gr.Column():
             output_image = gr.Image(type="pil", label="Result (Transparent Background)")
     
+    # Simple click handler to avoid complex API handling
     submit_btn.click(
         fn=remove_background,
         inputs=input_image,