Update main.py

main.py

@@ -1,14 +1,4 @@
-###############################################################################
-# 1. Environment Setup
-###############################################################################
 import os
-
-os.environ["GRADIO_SERVER_NAME"] = "0.0.0.0"
-os.environ["GRADIO_SERVER_PORT"] = "7860"
-
-###############################################################################
-# 2. Imports
-###############################################################################
 import uuid
 import base64
 import numpy as np
@@ -26,29 +16,35 @@ import gradio as gr
 import uvicorn
 from datetime import datetime
 import huggingface_hub as hf
-from huggingface_hub import spaces
 
-# Import spaces directly without conditional
-try:
-    from spaces import GPU
-    print("✅ spaces.GPU imported successfully")
-    USE_SPACES_GPU = True
-except ImportError:
-    print("⚠️ spaces.GPU not available - using CPU fallback")
-    USE_SPACES_GPU = False
+# Set up environment variables
+os.environ["GRADIO_SERVER_NAME"] = "0.0.0.0"
+os.environ["GRADIO_SERVER_PORT"] = "7860"
 
-print("📦 main.py loaded. ZeroGPU debugging enabled.")
+# Define the FastAPI app
+app = FastAPI(
+    title="Background Removal API",
+    description="Background removal API with Gradio interface",
+    version="2.0.0",
+    docs_url="/api/docs",
+    redoc_url="/api/redoc",
+    openapi_url="/api/openapi.json",
+    contact={
+        "name": "Background Removal API",
+        "url": "https://github.com/yourusername/background-removal-api",
+    },
+    license_info={
+        "name": "MIT",
+        "url": "https://opensource.org/licenses/MIT",
+    },
+)
 
-###############################################################################
-# 3. Pydantic Models for API Documentation
-###############################################################################
+# Define Pydantic models for API documentation
 class BackgroundRemovalRequest(BaseModel):
-    """Request model for background removal"""
     image_format: Optional[str] = Field(default="PNG", description="Output image format (PNG, JPEG)")
     quality: Optional[int] = Field(default=95, ge=1, le=100, description="Output quality for JPEG (1-100)")
-    
+
 class BackgroundRemovalResponse(BaseModel):
-    """Response model for successful background removal"""
     status: str = Field(..., description="Status of the operation")
     image_code: str = Field(..., description="Base64 encoded image with data URI prefix")
     processing_time: float = Field(..., description="Processing time in seconds")
@@ -56,164 +52,24 @@ class BackgroundRemovalResponse(BaseModel):
     output_format: str = Field(..., description="Output image format")
 
 class ErrorResponse(BaseModel):
-    """Error response model"""
     status: str = Field(..., description="Status of the operation")
     message: str = Field(..., description="Error message")
     error_code: Optional[str] = Field(None, description="Specific error code")
 
 class HealthResponse(BaseModel):
-    """Health check response"""
     status: str = Field(..., description="Service status")
     timestamp: str = Field(..., description="Current timestamp")
     version: str = Field(..., description="API version")
     gpu_available: bool = Field(..., description="Whether GPU is available")
     model_loaded: bool = Field(..., description="Whether the model is loaded")
 
