id-maker/core/restoration.py

import os
import requests
import io
from PIL import Image
from pathlib import Path
from layout_engine import load_settings

# The URL should be configured in settings.json or as an environment variable
def get_api_url():
    settings = load_settings()
    return os.environ.get("CODEFORMER_API_URL", settings.get("api", {}).get("codeformer_url", "http://localhost:7860/api/restore"))

def restore_image(img, fidelity=0.5, upscale=1, return_pil=True):
    """
    Calls external CodeFormer API to restore the image.
    img: PIL Image or path to image
    """
    url = get_api_url()
    
    # 1. Prepare the image buffer
    img_buffer = io.BytesIO()
    if isinstance(img, (str, Path)):
        with open(img, 'rb') as f:
            img_data = f.read()
            # Basic mime type guess
            ext = str(img).lower()
            mime = 'image/png' if ext.endswith('.png') else 'image/jpeg'
            fname = 'input.png' if ext.endswith('.png') else 'input.jpg'
    elif hasattr(img, 'save'): # PIL Image
        # Send as PNG to avoid compression artifacts before restoration
        img.save(img_buffer, format='PNG')
        img_data = img_buffer.getvalue()
        mime = 'image/png'
        fname = 'input.png'
    else:
        raise ValueError("Unsupported image type for restoration")

    # 2. Call the API
    try:
        print(f"Calling CodeFormer API at {url}...")
        files = {'image': (fname, img_data, mime)}
        data = {
            'fidelity': str(fidelity), 
            'upscale': str(upscale),
            'background_enhance': 'false',
            'face_upsample': 'false'
        }
        
        # The API returns JSON with an image_url
        response = requests.post(url, files=files, data=data, timeout=120)
        response.raise_for_status()
        result_json = response.json()

        if result_json.get('status') == 'success' and result_json.get('results'):
            restored_url = result_json['results'][0]['image_url']
            print(f"Restoration successful. Downloading result from: {restored_url}")
            
            # Download the actual image bytes
            img_res = requests.get(restored_url, timeout=60)
            img_res.raise_for_status()
            restored_pil = Image.open(io.BytesIO(img_res.content))
        else:
            raise ValueError(f"API Error: {result_json.get('message', 'Unknown error')}")
        
        if return_pil:
            return restored_pil
        else:
            import numpy as np
            import cv2
            return cv2.cvtColor(np.array(restored_pil), cv2.COLOR_RGB2BGR)

    except Exception as e:
        print(f"API Restoration Error: {e}")
        # Fallback to original image on failure
        if hasattr(img, 'convert'):
            return img
        else:
            return Image.open(io.BytesIO(img_data))

def init_restoration_model(device_str=None):
    """No-op for API-based restoration, but kept for compatibility."""
    print("Using Remote CodeFormer API (No local model loaded)")
    return None, "remote"

def batch_restore(input_folder, output_folder, fidelity=0.5):
    """CLI compatibility for remote restoration."""
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)
        
    files = [f for f in os.listdir(input_folder) if f.lower().endswith(('.jpg', '.jpeg', '.png'))]
    
    for filename in files:
        print(f"Restoring {filename} via API...")
        img_path = os.path.join(input_folder, filename)
        restored_img = restore_image(img_path, fidelity=fidelity, return_pil=True)
        restored_img.save(os.path.join(output_folder, filename), quality=95)