Add compressed image support with 2-3MB target file size - Add Compressed_Image_URL field to colorize endpoint response - Implement compress_image_to_target_size function with iterative quality/dimension adjustment - Original download_url returns uncompressed PNG (model output) - Compressed_Image_URL returns JPEG compressed to 2-3MB file size

app/main.py

@@ -1,5 +1,5 @@
 from fastapi import FastAPI, File, UploadFile, HTTPException, Header, Request, Form, Depends, Body
-from typing import Optional
+from typing import Optional, Tuple
 from fastapi.responses import FileResponse
 from huggingface_hub import hf_hub_download
 import uuid
@@ -81,8 +81,10 @@ else:
 # -------------------------------------------------
 UPLOAD_DIR = "/tmp/uploads"
 RESULTS_DIR = "/tmp/results"
+COMPRESSED_DIR = "/tmp/compressed"
 os.makedirs(UPLOAD_DIR, exist_ok=True)
 os.makedirs(RESULTS_DIR, exist_ok=True)
+os.makedirs(COMPRESSED_DIR, exist_ok=True)
 
 MEDIA_CLICK_DEFAULT_CATEGORY = os.getenv("DEFAULT_CATEGORY_FALLBACK", "69368fcd2e46bd68ae1889b2")
 
@@ -182,6 +184,105 @@ def colorize_image(img: Image.Image, model_type: str = "gan", cco_model: str = "
     else:
         return colorize_image_gan(img)
 
+def compress_image_to_target_size(img: Image.Image, target_size_mb: float = 2.5, min_size_mb: float = 2.0, max_size_mb: float = 3.0) -> Tuple[Image.Image, str]:
+    """
+    Compress image to target file size (2-3MB) by adjusting quality and dimensions
+    
+    Args:
+        img: PIL Image to compress
+        target_size_mb: Target file size in MB (default: 2.5MB)
+        min_size_mb: Minimum acceptable file size in MB (default: 2.0MB)
+        max_size_mb: Maximum acceptable file size in MB (default: 3.0MB)
+    
+    Returns:
+        Tuple of (compressed PIL Image, path to saved compressed image)
+    """
+    import tempfile
+    
+    # Ensure image is RGB
+    if img.mode != "RGB":
+        img = img.convert("RGB")
+    
+    # Get original dimensions
+    original_size = img.size
+    current_img = img.copy()
+    
+    # Start with high quality and reduce if needed
+    quality = 85
+    scale_factor = 1.0
+    
+    # Try different quality levels and scales to achieve target size
+    for attempt in range(20):  # Max 20 attempts
+        # Create temporary file to check size
+        with tempfile.NamedTemporaryFile(suffix='.jpg', delete=False) as tmp_file:
+            temp_path = tmp_file.name
+        
+        try:
+            # Resize if needed
+            if scale_factor < 1.0:
+                new_width = int(original_size[0] * scale_factor)
+                new_height = int(original_size[1] * scale_factor)
+                resized_img = current_img.resize((new_width, new_height), Image.Resampling.LANCZOS)
+            else:
+                resized_img = current_img
+            
+            # Save with current quality
+            resized_img.save(temp_path, "JPEG", quality=quality, optimize=True)
+            
+            # Check file size
+            file_size_bytes = os.path.getsize(temp_path)
+            file_size_mb = file_size_bytes / (1024 * 1024)
+            
+            logger.info("Compression attempt %d: quality=%d, scale=%.2f, size=%.2f MB", 
+                       attempt + 1, quality, scale_factor, file_size_mb)
+            
+            # If size is within target range, we're done
+            if min_size_mb <= file_size_mb <= max_size_mb:
+                logger.info("Target size achieved: %.2f MB", file_size_mb)
+                return resized_img, temp_path
+            
+            # If too large, reduce quality or scale
+            if file_size_mb > max_size_mb:
+                if quality > 30:
+                    quality -= 5  # Reduce quality
+                elif scale_factor > 0.5:
+                    scale_factor -= 0.05  # Reduce dimensions
+                    quality = 75  # Reset quality when scaling
+                else:
+                    # Already at minimum, accept current result
+                    logger.warning("Reached minimum compression, file size: %.2f MB", file_size_mb)
+                    return resized_img, temp_path
+            
+            # If too small but still reasonable (above 1MB), try to increase quality slightly
+            elif file_size_mb < min_size_mb:
+                if file_size_mb < 1.0:
+                    # Very small file, increase quality more aggressively
+                    if quality < 90:
+                        quality += 5
+                    elif scale_factor < 1.0:
+                        # Increase scale if we reduced it
+                        scale_factor = min(1.0, scale_factor + 0.1)
+                elif quality < 95:
+                    # Close to target, fine-tune quality
+                    quality += 2
+                    if quality > 95:
+                        quality = 95
+                else:
+                    # Already at max quality and scale, accept current result
+                    logger.info("File size %.2f MB is below minimum but at max quality, accepting", file_size_mb)
+                    return resized_img, temp_path
+            
+        finally:
+            # Clean up temp file if we're continuing
+            if os.path.exists(temp_path) and attempt < 19:
+                try:
+                    os.unlink(temp_path)
+                except:
+                    pass
+    
+    # Return the last attempt's result
+    return resized_img, temp_path
+
 # -------------------------------------------------
 # 🗄️ MongoDB Initialization
 # -------------------------------------------------
@@ -737,8 +838,35 @@ async def colorize(
 
         result_id = f"{uuid.uuid4()}.png"
         output_path = os.path.join(RESULTS_DIR, result_id)
+        # Save original image as PNG (uncompressed) - this is what the model produces
         output_img.save(output_path, "PNG")
 
+        # Create compressed version targeting 2-3MB file size
+        logger.info("Creating compressed version targeting 2-3MB file size...")
+        compressed_img, temp_compressed_path = compress_image_to_target_size(
+            output_img, 
+            target_size_mb=2.5, 
+            min_size_mb=2.0, 
+            max_size_mb=3.0
+        )
+        
+        # Move compressed image to final location
+        compressed_filename = result_id.replace(".png", "_compressed.jpg")
+        compressed_path = os.path.join(COMPRESSED_DIR, compressed_filename)
+        
+        # If temp file exists, move it; otherwise save the compressed image
+        if os.path.exists(temp_compressed_path):
+            import shutil
+            shutil.move(temp_compressed_path, compressed_path)
+        else:
+            compressed_img.save(compressed_path, "JPEG", quality=75, optimize=True)
+        
+        # Log compressed file size
+        compressed_size_mb = os.path.getsize(compressed_path) / (1024 * 1024)
+        original_size_mb = os.path.getsize(output_path) / (1024 * 1024)
+        logger.info("Original image size: %.2f MB, Compressed image size: %.2f MB", 
+                   original_size_mb, compressed_size_mb)
+
         base_url = "https://logicgoinfotechspaces-text-guided-image-colorization.hf.space"
 
         result_id_clean = result_id.replace(".png", "")
@@ -752,7 +880,8 @@ async def colorize(
             "download_url": f"{base_url}/results/{result_id}",
             "api_download_url": f"{base_url}/download/{result_id_clean}",
             "filename": result_id,
-            "caption": caption
+            "caption": caption,
+            "Compressed_Image_URL": f"{base_url}/compressed/{compressed_filename}"
         }
         
         # Log to MongoDB (colorization_db -> colorizations)
@@ -915,6 +1044,14 @@ def get_upload(request: Request, filename: str):
         )
         raise HTTPException(status_code=404, detail="File 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/jpeg"  # Default to JPEG
+    
     log_api_call(
         endpoint=f"/uploads/{filename}",
         method="GET",
@@ -924,3 +1061,34 @@ def get_upload(request: Request, filename: str):
     )
     
     return FileResponse(path, media_type=media_type)
+
+# -------------------------------------------------
+# 🌐 Public Compressed File
+# -------------------------------------------------
+@app.get("/compressed/{filename}")
+def get_compressed(request: Request, filename: str):
+    ip_address = request.client.host if request.client else None
+    
+    path = os.path.join(COMPRESSED_DIR, filename)
+    if not os.path.exists(path):
+        log_api_call(
+            endpoint=f"/compressed/{filename}",
+            method="GET",
+            status_code=404,
+            error="Compressed file not found",
+            ip_address=ip_address
+        )
+        raise HTTPException(status_code=404, detail="Compressed file not found")
+    
+    # Compressed images are JPEG
+    media_type = "image/jpeg"
+    
+    log_api_call(
+        endpoint=f"/compressed/{filename}",
+        method="GET",
+        status_code=200,
+        request_data={"filename": filename},
+        ip_address=ip_address
+    )
+    
+    return FileResponse(path, media_type=media_type)