Add PicPro Gradio backend with modular services and weights

app.py

@@ -0,0 +1,99 @@
+import gradio as gr
+import os
+import sys
+
+# Add current directory to path so we can import services
+sys.path.append(os.path.dirname(os.path.abspath(__file__)))
+
+from services.enhancer import ImageEnhancer
+from services.background_remover import BackgroundRemover
+from utils.download_weights import main as download_weights
+
+# --- Initialization ---
+print("Initializing PicPro AI Backend...")
+
+# Ensure weights are present
+try:
+    download_weights()
+except Exception as e:
+    print(f"Weight download warning: {e}")
+
+# Global Instances
+enhancer = ImageEnhancer()
+bg_remover = BackgroundRemover()
+
+# Pre-load models (optional, can be lazy)
+# enhancer.load_models()
+# bg_remover.load_model()
+
+# --- Wrapper Functions for Gradio ---
+
+def enhance_image(image):
+    """
+    Args:
+        image: numpy array (RGB) from Gradio
+    Returns:
+        numpy array (RGB)
+    """
+    if image is None:
+        return None
+    
+    try:
+        result = enhancer.process(image)
+        return result
+    except Exception as e:
+        raise gr.Error(f"Enhancement failed: {str(e)}")
+
+def remove_background(image):
+    """
+    Args:
+        image: PIL Image from Gradio
+    Returns:
+        PIL Image
+    """
+    if image is None:
+        return None
+    
+    try:
+        result = bg_remover.process(image)
+        return result
+    except Exception as e:
+        raise gr.Error(f"Background removal failed: {str(e)}")
+
+# --- Gradio Interface ---
+
+with gr.Blocks(title="PicPro AI Backend") as app:
+    gr.Markdown("# PicPro AI Backend Services")
+    gr.Markdown("GPU-accelerated Image Enhancement and Background Removal.")
+    
+    with gr.Tab("Enhance Image"):
+        gr.Markdown("### Real-ESRGAN + GFPGAN")
+        with gr.Row():
+            enhance_input = gr.Image(label="Input Image", type="numpy")
+            enhance_output = gr.Image(label="Enhanced Image", type="numpy")
+        
+        enhance_btn = gr.Button("Enhance", variant="primary")
+        enhance_btn.click(
+            fn=enhance_image,
+            inputs=enhance_input,
+            outputs=enhance_output,
+            api_name="enhance_image"
+        )
+        
+    with gr.Tab("Remove Background"):
+        gr.Markdown("### MODNet / U-2-Net (rembg)")
+        with gr.Row():
+            bg_input = gr.Image(label="Input Image", type="pil")
+            bg_output = gr.Image(label="No Background", type="pil", image_mode="RGBA")
+            
+        bg_btn = gr.Button("Remove Background", variant="primary")
+        bg_btn.click(
+            fn=remove_background,
+            inputs=bg_input,
+            outputs=bg_output,
+            api_name="remove_background"
+        )
+
+# Launch
+if __name__ == "__main__":
+    app.launch(server_name="0.0.0.0", server_port=7860)

requirements.txt

@@ -0,0 +1,12 @@
+gradio>=4.0.0
+torch>=2.0.0
+torchvision>=0.15.0
+numpy
+opencv-python-headless
+Pillow
+gfpgan>=1.3.8
+realesrgan>=0.3.0
+basicsr>=1.4.2
+rembg[gpu]
+requests
+aiofiles

services/background_remover.py

@@ -0,0 +1,35 @@
+from PIL import Image
+try:
+    from rembg import remove, new_session
+except ImportError:
+    print("Warning: rembg not found.")
+    remove = None
+    new_session = None
+
+class BackgroundRemover:
+    def __init__(self):
+        self.session = None
+
+    def load_model(self):
+        if self.session is not None:
+            return
+
+        if new_session is None:
+            raise ImportError("rembg not installed")
+            
+        # Initialize session
+        # 'u2net' is robust. 
+        self.session = new_session("u2net")
+        print("Background Remover loaded")
+
+    def process(self, input_image: Image.Image) -> Image.Image:
+        """
+        Process a PIL Image.
+        Returns: PIL Image with background removed.
+        """
+        if self.session is None:
+            self.load_model()
+
+        # Process
+        output_image = remove(input_image, session=self.session)
+        return output_image

services/enhancer.py

