Update app.py

app.py

@@ -1,17 +1,25 @@
 import os
 import cv2
-import uuid
 import numpy
 import base64
 from io import BytesIO
 from PIL import Image
-from flask import Flask, request, jsonify
+from fastapi import FastAPI, File, Form, UploadFile, HTTPException
+from fastapi.responses import JSONResponse
+from typing import Optional, List, Dict, Any, Union
+from pydantic import BaseModel
+import uvicorn
 from basicsr.archs.rrdbnet_arch import RRDBNet
 from basicsr.utils.download_util import load_file_from_url
 from realesrgan import RealESRGANer
 from realesrgan.archs.srvgg_arch import SRVGGNetCompact
 
-app = Flask(__name__)
+# Create FastAPI app
+app = FastAPI(
+    title="Image Enhancement API",
+    description="API for enhancing and upscaling images using Real-ESRGAN models",
+    version="1.0.0"
+)
 
 # Create weights directory if it doesn't exist
 os.makedirs('weights', exist_ok=True)
@@ -19,7 +27,59 @@ os.makedirs('weights', exist_ok=True)
 # Global variable to track image mode
 img_mode = "RGBA"
 
-def process_image(img_data, model_name, denoise_strength, face_enhance, outscale):
+# Models information
+AVAILABLE_MODELS = [
+    {
+        "name": "RealESRGAN_x4plus",
+        "description": "General purpose 4x upscaling model",
+        "scale": 4
+    },
+    {
+        "name": "RealESRNet_x4plus",
+        "description": "Alternative 4x upscaling model",
+        "scale": 4
+    },
+    {
+        "name": "RealESRGAN_x4plus_anime_6B",
+        "description": "Specialized for anime/cartoon images, 4x upscaling",
+        "scale": 4
+    },
+    {
+        "name": "RealESRGAN_x2plus",
+        "description": "2x upscaling model",
+        "scale": 2
+    },
+    {
+        "name": "realesr-general-x4v3",
+        "description": "General purpose 4x upscaling model with denoise control",
+        "scale": 4
+    }
+]
+
+# Pydantic models for API documentation
+class HealthResponse(BaseModel):
+    status: str
+    message: str
+
+class ModelInfo(BaseModel):
+    name: str
+    description: str
+    scale: int
+
+class ModelsResponse(BaseModel):
+    models: List[ModelInfo]
+
+class ImageProperties(BaseModel):
+    width: int
+    height: int
+    mode: str
+
+class EnhancementResponse(BaseModel):
+    enhanced_image: str
+    properties: ImageProperties
+    model_used: str
+
+async def process_image(img_data, model_name, denoise_strength, face_enhance, outscale):
     """Real-ESRGAN function to restore (and upscale) images."""
     global img_mode
     
@@ -48,7 +108,7 @@ def process_image(img_data, model_name, denoise_strength, face_enhance, outscale
             'https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-x4v3.pth'
         ]
     else:
-        return None, "Invalid model name"
+        raise HTTPException(status_code=400, detail=f"Invalid model name: {model_name}")
 
     # Download model if not already available
     model_path = os.path.join('weights', model_name + '.pth')
