Create gfpgan_cpu.py

gfpgan_cpu.py

@@ -0,0 +1,185 @@
+import os
+import argparse
+import cv2
+import glob
+import numpy as np
+import torch
+from tqdm import tqdm
+from pathlib import Path
+from basicsr.utils import imwrite
+
+# This is a simple implementation of GFPGAN for CPU usage on Hugging Face
+def download_model():
+    """Download the GFPGAN model if not already present"""
+    import urllib.request
+    import os
+    
+    os.makedirs('experiments/pretrained_models', exist_ok=True)
+    model_path = 'experiments/pretrained_models/GFPGANv1.3.pth'
+    
+    if not os.path.exists(model_path):
+        print("Downloading GFPGANv1.3 model...")
+        url = 'https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.3.pth'
+        urllib.request.urlretrieve(url, model_path)
+        print(f"Model downloaded to {model_path}")
+    
+    return model_path
+
+def setup_gfpgan():
+    """Set up GFPGAN with the required dependencies"""
+    # Install required packages if not already installed
+    try:
+        import basicsr
+    except ImportError:
+        print("Installing basicsr...")
+        os.system('pip install basicsr')
+    
+    try:
+        import facexlib
+    except ImportError:
+        print("Installing facexlib...")
+        os.system('pip install facexlib')
+    
+    try:
+        import gfpgan
+    except ImportError:
+        print("Installing GFPGAN...")
+        os.system('pip install gfpgan')
+    
+    from gfpgan import GFPGANer
+    
+    # Download the model
+    model_path = download_model()
+    
+    # Initialize GFPGAN for CPU usage
+    device = torch.device('cpu')
+    
+    # Set up the restorer - note we're using CPU mode with half=False
+    restorer = GFPGANer(
+        model_path=model_path,
+        upscale=2,
+        arch='clean',
+        channel_multiplier=2,
+        bg_upsampler=None,  # No background upsampler for CPU
+        device=device
+    )
+    
+    return restorer
+
+def process_image(restorer, img_path, output_dir='results'):
+    """Process a single image with GFPGAN"""
+    os.makedirs(output_dir, exist_ok=True)
+    os.makedirs(os.path.join(output_dir, 'restored_faces'), exist_ok=True)
+    os.makedirs(os.path.join(output_dir, 'restored_imgs'), exist_ok=True)
+    os.makedirs(os.path.join(output_dir, 'cmp'), exist_ok=True)
+    
+    # Read image
+    img_name = os.path.basename(img_path)
+    print(f'Processing {img_name} ...')
+    
+    basename, ext = os.path.splitext(img_name)
+    input_img = cv2.imread(img_path, cv2.IMREAD_COLOR)
+    
+    if input_img is None:
+        print(f"Warning: Cannot read image {img_path}")
+        return
+    
+    # Restore faces and background
+    cropped_faces, restored_faces, restored_img = restorer.enhance(
+        input_img, has_aligned=False, only_center_face=False, paste_back=True)
+    
+    # Save faces
+    for idx, (cropped_face, restored_face) in enumerate(zip(cropped_faces, restored_faces)):
+        # Save restored face
+        save_face_name = f'{basename}_{idx:02d}.png'
+        save_restore_path = os.path.join(output_dir, 'restored_faces', save_face_name)
+        imwrite(restored_face, save_restore_path)
+        
+        # Save comparison image
+        cmp_img = np.concatenate((cropped_face, restored_face), axis=1)
+        imwrite(cmp_img, os.path.join(output_dir, 'cmp', f'{basename}_{idx:02d}.png'))
+    
+    # Save restored image
+    if restored_img is not None:
+        extension = ext[1:] if ext else 'png'
+        save_restore_path = os.path.join(output_dir, 'restored_imgs', f'{basename}.{extension}')
+        imwrite(restored_img, save_restore_path)
+    
+    return os.path.join(output_dir, 'restored_imgs', f'{basename}.{extension}')
+
+def main():
+    """Main function to run GFPGAN on CPU"""
+    parser = argparse.ArgumentParser(description='GFPGAN for CPU')
+    parser.add_argument('--input', type=str, default='inputs', help='Input image or folder')
+    parser.add_argument('--output', type=str, default='results', help='Output folder')
+    args = parser.parse_args()
+    
+    # Set up GFPGAN
+    restorer = setup_gfpgan()
+    
+    # Process images
+    input_path = args.input
+    output_dir = args.output
+    
+    if os.path.isfile(input_path):
+        # Single image
+        process_image(restorer, input_path, output_dir)
+    else:
+        # Directory of images
+        os.makedirs(input_path, exist_ok=True)
+        img_list = sorted(glob.glob(os.path.join(input_path, '*.[jp][pn]g')))
+        for img_path in tqdm(img_list):
+            process_image(restorer, img_path, output_dir)
+    
+    print(f'Results are saved in {output_dir}')
+
+# For Hugging Face Spaces (Gradio interface)
+def create_gradio_app():
+    import gradio as gr
+    
+    restorer = setup_gfpgan()
+    
+    def process_image_gradio(image):
+        # Save input image temporarily
+        temp_input = 'temp_input.jpg'
+        cv2.imwrite(temp_input, image[:, :, ::-1])  # Convert RGB to BGR for OpenCV
+        
+        # Process the image
+        output_path = process_image(restorer, temp_input, 'results')
+        
+        # Read the output image
+        restored_img = cv2.imread(output_path)
+        
+        # Convert back to RGB for Gradio
+        if restored_img is not None:
+            restored_img = restored_img[:, :, ::-1]
+            return restored_img
+        else:
+            return image  # Return original if processing failed
+    
+    # Create Gradio interface
+    app = gr.Interface(
+        fn=process_image_gradio,
+        inputs=gr.Image(),
+        outputs=gr.Image(),
+        title="GFPGAN - Face Restoration",
+        description="Upload an image to improve facial details with GFPGAN running on CPU"
+    )
+    
+    return app
+
+if __name__ == '__main__':
+    import sys
+    
+    # Check if running in a Hugging Face Space
+    if os.getenv('SPACE_ID'):
+        try:
+            import gradio as gr
+        except ImportError:
+            os.system('pip install gradio')
+            import gradio as gr
+        
+        app = create_gradio_app()
+        app.launch()
+    else:
+        main()