@@ -0,0 +1,88 @@
+import cv2
+import numpy as np
+import torch
+import os
+try:
+    from basicsr.archs.rrdbnet_arch import RRDBNet
+    from realesrgan import RealESRGANer
+    from gfpgan import GFPGANer
+except ImportError:
+    print("Warning: AI modules not found. Ensure requirements are installed.")
+    RRDBNet = None
+    RealESRGANer = None
+    GFPGANer = None
+
+class ImageEnhancer:
+    def __init__(self):
+        self.upsampler = None
+        self.face_enhancer = None
+        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+    def load_models(self):
+        if self.upsampler is not None:
+            return
+
+        if RRDBNet is None:
+            raise ImportError("AI modules not installed")
+
+        # Paths
+        # In HF Spaces, we can download weights if missing or rely on cache
+        weights_dir = 'weights'
+        os.makedirs(weights_dir, exist_ok=True)
+        
+        realesrgan_path = os.path.join(weights_dir, 'RealESRGAN_x4plus.pth')
+        gfpgan_path = os.path.join(weights_dir, 'GFPGANv1.4.pth')
+
+        # Simple check/download if missing (handled by utils/download_weights.py usually)
+        # For robustness in Gradio Space, we assume they exist or we rely on the helper
+        if not os.path.exists(realesrgan_path) or not os.path.exists(gfpgan_path):
+             print("Weights not found locally. Ensure they are downloaded.")
+             # Fallback or error handled by caller
+
+        # Load Real-ESRGAN
+        model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)
+        
+        self.upsampler = RealESRGANer(
+            scale=4,
+            model_path=realesrgan_path,
+            model=model,
+            tile=0,
+            tile_pad=10,
+            pre_pad=0,
+            half=True if self.device.type == 'cuda' else False,
+            device=self.device,
+        )
+
+        # Load GFPGAN
+        self.face_enhancer = GFPGANer(
+            model_path=gfpgan_path,
+            upscale=4,
+            arch='clean',
+            channel_multiplier=2,
+            bg_upsampler=self.upsampler
+        )
+        print(f"Enhancer loaded on {self.device}")
+
+    def process(self, img_rgb):
+        """
+        Process an image (numpy array, RGB).
+        Returns: Enhanced image (numpy array, RGB).
+        """
+        if self.face_enhancer is None:
+            self.load_models()
+
+        # Convert RGB (Gradio) to BGR (OpenCV/GFPGAN)
+        img_bgr = cv2.cvtColor(img_rgb, cv2.COLOR_RGB2BGR)
+        
+        # Inference
+        # has_aligned=False means it will detect faces
+        _, _, output_bgr = self.face_enhancer.enhance(
+            img_bgr,
+            has_aligned=False,
+            only_center_face=False,
+            paste_back=True
+        )
+        
+        # Convert BGR back to RGB for Gradio
+        output_rgb = cv2.cvtColor(output_bgr, cv2.COLOR_BGR2RGB)
+        return output_rgb

utils/download_weights.py

@@ -0,0 +1,40 @@
+import os
+import requests
+import sys
+
+def download_file(url, dest_path):
+    if os.path.exists(dest_path):
+        print(f"File already exists: {dest_path}")
+        return
+
+    print(f"Downloading {url} to {dest_path}...")
+    try:
+        response = requests.get(url, stream=True)
+        response.raise_for_status()
+        with open(dest_path, 'wb') as f:
+            for chunk in response.iter_content(chunk_size=8192):
+                f.write(chunk)
+        print("Download complete.")
+    except Exception as e:
+        print(f"Error downloading {url}: {e}")
+        if os.path.exists(dest_path):
+            os.remove(dest_path)
+
+def main():
+    weights_dir = os.path.join(os.path.dirname(__file__), '..', 'weights')
+    os.makedirs(weights_dir, exist_ok=True)
+
+    # Real-ESRGAN
+    realesrgan_url = "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth"
+    download_file(realesrgan_url, os.path.join(weights_dir, "RealESRGAN_x4plus.pth"))
+
+    # GFPGAN
+    gfpgan_url = "https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.4.pth"
+    download_file(gfpgan_url, os.path.join(weights_dir, "GFPGANv1.4.pth"))
+    
+    # rembg will download its own models to ~/.u2net by default, 
+    # but for HF Spaces we might want to cache it. 
+    # For now, rembg handles it on first run.
+
+if __name__ == "__main__":
+    main()

weights/GFPGANv1.4.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f63b495992c0874c28e6c627fce803896264687beed2b44b9957bcfd890edc78
+size 4243456

weights/RealESRGAN_x4plus.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4fa0d38905f75ac06eb49a7951b426670021be3018265fd191d2125df9d682f1
+size 67040989