Update app.py

app.py

@@ -6,7 +6,6 @@ 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
 
@@ -24,14 +23,14 @@ def setup_environment():
     model_path = "weights/CodeFormer/codeformer.pth"
     download_file(model_url, model_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 facexlib model (for face detection)
+    facelib_url = "https://github.com/xinntao/facexlib/releases/download/v0.1.0/detection_Resnet50_Final.pth"
+    facelib_path = "weights/facelib/detection_Resnet50_Final.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"
+    realesrgan_url = "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth"
+    realesrgan_path = "weights/realesrgan/RealESRGAN_x4plus.pth"
     download_file(realesrgan_url, realesrgan_path)
 
 # Load CodeFormer model
@@ -43,6 +42,23 @@ def load_codeformer():
     model = model.to('cpu')  # Force CPU
     return model
 
+# 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
+
 # Inference function
 def enhance_image(input_image, fidelity_weight=0.5, background_enhance=True, face_upsample=False):
     # Convert PIL image to OpenCV format
@@ -84,7 +100,7 @@ def enhance_image(input_image, fidelity_weight=0.5, background_enhance=True, fac
         from realesrgan import RealESRGANer
         upsampler = RealESRGANer(
             scale=4,
-            model_path="weights/RealESRGAN_x4plus.pth",
+            model_path="weights/realesrgan/RealESRGAN_x4plus.pth",
             device='cpu'
         )
         restored_img, _ = upsampler.enhance(restored_img, outscale=4)