Update app.py

app.py

@@ -1,248 +1,101 @@
 import io
 import time
 import logging
+import os
 from pathlib import Path
 from fastapi import FastAPI, File, UploadFile, HTTPException
-from fastapi.responses import StreamingResponse, JSONResponse
+from fastapi.responses import StreamingResponse
 from fastapi.middleware.cors import CORSMiddleware
 from PIL import Image
-import os
 
 # Setup logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-# Set environment variables BEFORE importing rembg
-os.environ["U2NET_HOME"] = "/tmp/.u2net"  # Tell rembg where to cache models
-os.environ["OMP_NUM_THREADS"] = "1"  # Optimize for single CPU
+# Set environment for Hugging Face
+os.environ["U2NET_HOME"] = "/tmp/.u2net"
 
 AI_AVAILABLE = False
+session = None
+remove = None
+
 try:
-    # Pre-download the model on startup
-    logger.info("🚀 Initializing Background Remover...")
-    
-    # First, create cache directory
-    cache_dir = "/tmp/.u2net"
-    os.makedirs(cache_dir, exist_ok=True)
-    
-    # Try to import with pre-load
+    logger.info("Loading AI model...")
     from rembg import remove, new_session
     
-    # Use smallest model for faster loading
-    logger.info("📥 Loading AI model (this might take a moment on first run)...")
-    session = new_session("u2netp")  # u2netp is smaller than u2net
+    # Create cache directory
+    os.makedirs("/tmp/.u2net", exist_ok=True)
+    
+    # Load model - use u2netp (smaller, faster)
+    session = new_session("u2netp")
     AI_AVAILABLE = True
-    logger.info("✅ Rembg AI Available and Model Loaded!")
+    logger.info("✅ AI model loaded successfully!")
     
 except Exception as e:
-    logger.error(f"❌ Failed to load AI: {e}")
+    logger.error(f"AI failed to load: {e}")
     AI_AVAILABLE = False
-    session = None
-    remove = None
-
-# Import CPU fallback
-import numpy as np
-import cv2
 
-app = FastAPI(title="Background Remover API")
+app = FastAPI(title="Background Remover")
 
-# Enable CORS
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
-    allow_credentials=True,
     allow_methods=["*"],
     allow_headers=["*"],
 )
 
-def cpu_fallback(input_image: Image.Image) -> Image.Image:
-    """A simple CPU-based fallback using OpenCV."""
-    try:
-        # Convert PIL to numpy array for OpenCV
-        img_np = np.array(input_image.convert('RGB'))
-        h, w = img_np.shape[:2]
-        
-        # If image is too large, resize for performance
-        if h > 1000 or w > 1000:
-            scale = 1000 / max(h, w)
-            new_h, new_w = int(h * scale), int(w * scale)
-            img_np = cv2.resize(img_np, (new_w, new_h))
-            h, w = new_h, new_w
-        
-        # Create a mask
-        mask = np.zeros((h, w), np.uint8)
-        bgd_model = np.zeros((1, 65), np.float64)
-        fgd_model = np.zeros((1, 65), np.float64)
-        
-        # Define rectangle (10% border)
-        border = 0.1
-        rect = (
-            int(w * border), 
-            int(h * border), 
-            int(w * (1 - border * 2)), 
-            int(h * (1 - border * 2))
-        )
-        
-        # Apply GrabCut
-        cv2.grabCut(img_np, mask, rect, bgd_model, fgd_model, 3, cv2.GC_INIT_WITH_RECT)
-        
-        # Create final mask
-        mask2 = np.where((mask == 2) | (mask == 0), 0, 1).astype('uint8')
-        
-        # Create output with white background
-        output = np.ones((h, w, 3), dtype=np.uint8) * 255
-        
-        # Apply mask to get foreground
-        foreground = img_np * mask2[:, :, np.newaxis]
-        
-        # Combine white background with foreground
-        output = output * (1 - mask2[:, :, np.newaxis]) + foreground
-        
-        return Image.fromarray(output.astype(np.uint8))
-        
-    except Exception as e:
-        logger.error(f"CPU fallback failed: {e}")
-        # Return original image with white background
-        return input_image.convert('RGB')
-
-@app.post("/api/process")
+@app.post("/process")
 async def process_image(file: UploadFile = File(...)):
-    """
-    Removes the background from an uploaded image.
-    """
-    start_time = time.time()
-    method_used = "unknown"
-    
+    """Process uploaded image and remove background"""
     try:
-        # Validate file
-        if not file.filename:
-            raise HTTPException(status_code=400, detail="No file provided")
-        
-        # Check file size (max 10MB)
-        file.file.seek(0, 2)  # Seek to end
-        file_size = file.file.tell()
-        file.file.seek(0)  # Reset to beginning
-        
-        if file_size > 10 * 1024 * 1024:  # 10MB
-            raise HTTPException(status_code=400, detail="File too large (max 10MB)")
-        
-        # Read and validate image
-        contents = await file.read()
-        if len(contents) == 0:
-            raise HTTPException(status_code=400, detail="Empty file")
-        
-        logger.info(f"📥 Processing: {file.filename} ({len(contents)} bytes)")
+        # Read image
+        image_data = await file.read()
+        if not image_data:
+            raise HTTPException(400, "No image data")
         
         # Open image