@@ -113,7 +173,7 @@ def process_image(img_data, model_name, denoise_strength, face_enhance, outscale
         else:
             output, _ = upsampler.enhance(img, outscale=outscale)
     except RuntimeError as error:
-        return None, f"Error: {str(error)}"
+        raise HTTPException(status_code=500, detail=f"Processing error: {str(error)}")
     
     # Convert back to appropriate format based on mode
     if img_mode == "RGBA":
@@ -133,30 +193,47 @@ def process_image(img_data, model_name, denoise_strength, face_enhance, outscale
     
     return output_img, properties
 
-@app.route('/enhancer', methods=['POST'])
-def enhance_image():
+@app.post("/enhancer", response_model=EnhancementResponse, summary="Enhance and upscale an image")
+async def enhance_image(
+    image: UploadFile = File(..., description="Image file to enhance"),
+    model: str = Form("RealESRGAN_x4plus", description="Model name to use for enhancement"),
+    denoise_strength: float = Form(0.5, description="Denoise strength (0-1)"),
+    outscale: int = Form(4, description="Output scale factor"),
+    face_enhance: bool = Form(False, description="Enable face enhancement")
+):
+    """
+    Enhance and upscale an image using Real-ESRGAN models.
+    
+    - **image**: Upload an image file (PNG, JPG, etc.)
+    - **model**: Select a model from the available options
+    - **denoise_strength**: Control the denoising strength (only for realesr-general-x4v3)
+    - **outscale**: Control the output resolution scaling
+    - **face_enhance**: Enable face enhancement using GFPGAN
+    
+    Returns the enhanced image as a base64 string along with image properties.
+    """
     try:
-        # Check if file is present in request
-        if 'image' not in request.files:
-            return jsonify({"error": "No image file provided"}), 400
+        # Validate model name
+        valid_models = [m["name"] for m in AVAILABLE_MODELS]
+        if model not in valid_models:
+            raise HTTPException(
+                status_code=400, 
+                detail=f"Invalid model. Choose from: {', '.join(valid_models)}"
+            )
         
-        # Get image file
-        image_file = request.files['image']
+        # Validate other parameters
+        if not (0 <= denoise_strength <= 1):
+            raise HTTPException(status_code=400, detail="Denoise strength must be between 0 and 1")
         
-        # Get parameters with defaults
-        model_name = request.form.get('model', 'RealESRGAN_x4plus')
-        denoise_strength = float(request.form.get('denoise_strength', 0.5))
-        outscale = int(request.form.get('outscale', 4))
-        face_enhance = request.form.get('face_enhance', 'false').lower() == 'true'
+        if not (1 <= outscale <= 8):
+            raise HTTPException(status_code=400, detail="Outscale must be between 1 and 8")
         
-        # Open image using PIL
-        img = Image.open(image_file)
+        # Read the image file
+        contents = await image.read()
+        img = Image.open(BytesIO(contents))
         
         # Process image
-        output_img, properties = process_image(img, model_name, denoise_strength, face_enhance, outscale)
-        
-        if output_img is None:
-            return jsonify({"error": properties}), 500
+        output_img, properties = await process_image(img, model, denoise_strength, face_enhance, outscale)
         
         # Convert to PIL Image and then to base64
         output_pil = Image.fromarray(output_img)
@@ -172,53 +249,35 @@ def enhance_image():
         img_str = base64.b64encode(buffer.getvalue()).decode('utf-8')
         
         # Return response
-        return jsonify({
+        return {
             "enhanced_image": img_str,
             "properties": properties,
-            "model_used": model_name
-        })
+            "model_used": model
+        }
     
+    except HTTPException as e:
+        raise e
     except Exception as e:
-        return jsonify({"error": str(e)}), 500
-
-# Add a simple health check route
-@app.route('/', methods=['GET'])
-def health_check():
-    return jsonify({"status": "healthy", "message": "Image enhancement server is running"})
-
-# Add a route to list available models
-@app.route('/models', methods=['GET'])
-def list_models():
-    models = [
-        {
-            "name": "RealESRGAN_x4plus",
-            "description": "General purpose 4x upscaling model",
-            "scale": 4
-        },
-        {
-            "name": "RealESRNet_x4plus",
-            "description": "Alternative 4x upscaling model",
-            "scale": 4
-        },
-        {
-            "name": "RealESRGAN_x4plus_anime_6B",
-            "description": "Specialized for anime/cartoon images, 4x upscaling",
-            "scale": 4
-        },
-        {
-            "name": "RealESRGAN_x2plus",
-            "description": "2x upscaling model",
-            "scale": 2
-        },
-        {
-            "name": "realesr-general-x4v3",
-            "description": "General purpose 4x upscaling model with denoise control",
-            "scale": 4
-        }
-    ]
-    return jsonify({"models": models})
+        raise HTTPException(status_code=500, detail=f"Server error: {str(e)}")
+
+@app.get("/", response_model=HealthResponse, summary="Check server health")
+async def health_check():
+    """Check if the image enhancement server is running."""
+    return {"status": "healthy", "message": "Image enhancement server is running"}
+
+@app.get("/models", response_model=ModelsResponse, summary="List available models")
+async def list_models():
+    """Get a list of all available enhancement models with descriptions."""
+    return {"models": AVAILABLE_MODELS}
+
+# Add startup event to print server info
+@app.on_event("startup")
+async def startup_event():
+    print("🚀 Image Enhancement API is starting up!")
+    print(f"📚 Available models: {', '.join(m['name'] for m in AVAILABLE_MODELS)}")
+    print("📋 API documentation available at /docs or /redoc")
 
 if __name__ == "__main__":
-    # Run server with debug mode enabled
+    # Run server with Uvicorn
     port = int(os.environ.get("PORT", 8000))
-    app.run(host="0.0.0.0", port=port)
+    uvicorn.run("app:app", host="0.0.0.0", port=port, reload=True)