Update app.py

app.py

@@ -45,7 +45,7 @@ logger = logging.getLogger(__name__)
 os.environ["U2NET_HOME"] = "/tmp/u2net_models"
 os.environ["OMP_NUM_THREADS"] = "1"
 os.environ["CUDA_VISIBLE_DEVICES"] = ""
-os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"  # Suppress TF warnings
+os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"
 os.environ["PYTHONUNBUFFERED"] = "1"
 
 # Create necessary directories
@@ -86,10 +86,8 @@ def download_model_with_retry(model_name: str, max_retries: int = 3) -> Optional
 try:
     logger.info("🚀 Initializing AI models...")
     
-    # Import rembg only when needed
     from rembg import remove, new_session
     
-    # Model configurations with priorities
     model_configs = [
         {"name": "u2net", "description": "Best quality - general purpose", "size": "176MB", "priority": 1},
         {"name": "u2netp", "description": "Fast processing - lightweight", "size": "4.7MB", "priority": 2},
@@ -97,10 +95,8 @@ try:
         {"name": "u2net_human_seg", "description": "Specialized for portraits", "size": "176MB", "priority": 4},
     ]
     
-    # Sort by priority
     model_configs.sort(key=lambda x: x["priority"])
     
-    # Load models
     for config in model_configs:
         session = download_model_with_retry(config["name"])
         if session:
@@ -134,7 +130,6 @@ app = FastAPI(
     version="8.2.0"
 )
 
