Create rembg_method.py

rembg_method.py

@@ -0,0 +1,40 @@
+from PIL import Image
+import io
+import logging
+from datetime import datetime
+
+logger = logging.getLogger(__name__)
+
+def remove_background(image_data):
+    start_time = datetime.now()
+    logger.info(f"Removing background - {start_time}")
+    
+    try:
+        # Try to import rembg
+        try:
+            from rembg import remove
+            logger.info("rembg module imported successfully")
+        except ImportError:
+            logger.error("rembg module not found!")
+            raise ImportError('rembg library not installed')
+            
+        input_image = Image.open(io.BytesIO(image_data))
+        logger.info(f"Image opened. Mode: {input_image.mode}, Size: {input_image.size}")
+        
+        if input_image.mode != 'RGB':
+            logger.info("Converting image to RGB...")
+            input_image = input_image.convert('RGB')
+            
+        logger.info("Removing background...")
+        output_image = remove(input_image)
+        
+        img_io = io.BytesIO()
+        output_image.save(img_io, 'PNG')
+        img_io.seek(0)
+        
+        duration = (datetime.now() - start_time).total_seconds()
+        logger.info(f"Successfully processed. Duration: {duration} seconds")
+        return img_io
+    except Exception as e:
+        logger.error(f"Background removal error: {str(e)}", exc_info=True)
+        raise