Add simple colorization fallback using LAB color space when model fails to load

app/main_fastai.py

@@ -28,6 +28,8 @@ from PIL import Image
 import torch
 import uvicorn
 import gradio as gr
+import numpy as np
+import cv2
 
 # FastAI imports
 from fastai.vision.all import *
@@ -211,6 +213,40 @@ async def health_check():
         response["model_error"] = model_load_error
     return response
 
+def simple_colorize_fallback(image: Image.Image) -> Image.Image:
+    """
+    Simple fallback colorization using LAB color space
+    This provides basic colorization when the model doesn't load
+    """
+    # Convert to LAB color space
+    if image.mode != "RGB":
+        image = image.convert("RGB")
+    
+    # Convert to numpy array
+    img_array = np.array(image)
+    
+    # Convert RGB to LAB
+    lab = cv2.cvtColor(img_array, cv2.COLOR_RGB2LAB)
+    
+    # Apply simple colorization: enhance the L channel and add some color hints
+    l, a, b = cv2.split(lab)
+    
+    # Enhance lightness
+    clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8,8))
+    l = clahe.apply(l)
+    
+    # Add subtle color hints to a and b channels
+    # This is a very basic approach - just adds some warm tones
+    a = np.clip(a.astype(np.float32) + 5, 0, 255).astype(np.uint8)
+    b = np.clip(b.astype(np.float32) + 10, 0, 255).astype(np.uint8)
+    
+    # Merge channels and convert back to RGB
+    lab_colored = cv2.merge([l, a, b])
+    colored_rgb = cv2.cvtColor(lab_colored, cv2.COLOR_LAB2RGB)
+    
+    return Image.fromarray(colored_rgb)
+
+
 def colorize_pil(image: Image.Image) -> Image.Image:
     """Run model prediction and return colorized image"""
     # Try FastAI first
@@ -251,7 +287,9 @@ def colorize_pil(image: Image.Image) -> Image.Image:
         return pytorch_colorizer.colorize(image)
     
     else:
-        raise RuntimeError("No colorization model loaded")
+        # Final fallback: simple colorization
+        logger.warning("No model loaded, using simple colorization fallback")
+        return simple_colorize_fallback(image)
 
 @app.post("/colorize")
 async def colorize_api(
@@ -262,8 +300,9 @@ async def colorize_api(
     Upload a black & white image -> returns colorized image.
     Requires Firebase authentication unless DISABLE_AUTH=true
     """
-    if learn is None and pytorch_colorizer is None:
-        raise HTTPException(status_code=503, detail="Colorization model not loaded")
+    # Allow fallback colorization even if model isn't loaded
+    # if learn is None and pytorch_colorizer is None:
+    #     raise HTTPException(status_code=503, detail="Colorization model not loaded")
     
     if not file.content_type or not file.content_type.startswith("image/"):
         raise HTTPException(status_code=400, detail="File must be an image")
@@ -299,8 +338,7 @@ def gradio_colorize(image):
     if image is None:
         return None
     try:
-        if learn is None and pytorch_colorizer is None:
-            return None
+        # Always try to colorize, even with fallback
         return colorize_pil(image)
     except Exception as e:
         logger.error("Gradio colorization error: %s", str(e))

requirements.txt

@@ -7,4 +7,6 @@ pillow
 firebase-admin
 fastai
 huggingface_hub
-pydantic-settings
+pydantic-settings
+opencv-python
+numpy