-# CORS configuration
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -146,29 +141,24 @@ app.add_middleware(
 # ========== UTILITY FUNCTIONS ==========
 
 def validate_image_file(file: UploadFile) -> None:
-    """Validate image file upload"""
     if not file.content_type or not file.content_type.startswith('image/'):
         raise HTTPException(400, f"File must be an image, got {file.content_type}")
     
-    # Check file extension
     allowed_extensions = {'.jpg', '.jpeg', '.png', '.webp', '.bmp', '.tiff', '.tif'}
     file_ext = Path(file.filename).suffix.lower()
     if file_ext not in allowed_extensions:
         raise HTTPException(400, f"Unsupported image format: {file_ext}")
 
 def validate_video_file(file: UploadFile) -> None:
-    """Validate video file upload"""
     if not file.content_type or not file.content_type.startswith('video/'):
         raise HTTPException(400, f"File must be a video, got {file.content_type}")
     
-    # Check file extension
     allowed_extensions = {'.mp4', '.avi', '.mov', '.webm', '.mkv', '.flv', '.wmv'}
     file_ext = Path(file.filename).suffix.lower()
     if file_ext not in allowed_extensions:
         raise HTTPException(400, f"Unsupported video format: {file_ext}")
 
 def cleanup_temp_files(*file_paths: str) -> None:
-    """Clean up temporary files"""
     for path in file_paths:
         if path and os.path.exists(path):
             try:
@@ -178,7 +168,6 @@ def cleanup_temp_files(*file_paths: str) -> None:
                 logger.warning(f"Failed to cleanup {path}: {e}")
 
 def get_file_size_mb(file_path: str) -> float:
-    """Get file size in megabytes"""
     if os.path.exists(file_path):
         return os.path.getsize(file_path) / (1024 * 1024)
     return 0.0
@@ -192,16 +181,12 @@ def process_image_sync(
     max_size: int = 1024,
     quality: int = 95
 ) -> Dict:
-    """Process single image synchronously"""
     try:
-        # Open and validate image
         image = Image.open(io.BytesIO(image_data))
         
-        # Convert to RGB if necessary
         if image.mode not in ['RGB', 'RGBA']:
             image = image.convert('RGB')
         
-        # Select model (fallback to first available)
         if model_name not in models:
             available_models = list(models.keys())
             if not available_models:
@@ -211,30 +196,24 @@ def process_image_sync(
         
         session = models[model_name]["session"]
         
-        # Resize if too large
         original_size = image.size
         if max(image.size) > max_size:
             image.thumbnail((max_size, max_size), Image.Resampling.LANCZOS)
             logger.info(f"Resized from {original_size} to {image.size}")
         
-        # Ensure RGBA for transparency support
         if image.mode != 'RGBA':
             image = image.convert('RGBA')
         
-        # Remove background
         result = remove_func(image, session=session)
         
-        # Handle transparency
         if not transparent:
             if result.mode != 'RGBA':
                 result = result.convert('RGBA')
             white_bg = Image.new('RGBA', result.size, (255, 255, 255, 255))
             result = Image.alpha_composite(white_bg, result)
         
-        # Optimize and save
         output_buffer = io.BytesIO()
         
-        # Convert to RGB for JPEG if not transparent
         if not transparent:
             result_rgb = result.convert('RGB')
             result_rgb.save(output_buffer, format="JPEG", quality=quality, optimize=True)
@@ -268,9 +247,7 @@ def process_frame_with_fallback(
     model_name: str = "silueta",
     transparent: bool = True
 ) -> Optional[np.ndarray]:
-    """Process a single frame with robust error handling"""
     try:
-        # Select available model
         if model_name not in models:
             available_models = list(models.keys())
             if not available_models:
@@ -279,7 +256,6 @@ def process_frame_with_fallback(
         
         session = models[model_name]["session"]
         
-        # Convert BGR to RGB for PIL
         if len(frame.shape) == 3:
             if frame.shape[2] == 3:
                 frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
@@ -288,33 +264,24 @@ def process_frame_with_fallback(
             else:
                 return None
         else:
-            # Grayscale to RGB
             frame_rgb = cv2.cvtColor(frame, cv2.COLOR_GRAY2RGB)
         
-        # Convert to PIL Image
         img = Image.fromarray(frame_rgb)
         if img.mode != 'RGBA':
             img = img.convert('RGBA')
         
-        # Remove background
         result = remove_func(img, session=session)
-        
-        # Convert back to numpy
         result_array = np.array(result)
         
         if transparent:
             if result_array.shape[2] == 4:
-                # RGBA to BGRA for OpenCV
                 return cv2.cvtColor(result_array, cv2.COLOR_RGBA2BGRA)
             else:
-                # Add alpha channel
                 bgr = cv2.cvtColor(result_array, cv2.COLOR_RGB2BGR)
                 alpha = np.ones((bgr.shape[0], bgr.shape[1], 1), dtype=np.uint8) * 255
                 return np.concatenate([bgr, alpha], axis=2)
         else:
-            # Add white background
             if result_array.shape[2] == 4:
-                # Create white background
                 alpha = result_array[:, :, 3:4] / 255.0
                 white_bg = np.ones_like(result_array[:, :, :3]) * 255
                 blended = (1 - alpha) * white_bg + alpha * result_array[:, :, :3]
@@ -324,7 +291,6 @@ def process_frame_with_fallback(
             
     except Exception as e:
         logger.error(f"Frame processing failed: {e}")
-        # Return original frame with alpha if needed
         if transparent and len(frame.shape) == 3 and frame.shape[2] == 3:
             return cv2.cvtColor(frame, cv2.COLOR_BGR2BGRA)
         return frame
@@ -338,9 +304,7 @@ def create_video_with_ffmpeg(
     height: int = 360,
     transparent: bool = True
 ) -> Tuple[bool, str]:
-    """Create video using ffmpeg with specified parameters"""
     try:
-        # Check if ffmpeg is available
         result = subprocess.run(['which', 'ffmpeg'], capture_output=True, text=True)
         if result.returncode != 0:
             logger.warning("ffmpeg not found, using OpenCV fallback")
@@ -348,10 +312,8 @@ def create_video_with_ffmpeg(
         
         actual_output_path = output_path
         
-        # Build ffmpeg command based on format and transparency
         if format.lower() in ['mp4', 'mov']:
             if transparent:
-                # MP4 doesn't support transparency well, use WebM instead
                 logger.info("MP4 doesn't support alpha channel, switching to WebM format")
                 actual_output_path = output_path.replace('.mp4', '.webm').replace('.mov', '.webm')
                 cmd = [
@@ -438,31 +400,25 @@ def process_video_reliable(
     transparent: bool = True,
     format: str = "mp4"
 ) -> Dict:
-    """Process video with reliable output"""
     try:
         logger.info(f"Processing video: {input_path} -> {output_path}")
         logger.info(f"Settings: model={model_name}, size={max_size}, fps={fps}, format={format}")
         
-        # Open video
         cap = cv2.VideoCapture(input_path)
         if not cap.isOpened():
             raise ValueError(f"Cannot open video: {input_path}")
         
-        # Get video properties
         width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
         height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
         orig_fps = cap.get(cv2.CAP_PROP_FPS)
         total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
         
-        # Calculate FPS
         target_fps = min(orig_fps if orig_fps > 0 else fps, 15)
         
-        # Calculate dimensions
         scale = min(max_size / max(width, height), 1.0)
         new_width = int(width * scale)
         new_height = int(height * scale)
         
-        # Ensure even dimensions
         if new_width % 2 != 0:
             new_width += 1
         if new_height % 2 != 0:
@@ -471,15 +427,13 @@ def process_video_reliable(
         logger.info(f"Video: {width}x{height} -> {new_width}x{new_height}")
         logger.info(f"FPS: {target_fps}, Total frames: {total_frames}")
         
-        # Create temp directory for frames
         frames_dir = tempfile.mkdtemp(prefix="video_frames_")
         logger.info(f"Frames directory: {frames_dir}")
         
-        # Process frames
         frame_count = 0
         processed_count = 0
         start_time = time.time()
-        max_processing_time = 300  # 5 minutes timeout
+        max_processing_time = 300
         
         while True:
             ret, frame = cap.read()
@@ -488,36 +442,28 @@ def process_video_reliable(
             
             frame_count += 1
             
-            # Skip frames if configured
             if frame_skip > 0 and frame_count % (frame_skip + 1) != 0:
                 continue
             
-            # Check timeout
             if time.time() - start_time > max_processing_time:
                 logger.warning(f"Processing timeout reached after {processed_count} frames")
                 break
             
             try:
-                # Resize frame
                 frame_resized = cv2.resize(frame, (new_width, new_height))
-                
-                # Process frame
                 processed_frame = process_frame_with_fallback(frame_resized, model_name, transparent)
                 
                 if processed_frame is not None:
-                    # Save frame
                     frame_path = os.path.join(frames_dir, f"frame_{processed_count:06d}.png")
                     cv2.imwrite(frame_path, processed_frame)
                     processed_count += 1
                 
-                # Log progress
                 if processed_count % 10 == 0:
                     elapsed = time.time() - start_time
                     logger.info(f"Processed {processed_count} frames ({elapsed:.1f}s)")
                     
             except Exception as e:
                 logger.error(f"Frame {frame_count} failed: {e}")
-                # Save original frame as fallback
                 try:
                     frame_path = os.path.join(frames_dir, f"frame_{processed_count:06d}.png")
                     frame_resized = cv2.resize(frame, (new_width, new_height))
@@ -535,7 +481,6 @@ def process_video_reliable(
         if processed_count == 0:
             raise ValueError("No frames processed")
         
-        # Create video
         actual_format = format
         if transparent and format.lower() in ['mp4', 'mov']:
             actual_format = 'webm'
@@ -547,7 +492,6 @@ def process_video_reliable(
         )
         
         if not ffmpeg_success:
-            # Fallback to OpenCV
             logger.info("Using OpenCV fallback")
             if transparent and actual_format.lower() == "avi":
                 fourcc = cv2.VideoWriter_fourcc(*'FFV1')
@@ -559,11 +503,9 @@ def process_video_reliable(
             out = cv2.VideoWriter(actual_output_path, fourcc, target_fps, (new_width, new_height), True)
             
             if not out.isOpened():
-                # Last resort fallback
                 fourcc = cv2.VideoWriter_fourcc(*'MJPG')
                 out = cv2.VideoWriter(actual_output_path, fourcc, target_fps, (new_width, new_height))
             
-            # Write frames
             for i in range(processed_count):
                 frame_path = os.path.join(frames_dir, f"frame_{i:06d}.png")
                 if os.path.exists(frame_path):
@@ -573,7 +515,6 @@ def process_video_reliable(
             
             out.release()
         
-        # Cleanup frames directory
         try:
             shutil.rmtree(frames_dir)
         except:
@@ -581,7 +522,6 @@ def process_video_reliable(
         
         processing_time = time.time() - start_time
         
-        # Verify output
         if not os.path.exists(actual_output_path) or os.path.getsize(actual_output_path) < 1024:
             logger.error(f"Output file too small or missing: {actual_output_path}")
             create_fallback_video(actual_output_path, new_width, new_height, target_fps)
@@ -608,20 +548,17 @@ def process_video_reliable(
         raise
 
 def create_fallback_video(output_path: str, width: int, height: int, fps: int):
-    """Create a simple video with error message"""
     try:
         fourcc = cv2.VideoWriter_fourcc(*'mp4v')
         out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
         
-        for i in range(fps * 5):  # 5 second video
+        for i in range(fps * 5):
             frame = np.zeros((height, width, 3), dtype=np.uint8)
             
-            # Gradient background
             for y in range(height):
                 color_value = int(40 + (y / height) * 40)
                 cv2.rectangle(frame, (0, y), (width, y+1), (color_value, color_value, color_value+20), -1)
             
-            # Add text
             if i < fps * 2:
                 text = "Processing Video..."
             else:
@@ -647,7 +584,6 @@ def create_fallback_video(output_path: str, width: int, height: int, fps: int):
 
 @app.get("/")
 async def home():
-    """Home page with API documentation"""
     html = """
     <!DOCTYPE html>
     <html lang="en">
@@ -656,12 +592,7 @@ async def home():
         <meta name="viewport" content="width=device-width, initial-scale=1.0">
         <title>🎬 Background Remover Pro</title>
         <style>
-            * {
-                margin: 0;
-                padding: 0;
-                box-sizing: border-box;
-            }
-            
+            * { margin: 0; padding: 0; box-sizing: border-box; }
             body {
                 font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, Oxygen, Ubuntu, sans-serif;
                 background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
@@ -673,7 +604,6 @@ async def home():
                 align-items: center;
                 justify-content: center;
             }
-            
             .container {
                 max-width: 900px;
                 margin: 0 auto;
@@ -684,7 +614,6 @@ async def home():
                 box-shadow: 0 20px 60px rgba(0, 0, 0, 0.3);
                 border: 1px solid rgba(255, 255, 255, 0.2);
             }
-            
             h1 {
                 font-size: 2.5em;
                 margin-bottom: 10px;
@@ -693,14 +622,12 @@ async def home():
                 -webkit-text-fill-color: transparent;
                 text-align: center;
             }
-            
             .subtitle {
                 text-align: center;
                 color: #cbd5e0;
                 margin-bottom: 30px;
                 font-size: 1.1em;
             }
-            
             .status {
                 background: linear-gradient(135deg, rgba(72, 187, 120, 0.2), rgba(72, 187, 120, 0.1));
                 padding: 15px 20px;
@@ -709,25 +636,17 @@ async def home():
                 margin: 20px 0;
                 border: 1px solid rgba(72, 187, 120, 0.3);
             }
-            
-            .status.warning {
-                background: linear-gradient(135deg, rgba(237, 137, 54, 0.2), rgba(237, 137, 54, 0.1));
-                border: 1px solid rgba(237, 137, 54, 0.3);
-            }
-            
             h2 {
                 margin: 25px 0 15px 0;
                 color: #e0e7ff;
                 font-size: 1.5em;
             }
-            
             .endpoint-grid {
                 display: grid;
                 grid-template-columns: repeat(auto-fit, minmax(250px, 1fr));
                 gap: 15px;
                 margin: 20px 0;
             }
-            
             .endpoint {
                 background: rgba(255, 255, 255, 0.05);
                 padding: 20px;
@@ -735,19 +654,16 @@ async def home():
                 border-left: 4px solid #667eea;
                 transition: transform 0.2s, background 0.2s;
             }
-            
             .endpoint:hover {
                 background: rgba(255, 255, 255, 0.1);
                 transform: translateY(-2px);
             }
-            
             .endpoint strong {
                 display: block;
                 font-size: 1.1em;
                 margin-bottom: 8px;
                 color: #a78bfa;
             }
-            
             .endpoint code {
                 background: rgba(0, 0, 0, 0.3);
                 padding: 4px 8px;
@@ -756,21 +672,18 @@ async def home():
                 display: inline-block;
                 margin: 5px 0;
             }
-            
             .feature-list {
                 display: grid;
                 grid-template-columns: repeat(auto-fit, minmax(200px, 1fr));
                 gap: 10px;
                 margin: 20px 0;
             }
-            
             .feature {
                 background: rgba(255, 255, 255, 0.05);
                 padding: 12px;
                 border-radius: 8px;
                 text-align: center;
             }
-            
             .footer {
                 text-align: center;
                 margin-top: 30px;
@@ -778,7 +691,6 @@ async def home():
                 border-top: 1px solid rgba(255, 255, 255, 0.1);
                 color: #a0aec0;
             }
-            
             .badge {
                 display: inline-block;
                 padding: 4px 10px;
@@ -787,15 +699,9 @@ async def home():
                 font-size: 0.85em;
                 margin: 5px 5px 0 0;
             }
-            
             @media (max-width: 768px) {
-                .container {
-                    padding: 20px;
-                }
-                
-                h1 {
-                    font-size: 2em;
-                }
+                .container { padding: 20px; }
+                h1 { font-size: 2em; }
             }
         </style>
     </head>
@@ -879,7 +785,6 @@ async def home():
 
 @app.get("/api/health")
 async def health_check():
-    """Health check endpoint"""
     return {
         "status": "healthy" if AI_AVAILABLE else "degraded",
         "ai_available": AI_AVAILABLE,
@@ -895,7 +800,6 @@ async def health_check():
 
 @app.get("/api/models")
 async def list_models():
-    """List available AI models"""
     model_info = {}
     for name, info in models.items():
         model_info[name] = {
@@ -919,18 +823,15 @@ async def process_image_api(
     max_size: int = Query(1024, ge=256, le=2048, description="Maximum dimension"),
     quality: int = Query(95, ge=50, le=100, description="Output quality (JPEG only)")
 ):
-    """Remove background from an image"""
     start_time = time.time()
     
     try:
-        # Validate
         validate_image_file(file)
         
-        # Read file
         contents = await file.read()
         if len(contents) == 0:
             raise HTTPException(400, "Empty file")
-        if len(contents) > 20 * 1024 * 1024:  # 20MB limit
+        if len(contents) > 20 * 1024 * 1024:
             raise HTTPException(400, "File too large (max 20MB)")
         
         logger.info(f"Processing image: {file.filename} ({len(contents)/1024:.1f}KB)")
@@ -938,7 +839,6 @@ async def process_image_api(
         if not AI_AVAILABLE:
             raise HTTPException(503, "AI service temporarily unavailable")
         
-        # Process image
         loop = asyncio.get_event_loop()
         process_func = functools.partial(
             process_image_sync,
@@ -953,7 +853,6 @@ async def process_image_api(
         
         processing_time = time.time() - start_time
         
-        # Prepare response headers
         filename = Path(file.filename).stem
         headers = {
             "X-Processing-Time": f"{processing_time:.3f}",
@@ -989,18 +888,15 @@ async def process_video_api(
     frame_skip: int = Query(1, ge=0, le=5, description="Skip frames (0=none)"),
     format: str = Query("mp4", description="Output format (mp4, avi, webm, mov)")
 ):
-    """Remove background from a video"""
     start_time = time.time()
     
     try:
-        # Validate
         validate_video_file(file)
         
-        # Read file
         contents = await file.read()
         if len(contents) == 0:
             raise HTTPException(400, "Empty file")
-        if len(contents) > 100 * 1024 * 1024:  # 100MB limit
+        if len(contents) > 100 * 1024 * 1024:
             raise HTTPException(400, "File too large (max 100MB)")
         
         logger.info(f"Processing video: {file.filename} ({len(contents)/1024/1024:.1f}MB)")
@@ -1008,7 +904,6 @@ async def process_video_api(
         if not AI_AVAILABLE:
             raise HTTPException(503, "AI service temporarily unavailable")
         
-        # Save uploaded file
         input_suffix = Path(file.filename).suffix or '.mp4'
         input_path = tempfile.mktemp(suffix=input_suffix, dir="/tmp/uploads")
         
@@ -1017,16 +912,14 @@ async def process_video_api(
         
         output_path = None
         try:
-            # Create output file
             output_ext = format.lower()
             if transparent and output_ext in ['mp4', 'mov']:
-                output_ext = 'webm'  # WebM supports transparency better
+                output_ext = 'webm'
                 format = 'webm'
                 logger.info("Switched to WebM format for transparency support")
             
             output_path = tempfile.mktemp(suffix=f'.{output_ext}', dir="/tmp/outputs")
             
-            # Process video
             loop = asyncio.get_event_loop()
             process_func = functools.partial(
                 process_video_reliable,
@@ -1044,10 +937,8 @@ async def process_video_api(
             
             processing_time = time.time() - start_time
             
-            # Schedule cleanup
             background_tasks.add_task(cleanup_temp_files, input_path)
             
-            # Determine media type
             media_types = {
                 "mp4": "video/mp4",
                 "avi": "video/x-msvideo",
@@ -1056,7 +947,6 @@ async def process_video_api(
             }
             media_type = media_types.get(result["format"].lower(), "video/mp4")
             
-            # Prepare headers
             filename = Path(file.filename).stem
             headers = {
                 "X-Processing-Time": f"{processing_time:.3f}",
@@ -1093,11 +983,8 @@ async def process_video_api(
         logger.error(f"Video API error: {str(e)}")
         raise HTTPException(500, f"Video processing failed: {str(e)[:100]}")
 
-# ========== ERROR HANDLERS ==========
-
 @app.exception_handler(HTTPException)
 async def http_exception_handler(request, exc):
-    """Custom HTTP exception handler"""
     return JSONResponse(
         status_code=exc.status_code,
         content={
@@ -1109,7 +996,6 @@ async def http_exception_handler(request, exc):
 
 @app.exception_handler(Exception)
 async def general_exception_handler(request, exc):
-    """General exception handler"""
     logger.error(f"Unhandled exception: {str(exc)}")
     return JSONResponse(
         status_code=500,
@@ -1121,11 +1007,8 @@ async def general_exception_handler(request, exc):
         }
     )
 
-# ========== STARTUP/SHUTDOWN ==========
-
 @app.on_event("startup")
 async def startup_event():
-    """Startup tasks"""
     logger.info("=" * 50)
     logger.info("🚀 Background Remover Pro v8.2.0 Starting...")
     logger.info(f"📊 AI Available: {AI_AVAILABLE}")
@@ -1135,11 +1018,9 @@ async def startup_event():
 
 @app.on_event("shutdown")
 async def shutdown_event():
-    """Cleanup on shutdown"""
     logger.info("Shutting down...")
     cpu_executor.shutdown(wait=True)
     
-    # Cleanup temp directories
     for dir_path in ["/tmp/uploads", "/tmp/outputs"]:
         try:
             if os.path.exists(dir_path):
@@ -1147,8 +1028,6 @@ async def shutdown_event():
         except:
             pass
 
-# ========== MAIN ==========
-
 if __name__ == "__main__":
     import uvicorn