-###############################################################################
-# 4. App Setup with Enhanced Documentation
-###############################################################################
-API_KEY = os.getenv("API_KEY", "demo-key-change-in-production")
-
-app = FastAPI(
-    title="🧠 Background Removal API",
-    description="""
-    # Background Removal API with Gradio Interface
-    
-    This API provides advanced background removal capabilities using ONNX models with optional GPU acceleration.
-    
-    ## Features
-    - 🖼️ High-quality background removal
-    - ⚡ GPU acceleration (when available)
-    - 🎨 Multiple output formats (PNG, JPEG)
-    - 📱 Gradio web interface
-    - 🔒 API key authentication
-    - 📊 Real-time processing metrics
-    
-    ## Usage
-    1. **Web Interface**: Visit the root path `/` for the interactive Gradio interface
-    2. **REST API**: Use `/api/remove-background` endpoint for programmatic access
-    3. **Documentation**: Visit `/api/docs` for this interactive documentation
-    
-    ## Authentication
-    - API endpoints require an API key provided via form data or header
-    - Set `API_KEY` environment variable for production use
-    
-    ## Gradio Integration
-    According to [Gradio's documentation](https://www.gradio.app/guides/querying-gradio-apps-with-curl), 
-    the Gradio interface automatically exposes REST API endpoints that can be accessed via cURL:
-    
-    ```bash
-    # Make prediction
-    curl -X POST {your-url}/call/remove_background_gpu \\
-         -H "Content-Type: application/json" \\
-         -d '{"data": [{"path": "https://example.com/image.jpg"}]}'
-    
-    # Get result
-    curl -N {your-url}/call/remove_background_gpu/{event_id}
-    ```
-    """,
-    version="2.0.0",
-    docs_url="/api/docs",
-    redoc_url="/api/redoc",
-    openapi_url="/api/openapi.json",
-    contact={
-        "name": "Background Removal API",
-        "url": "https://github.com/yourusername/background-removal-api",
-    },
-    license_info={
-        "name": "MIT",
-        "url": "https://opensource.org/licenses/MIT",
-    },
-)
-
-# Security
-security = HTTPBearer()
-
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=["*"],
-    allow_credentials=True,
-    allow_methods=["*"],
-    allow_headers=["*"]
-)
-
-TMP_FOLDER = "tmp"
-os.makedirs(TMP_FOLDER, exist_ok=True)
-app.mount("/tmp", StaticFiles(directory=TMP_FOLDER), name="tmp")
-
-model_path = "BiRefNet-portrait-epoch_150.onnx"
-input_size = (1024, 1024)
-
-###############################################################################
-# 5. Authentication Helper
-###############################################################################
-async def verify_api_key(api_key: str = Form(...)):
-    """Verify API key from form data"""
-    if api_key != API_KEY:
-        raise HTTPException(
-            status_code=401, 
-            detail="Invalid API key",
-            headers={"WWW-Authenticate": "Bearer"},
-        )
-    return api_key
-
-async def verify_api_key_header(credentials: HTTPAuthorizationCredentials = Depends(security)):
-    """Verify API key from Authorization header"""
-    if credentials.credentials != API_KEY:
-        raise HTTPException(
-            status_code=401,
-            detail="Invalid API key",
-            headers={"WWW-Authenticate": "Bearer"},
-        )
-    return credentials.credentials
-
-###############################################################################
-# 6. Preprocess & Postprocess Utilities
-###############################################################################
-def preprocess_image(image):
-    """Handle both bytes and numpy arrays"""
-    if isinstance(image, bytes):
-        nparr = np.frombuffer(image, np.uint8)
-        img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
-    else:
-        img = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
-    
-    original_img = img.copy()
-    original_shape = img.shape[:2]
-    rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
-    resized = cv2.resize(rgb, input_size)
-    normalized = resized.astype(np.float32) / 255.0
-    normalized = (normalized - 0.5) / 0.5
-    transposed = np.transpose(normalized, (2, 0, 1))
-    input_tensor = np.expand_dims(transposed, axis=0).astype(np.float32)
-    return input_tensor, original_shape, original_img
-
-def apply_mask(original_img, mask_array, original_shape):
-    mask = np.squeeze(mask_array)
-    mask = cv2.resize(mask, (original_shape[1], original_shape[0]))
-    binary_mask = (mask > 0.5).astype(np.uint8)
-    masked_img = cv2.bitwise_and(original_img, original_img, mask=binary_mask)
-    alpha = (binary_mask * 255).astype(np.uint8)
-    bgra = cv2.cvtColor(masked_img, cv2.COLOR_BGR2BGRA)
-    bgra[:, :, 3] = alpha
-    return bgra
-
-###############################################################################
-# 7. Core Processing Function
-###############################################################################
-@hf.hub.git_repo
+# Define the core processing function
 def process_image_core(image_data, use_gpu=False, output_format="PNG", quality=95):
