Update app.py

app.py

@@ -6,6 +6,9 @@ import cv2
 import numpy as np
 from PIL import Image
 import urllib.request
+from basicsr.utils import img2tensor, tensor2img
+from facexlib.utils.face_restoration_helper import FaceRestoreHelper
+from codeformer_arch import CodeFormer
 
 # Function to download a file from a URL
 def download_file(url, dest):
@@ -14,65 +17,56 @@ def download_file(url, dest):
         urllib.request.urlretrieve(url, dest)
         print(f"Downloaded {dest}")
 
-# Download pretrained model and necessary files
+# Download pretrained models
 def setup_environment():
     # Download CodeFormer pretrained model
     model_url = "https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/codeformer.pth"
-    model_path = "weights/codeformer.pth"
+    model_path = "weights/CodeFormer/codeformer.pth"
     download_file(model_url, model_path)
 
-    # Download facexlib detection models
-    retinaface_url = "https://github.com/xinntao/facexlib/releases/download/v0.1.0/detection_Resnet50_Final.pth"
-    retinaface_path = "weights/detection_Resnet50_Final.pth"
-    download_file(retinaface_url, retinaface_path)
+    # Download facelib model (for face detection)
+    facelib_url = "https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/facelib.pth"
+    facelib_path = "weights/facelib.pth"
+    download_file(facelib_url, facelib_path)
+
+    # Download Real-ESRGAN model for background upsampling (optional)
+    realesrgan_url = "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/RealESRGAN_x4plus.pth"
+    realesrgan_path = "weights/RealESRGAN_x4plus.pth"
+    download_file(realesrgan_url, realesrgan_path)
 
 # Load CodeFormer model
 def load_codeformer():
     setup_environment()
-    from codeformer_arch import CodeFormer
-    model_path = "weights/codeformer.pth"
-    net = CodeFormer(dim_embd=512, codebook_size=1024, n_head=8, n_layer=9, connect_list=['32', '64', '128', '256']).to('cpu')
-    checkpoint = torch.load(model_path, map_location='cpu')
-    net.load_state_dict(checkpoint)
-    net.eval()
-    return net
-
-# Image processing utilities (mimicking basicsr.utils)
-def img2tensor(img, bgr2rgb=True, float32=True):
-    if bgr2rgb:
-        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
-    img = torch.from_numpy(img.transpose(2, 0, 1)).float()
-    if float32:
-        img = img / 255.0
-    return img
-
-def tensor2img(tensor, rgb2bgr=True, min_max=(-1, 1)):
-    tensor = tensor.squeeze().float().cpu().clamp_(*min_max)
-    tensor = (tensor - min_max[0]) / (min_max[1] - min_max[0]) * 255.0
-    img = tensor.numpy().transpose(1, 2, 0).astype(np.uint8)
-    if rgb2bgr:
-        img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
-    return img
+    model = CodeFormer(dim_embd=512, codebook_size=1024, n_head=8, n_layer=9, connect_list=['32', '64', '128', '256'])
+    model.load_state_dict(torch.load("weights/CodeFormer/codeformer.pth", map_location='cpu'))
+    model.eval()
+    model = model.to('cpu')  # Force CPU
+    return model
 
 # Inference function
-def enhance_image(image, fidelity_weight=0.5):
-    from facexlib.utils.face_restoration_helper import FaceRestoreHelper
-    
-    # Load model
-    net = load_codeformer()
-    
+def enhance_image(input_image, fidelity_weight=0.5, background_enhance=True, face_upsample=False):
     # Convert PIL image to OpenCV format
-    img = np.array(image)
+    img = np.array(input_image)
     img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
-    
+
     # Initialize face helper
-    face_helper = FaceRestoreHelper(upscale_factor=1, face_size=512, crop_ratio=(1, 1), det_model='retinaface_resnet50', save_ext='png', device='cpu')
+    face_helper = FaceRestoreHelper(
+        upscale_factor=1 if not face_upsample else 2,
+        face_size=512,
+        crop_ratio=(1, 1),
+        det_model='retinaface_resnet50',
+        save_ext='png',
+        device='cpu'
+    )
     face_helper.clean_all()
     face_helper.read_image(img)
     face_helper.get_face_landmarks_5()
     face_helper.align_warp_face()
-    
-    # Enhance face with CodeFormer
+
+    # Load CodeFormer model
+    net = load_codeformer()
+
+    # Enhance face
     for cropped_face in face_helper.cropped_faces:
         cropped_face_t = img2tensor(cropped_face, bgr2rgb=True, float32=True)
         with torch.no_grad():
@@ -80,11 +74,21 @@ def enhance_image(image, fidelity_weight=0.5):
         restored_face = tensor2img(output, rgb2bgr=True, min_max=(-1, 1))
         restored_face = restored_face.astype('uint8')
         face_helper.add_restored_face(restored_face)
-    
-    # Get final restored image
+
+    # Get restored image
     face_helper.get_inverse_affine(None)
     restored_img = face_helper.paste_faces_to_input_image()
-    
+
+    # Background enhancement with Real-ESRGAN (optional)
+    if background_enhance:
+        from realesrgan import RealESRGANer
+        upsampler = RealESRGANer(
+            scale=4,
+            model_path="weights/RealESRGAN_x4plus.pth",
+            device='cpu'
+        )
+        restored_img, _ = upsampler.enhance(restored_img, outscale=4)
+
     # Convert back to PIL for Gradio
     restored_img = cv2.cvtColor(restored_img, cv2.COLOR_BGR2RGB)
     return Image.fromarray(restored_img)
@@ -93,17 +97,19 @@ def enhance_image(image, fidelity_weight=0.5):
 with gr.Blocks() as demo:
     gr.Markdown("# CodeFormer Face Restoration (CPU)")
     gr.Markdown("Upload an image to enhance faces using CodeFormer. Runs on CPU in Hugging Face Spaces.")
-    
+
     with gr.Row():
         input_image = gr.Image(type="pil", label="Input Image")
         output_image = gr.Image(type="pil", label="Enhanced Image")
-    
-    fidelity_slider = gr.Slider(0, 1, value=0.5, step=0.1, label="Fidelity Weight (0 = more restoration, 1 = more original)")
+
+    fidelity_slider = gr.Slider(0, 1, value=0.5, step=0.01, label="Fidelity Weight (0 = more restoration, 1 = more original)")
+    background_enhance = gr.Checkbox(label="Enhance Background (Real-ESRGAN)", value=True)
+    face_upsample = gr.Checkbox(label="Upsample Restored Faces", value=False)
     submit_btn = gr.Button("Enhance")
-    
+
     submit_btn.click(
         fn=enhance_image,
-        inputs=[input_image, fidelity_slider],
+        inputs=[input_image, fidelity_slider, background_enhance, face_upsample],
         outputs=output_image
     )