Update response formats for /colorize and /health endpoints - /colorize: Add filename, caption, change to PNG format, api_download_url - /health: Add model_type and provider fields

app/main.py

@@ -190,7 +190,20 @@ async def shutdown_event():
 # -------------------------------------------------
 @app.get("/health")
 def health_check(request: Request):
-    response = {"status": "healthy", "model_loaded": True}
+    # Determine model type based on what's available
+    if CCO_AVAILABLE:
+        model_type = "cco-eccv16"
+        provider = "local"
+    else:
+        model_type = "gan"
+        provider = "local"
+    
+    response = {
+        "status": "healthy",
+        "model_loaded": True,
+        "model_type": model_type,
+        "provider": provider
+    }
     
     # Log API call
     log_api_call(
@@ -408,19 +421,27 @@ async def colorize(
 
         processing_time = time.time() - start_time
 
-        result_id = f"{uuid.uuid4()}.jpg"
+        result_id = f"{uuid.uuid4()}.png"
         output_path = os.path.join(RESULTS_DIR, result_id)
-        output_img.save(output_path, "JPEG", quality=95)
+        output_img.save(output_path, "PNG")
 
         base_url = "https://logicgoinfotechspaces-text-guided-image-colorization.hf.space"
 
-        result_id_clean = result_id.replace(".jpg", "")
+        result_id_clean = result_id.replace(".png", "")
+        
+        # Generate caption based on model used
+        if model_type == "cco":
+            caption = f"Colorized using CCO {cco_model} model with vibrant natural colors, high quality"
+        else:
+            caption = "Colorized using GAN model with vibrant natural colors, high quality"
         
         response_data = {
             "success": True,
             "result_id": result_id_clean,
             "download_url": f"{base_url}/results/{result_id}",
-            "api_download": f"{base_url}/download/{result_id_clean}"
+            "api_download_url": f"{base_url}/download/{result_id_clean}",
+            "filename": result_id,
+            "caption": caption
         }
         
         # Log to MongoDB (colorization_db -> colorizations)
@@ -489,8 +510,25 @@ def download_result(
     
     ip_address = request.client.host if request.client else None
 
-    filename = f"{file_id}.jpg"
-    path = os.path.join(RESULTS_DIR, filename)
+    # Try PNG first, then JPG for backward compatibility
+    filename_png = f"{file_id}.png"
+    filename_jpg = f"{file_id}.jpg"
+    path_png = os.path.join(RESULTS_DIR, filename_png)
+    path_jpg = os.path.join(RESULTS_DIR, filename_jpg)
+    
+    # Check which file exists
+    if os.path.exists(path_png):
+        filename = filename_png
+        path = path_png
+        media_type = "image/png"
+    elif os.path.exists(path_jpg):
+        filename = filename_jpg
+        path = path_jpg
+        media_type = "image/jpeg"
+    else:
+        filename = filename_png
+        path = path_png
+        media_type = "image/png"
 
     if not os.path.exists(path):
         log_api_call(
@@ -510,7 +548,7 @@ def download_result(
         ip_address=ip_address
     )
 
-    return FileResponse(path, media_type="image/jpeg")
+    return FileResponse(path, media_type=media_type)
 
 # -------------------------------------------------
 # 🌐 Public Result File
@@ -530,6 +568,14 @@ def get_result(request: Request, filename: str):
         )
         raise HTTPException(status_code=404, detail="Result not found")
     
+    # Determine media type based on file extension
+    if filename.lower().endswith('.png'):
+        media_type = "image/png"
+    elif filename.lower().endswith('.jpg') or filename.lower().endswith('.jpeg'):
+        media_type = "image/jpeg"
+    else:
+        media_type = "image/png"  # Default to PNG
+    
     log_api_call(
         endpoint=f"/results/{filename}",
         method="GET",
@@ -538,7 +584,7 @@ def get_result(request: Request, filename: str):
         ip_address=ip_address
     )
     
-    return FileResponse(path, media_type="image/jpeg")
+    return FileResponse(path, media_type=media_type)
 
 # -------------------------------------------------
 # 🌐 Public Uploaded File
@@ -566,4 +612,4 @@ def get_upload(request: Request, filename: str):
         ip_address=ip_address
     )
     
-    return FileResponse(path, media_type="image/jpeg")
+    return FileResponse(path, media_type=media_type)