-    """Core image processing logic with enhanced options"""
-    import time
-    start_time = time.time()
-    
-    try:
-        providers = ["CUDAExecutionProvider", "CPUExecutionProvider"] if use_gpu else ["CPUExecutionProvider"]
-        session = ort.InferenceSession(model_path, providers=providers)
-    except ort.OrtSessionException as e:
-        print(f"⚠️ GPU session failed: {str(e)}")
-        session = ort.InferenceSession(model_path, providers=["CPUExecutionProvider"])
-    
+    # Load the model and preprocess the image
+    model_path = "BiRefNet-portrait-epoch_150.onnx"
+    input_size = (1024, 1024)
+    providers = ["CUDAExecutionProvider", "CPUExecutionProvider"] if use_gpu else ["CPUExecutionProvider"]
+    session = ort.InferenceSession(model_path, providers=providers)
     input_name = session.get_inputs()[0].name
     input_tensor, original_shape, original_img = preprocess_image(image_data)
     output = session.run(None, {input_name: input_tensor})
@@ -221,246 +77,46 @@ def process_image_core(image_data, use_gpu=False, output_format="PNG", quality=9
     result_img = apply_mask(original_img, mask, original_shape)
     result_pil = Image.fromarray(cv2.cvtColor(result_img, cv2.COLOR_BGRA2RGBA))
     
-    processing_time = time.time() - start_time
-    
-    return result_pil, processing_time, original_shape
+    # Save the result
+    result_filename = f"{uuid.uuid4()}.{output_format.lower()}"
+    output_path = f"tmp/{result_filename}"
+    if output_format.upper() == "PNG":
+        result_pil.save(output_path, "PNG")
+    else:
+        if result_pil.mode == 'RGBA':
+            rgb_img = Image.new('RGB', result_pil.size, (255, 255, 255))
+            rgb_img.paste(result_pil, mask=result_pil.split()[-1])
+            rgb_img.save(output_path, "JPEG", quality=quality)
+        else:
+            result_pil.save(output_path, "JPEG", quality=quality)
+
+    # Encode the result to base64
+    with open(output_path, "rb") as img_file:
+        base64_image = base64.b64encode(img_file.read()).decode("utf-8")
 
-###############################################################################
-# 8. GPU Gradio Function (Direct Decoration)
-###############################################################################
+    # Clean up temporary file
+    os.remove(output_path)
+
+    return base64_image
+
+# Define the Gradio interface
 def gradio_processor(image_np):
-    """Process image for Gradio interface"""
     # Convert numpy array to bytes
     success, img_encoded = cv2.imencode('.jpg', cv2.cvtColor(image_np, cv2.COLOR_RGB2BGR))
     if not success:
         raise ValueError("Failed to encode image")
     image_bytes = img_encoded.tobytes()
     
-    result, _, _ = process_image_core(image_bytes, use_gpu=USE_SPACES_GPU)
-    return result
-
-# Apply GPU decorator directly
-if USE_SPACES_GPU:
-    print("💻 Applying @GPU decorator directly")
-    gradio_processor = GPU(duration=240)(gradio_processor)
-    print("✅ @GPU decorator applied")
-else:
-    print("💻 Skipping GPU decorator")
-
-print("🧠 gradio_processor has GPU metadata:", getattr(gradio_processor, "_spaces_gpu", None))
-
-###############################################################################
-# 9. FastAPI Endpoints
-###############################################################################
-
-@app.get("/", response_class=HTMLResponse, include_in_schema=False)
-async def root():
-    """Root endpoint redirects to Gradio interface"""
-    return """
-    <!DOCTYPE html>
-    <html>
-    <head>
-        <title>Background Removal Service</title>
-        <meta http-equiv="refresh" content="0; url=/gradio">
-    </head>
-    <body>
-        <p>Redirecting to Gradio interface...</p>
-        <p><a href="/gradio">Click here if not redirected automatically</a></p>
-    </body>
-    </html>
-    """
-
-@app.get("/api/health", response_model=HealthResponse, tags=["Health"])
-async def health_check():
-    """
-    Health check endpoint to verify service status
-    
-    Returns:
-        HealthResponse: Current service status and information
-    """
-    try:
-        # Test model loading
-        session = ort.InferenceSession(model_path, providers=["CPUExecutionProvider"])
-        model_loaded = True
-    except:
-        model_loaded = False
-    
-    return HealthResponse(
-        status="healthy",
-        timestamp=datetime.now().isoformat(),
-        version="2.0.0",
-        gpu_available=USE_SPACES_GPU,
-        model_loaded=model_loaded
-    )
+    # Process the image
+    result_img = process_image_core(image_bytes, use_gpu=False, output_format="PNG")
+    return result_img
 
-@app.post(
-    "/api/remove-background",
-    response_model=BackgroundRemovalResponse,
-    responses={
-        401: {"model": ErrorResponse, "description": "Invalid API key"},
-        400: {"model": ErrorResponse, "description": "Invalid input"},
-        500: {"model": ErrorResponse, "description": "Processing error"}
-    },
-    tags=["Background Removal"]
-)
-async def remove_background_api(
-    image: UploadFile = File(..., description="Image file to process"),
-    api_key: str = Depends(verify_api_key),
-    output_format: str = Form(default="PNG", description="Output format: PNG or JPEG"),
-    quality: int = Form(default=95, ge=1, le=100, description="JPEG quality (1-100, ignored for PNG)")
-):
-    """
-    Remove background from uploaded image
-    
-    This endpoint processes an uploaded image and returns the result with background removed.
-    
-    **Parameters:**
-    - **image**: Image file to process (supports common formats: JPEG, PNG, WebP, etc.)
-    - **api_key**: Authentication key
-    - **output_format**: Output image format (PNG recommended for transparency)
-    - **quality**: JPEG compression quality (1-100, only used for JPEG output)
-    
-    **Returns:**
-    - Base64 encoded image with transparent background
-    - Processing time and metadata
-    
-    **Example Usage:**
-    ```bash
-    curl -X POST "http://localhost:7860/api/remove-background" \\
-         -F "api_key=demo-key-change-in-production" \\
-         -F "image=@your-image.jpg" \\
-         -F "output_format=PNG"
-    ```
-    """
-    try:
-        # Validate image file
-        if not image.content_type.startswith('image/'):
-            raise HTTPException(
-                status_code=400,
-                detail=f"Invalid file type: {image.content_type}. Please upload an image file."
-            )
-        
-        # Validate output format
-        if output_format.upper() not in ["PNG", "JPEG", "JPG"]:
-            raise HTTPException(
-                status_code=400,
-                detail="Invalid output format. Use 'PNG' or 'JPEG'."
-            )
-        
-        # Process image
-        image_data = await image.read()
-        result_img, processing_time, original_shape = process_image_core(
-            image_data, 
-            use_gpu=False,  # CPU for API endpoint
-            output_format=output_format.upper(),
-            quality=quality
-        )
-
-        # Save result
-        result_filename = f"{uuid.uuid4()}.{output_format.lower()}"
-        output_path = f"{TMP_FOLDER}/{result_filename}"
-        
-        if output_format.upper() == "PNG":
-            result_img.save(output_path, "PNG")
-        else:
-            # Convert RGBA to RGB for JPEG
-            if result_img.mode == 'RGBA':
-                rgb_img = Image.new('RGB', result_img.size, (255, 255, 255))
-                rgb_img.paste(result_img, mask=result_img.split()[-1])
-                rgb_img.save(output_path, "JPEG", quality=quality)
-            else:
-                result_img.save(output_path, "JPEG", quality=quality)
-
-        # Encode to base64
-        with open(output_path, "rb") as img_file:
-            base64_image = base64.b64encode(img_file.read()).decode("utf-8")
-
-        # Clean up temporary file
-        os.remove(output_path)
-
-        return BackgroundRemovalResponse(
-            status="success",
-            image_code=f"data:image/{output_format.lower()};base64,{base64_image}",
-            processing_time=round(processing_time, 3),
-            original_size=[original_shape[1], original_shape[0]],  # width, height
-            output_format=output_format.upper()
-        )
-
-    except HTTPException:
-        raise
-    except Exception as e:
-        raise HTTPException(
-            status_code=500,
-            detail=f"Processing error: {str(e)}"
-        )
-
-@app.post(
-    "/api/remove-background-file",
-    response_class=FileResponse,
-    tags=["Background Removal"]
-)
-async def remove_background_file(
-    image: UploadFile = File(..., description="Image file to process"),
-    api_key: str = Depends(verify_api_key),
-    output_format: str = Form(default="PNG", description="Output format: PNG or JPEG")
-):
-    """
-    Remove background and return image file directly
-    
-    Similar to `/remove-background` but returns the processed image file directly
-    instead of base64 encoded data.
-    """
-    try:
-        if not image.content_type.startswith('image/'):
-            raise HTTPException(status_code=400, detail="Invalid file type")
-        
-        image_data = await image.read()
-        result_img, _, _ = process_image_core(image_data, use_gpu=False)
-
-        result_filename = f"processed_{uuid.uuid4()}.{output_format.lower()}"
-        output_path = f"{TMP_FOLDER}/{result_filename}"
-        
-        if output_format.upper() == "PNG":
-            result_img.save(output_path, "PNG")
-        else:
-            if result_img.mode == 'RGBA':
-                rgb_img = Image.new('RGB', result_img.size, (255, 255, 255))
-                rgb_img.paste(result_img, mask=result_img.split()[-1])
-                rgb_img.save(output_path, "JPEG", quality=95)
-            else:
-                result_img.save(output_path, "JPEG", quality=95)
-
-        return FileResponse(
-            output_path, 
-            media_type=f"image/{output_format.lower()}",
-            filename=result_filename,
-            background=lambda: os.remove(output_path)  # Clean up after response
-        )
-
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
-
-###############################################################################
-# 10. Gradio Interface
-###############################################################################
 interface = gr.Interface(
     fn=gradio_processor,
     inputs=gr.Image(type="numpy", label="Upload Image"),
     outputs=gr.Image(type="pil", label="Processed Image"),
-    title="🧠 Background Removal Service",
-    description="""
-    ## Upload an image to remove its background
-    
-    - **GPU Acceleration**: Uses ZeroGPU when available
-    - **High Quality**: ONNX model for precise background removal
-    - **Fast Processing**: Optimized for real-time use
-    
-    ### API Access
-    This interface is also available via REST API:
-    - **Documentation**: [/api/docs](/api/docs)
-    - **Health Check**: [/api/health](/api/health)
-    - **Background Removal**: POST `/api/remove-background`
-    """,
+    title="Background Removal Service",
+    description="Upload an image to remove its background",
     examples=[
         # Add example images if you have them
     ],
@@ -468,21 +124,11 @@ interface = gr.Interface(
     allow_flagging="never"
 )
 
-interface.api_name = "remove_background_gpu"
-
-###############################################################################
-# 11. Mount Gradio App
-###############################################################################
+# Mount the Gradio app
 app = gr.mount_gradio_app(app, interface, path="/gradio", ssr_mode=False)
 
-###############################################################################
-# 12. Run Local Dev Server
-###############################################################################
+# Run the local dev server
 if __name__ == "__main__":
-    print("🚀 Starting Background Removal Service")
-    print(f"📖 API Documentation: http://localhost:7860/api/docs")
-    print(f"🎨 Gradio Interface: http://localhost:7860/gradio")
-    print(f"❤️ Health Check: http://localhost:7860/api/health")
     uvicorn.run(
         "main:app", 
         host="0.0.0.0",