Update main.py

main.py

@@ -15,12 +15,16 @@ import numpy as np
 import cv2
 import onnxruntime as ort
 from PIL import Image
-from fastapi import FastAPI, File, UploadFile, Form, HTTPException
-from fastapi.responses import JSONResponse
+from fastapi import FastAPI, File, UploadFile, Form, HTTPException, Query, Depends
+from fastapi.responses import JSONResponse, FileResponse, HTMLResponse
 from fastapi.staticfiles import StaticFiles
 from fastapi.middleware.cors import CORSMiddleware
+from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
+from pydantic import BaseModel, Field
+from typing import Optional, List, Union
 import gradio as gr
 import uvicorn
+from datetime import datetime
 
 # Import spaces directly without conditional
 try:
@@ -34,14 +38,99 @@ except ImportError:
 print("📦 main.py loaded. ZeroGPU debugging enabled.")
 
 ###############################################################################
-# 3. App Setup
+# 3. Pydantic Models for API Documentation
 ###############################################################################
-API_KEY = os.getenv("API_KEY")
+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")
+    original_size: List[int] = Field(..., description="Original image dimensions [width, height]")
+    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 = FastAPI(docs_url="/api/docs", openapi_url="/api/openapi.json")
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
+    allow_credentials=True,
     allow_methods=["*"],
     allow_headers=["*"]
 )
@@ -54,7 +143,30 @@ model_path = "BiRefNet-portrait-epoch_150.onnx"
 input_size = (1024, 1024)
 
 ###############################################################################
-# 4. Preprocess & Postprocess Utilities
+# 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"""
@@ -85,10 +197,13 @@ def apply_mask(original_img, mask_array, original_shape):
     return bgra
 
 ###############################################################################
-# 5. Core Processing Function
+# 7. Core Processing Function
 ###############################################################################
-def process_image_core(image_data, use_gpu=False):
-    """Core image processing logic"""
+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)
@@ -101,10 +216,14 @@ def process_image_core(image_data, use_gpu=False):
     output = session.run(None, {input_name: input_tensor})
     mask = output[0]
     result_img = apply_mask(original_img, mask, original_shape)
-    return Image.fromarray(cv2.cvtColor(result_img, cv2.COLOR_BGRA2RGBA))
+    result_pil = Image.fromarray(cv2.cvtColor(result_img, cv2.COLOR_BGRA2RGBA))
+    
+    processing_time = time.time() - start_time
+    
+    return result_pil, processing_time, original_shape
 
 ###############################################################################
-# 6. GPU Gradio Function (Direct Decoration)
+# 8. GPU Gradio Function (Direct Decoration)
 ###############################################################################
 def gradio_processor(image_np):
     """Process image for Gradio interface"""
@@ -114,7 +233,8 @@ def gradio_processor(image_np):
         raise ValueError("Failed to encode image")
     image_bytes = img_encoded.tobytes()
     
-    return process_image_core(image_bytes, use_gpu=USE_SPACES_GPU)
+    result, _, _ = process_image_core(image_bytes, use_gpu=USE_SPACES_GPU)
+    return result
 
 # Apply GPU decorator directly
 if USE_SPACES_GPU:
@@ -127,59 +247,242 @@ else:
 print("🧠 gradio_processor has GPU metadata:", getattr(gradio_processor, "_spaces_gpu", None))
 
 ###############################################################################
-# 7. Gradio Interface
+# 9. FastAPI Endpoints
 ###############################################################################
-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 (GPU)",
-    description="Upload an image to remove its background using ZeroGPU-powered ONNX.",
-    flagging_options=None
-)
 
-interface.api_name = "remove_background_gpu"
+@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>
+    """
 
-###############################################################################
-# 8. FastAPI REST API (CPU Inference)
-###############################################################################
-@app.post("/api/remove-background")
+@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
+    )
+
+@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(
-    api_key: str = Form(...), 
-    image: UploadFile = 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"),
+    quality: int = Form(default=95, ge=1, le=100, description="JPEG quality (1-100, ignored for PNG)")
 ):
-    if api_key != API_KEY:
-        raise HTTPException(status_code=401, detail="Invalid API key")
-
+    """
+    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 = process_image_core(image_data, use_gpu=False)
+        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
+        )
 
-        result_filename = f"{uuid.uuid4()}.png"
+        # Save result
+        result_filename = f"{uuid.uuid4()}.{output_format.lower()}"
         output_path = f"{TMP_FOLDER}/{result_filename}"
-        result_img.save(output_path, "PNG")
+        
+        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")
 
-        return JSONResponse(content={
-            "status": "success",
-            "image_code": f"data:image/png;base64,{base64_image}"
-        })
+        # 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:
-        return JSONResponse(content={
-            "status": "failure",
-            "message": str(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`
+    """,
+    examples=[
+        # Add example images if you have them
+    ],
+    flagging_options=None,
+    allow_flagging="never"
+)
+
+interface.api_name = "remove_background_gpu"
 
 ###############################################################################
-# 9. Mount Gradio at Root Path (UI)
+# 11. Mount Gradio App
 ###############################################################################
-app = gr.mount_gradio_app(app, interface, path="/", ssr_mode=False)
+app = gr.mount_gradio_app(app, interface, path="/gradio", ssr_mode=False)
 
 ###############################################################################
-# 10. Run Local Dev Server
+# 12. Run Local Dev Server
 ###############################################################################
 if __name__ == "__main__":
-    uvicorn.run("main:app", host="0.0.0.0", port=7860, reload=True)
+    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", 
+        port=7860, 
+        reload=False  # Disable reload for production
+    )