-        try:
-            image = Image.open(io.BytesIO(contents))
-        except Exception:
-            raise HTTPException(status_code=400, detail="Invalid image file")
+        image = Image.open(io.BytesIO(image_data))
         
-        # Process image
-        output_image = None
+        # Resize if too large (for performance)
+        max_size = 1024
+        if max(image.size) > max_size:
+            image.thumbnail((max_size, max_size), Image.LANCZOS)
         
+        # Process with AI if available
         if AI_AVAILABLE and session and remove:
-            try:
-                logger.info("🤖 Using AI method...")
-                # Resize if too large for AI
-                max_size = 1024
-                if max(image.size) > max_size:
-                    image.thumbnail((max_size, max_size), Image.Resampling.LANCZOS)
-                
-                # Remove background
-                output_image = remove(image, session=session)
-                method_used = "AI (u2netp)"
-                
-            except Exception as ai_error:
-                logger.warning(f"AI failed: {ai_error}, using CPU fallback")
-                output_image = cpu_fallback(image)
-                method_used = "CPU (GrabCut - AI failed)"
+            output = remove(image, session=session)
         else:
-            logger.info("⚙️ Using CPU method (AI not available)...")
-            output_image = cpu_fallback(image)
-            method_used = "CPU (GrabCut)"
+            raise HTTPException(503, "AI service not available")
         
         # Convert to bytes
         output_buffer = io.BytesIO()
-        output_image.save(output_buffer, format="PNG")
+        output.save(output_buffer, format="PNG")
         output_buffer.seek(0)
         
-        processing_time = time.time() - start_time
-        logger.info(f"✅ Success in {processing_time:.2f}s using {method_used}")
-        
-        # Return the image
         return StreamingResponse(
             output_buffer,
             media_type="image/png",
             headers={
-                "X-Processing-Method": method_used,
-                "X-Processing-Time": f"{processing_time:.2f}s",
-                "Content-Disposition": f'attachment; filename="no_bg_{Path(file.filename).stem}.png"'
+                "Content-Disposition": f"attachment; filename=no_bg_{Path(file.filename).stem}.png"
             }
         )
         
-    except HTTPException:
-        raise
     except Exception as e:
-        logger.error(f"❌ Processing error: {e}")
-        raise HTTPException(status_code=500, detail=f"Processing failed: {str(e)[:100]}")
-
-@app.post("/api/process_simple")
-async def process_image_simple(file: UploadFile = File(...)):
-    """
-    Simplified version without CPU fallback - for testing
-    """
-    try:
-        if AI_AVAILABLE and session and remove:
-            contents = await file.read()
-            image = Image.open(io.BytesIO(contents))
-            
-            # Resize for speed
-            image.thumbnail((512, 512))
-            
-            result = remove(image, session=session)
-            
-            output_buffer = io.BytesIO()
-            result.save(output_buffer, format="PNG")
-            output_buffer.seek(0)
-            
-            return StreamingResponse(output_buffer, media_type="image/png")
-        else:
-            return JSONResponse(
-                status_code=503,
-                content={"error": "AI service unavailable"}
-            )
-    except Exception as e:
-        logger.error(f"Simple process error: {e}")
-        raise HTTPException(status_code=500, detail=str(e))
-
-@app.get("/health")
-async def health_check():
-    """Health check endpoint"""
-    return {
-        "status": "healthy" if AI_AVAILABLE else "degraded",
-        "ai_available": AI_AVAILABLE,
-        "session_loaded": session is not None,
-        "timestamp": time.time()
-    }
+        logger.error(f"Error: {e}")
+        raise HTTPException(500, f"Processing failed: {str(e)}")
 
 @app.get("/")
-async def root():
-    """Simple homepage with form"""
+async def home():
+    """Simple homepage"""
     return {
         "message": "Background Remover API",
         "status": "running",
-        "ai": "available" if AI_AVAILABLE else "unavailable",
-        "endpoints": {
-            "POST /api/process": "Remove background from image",
-            "GET /health": "Health check"
-        },
-        "instructions": "Upload an image to /api/process to remove background"
+        "ai_available": AI_AVAILABLE,
+        "endpoint": "POST /process",
+        "example": "curl -X POST -F 'file=@your_image.jpg' https://your-space.hf.space/process -o result.png"
     }
 
-# Optional: Add a test endpoint
-@app.get("/test")
-async def test_endpoint():
-    """Test that the server is responding"""
-    return {"status": "ok", "message": "Server is running"}
+@app.get("/health")
+async def health():
+    """Health check"""
+    return {"status": "ok", "ai_loaded": AI_AVAILABLE}