Spaces:
Sleeping
Sleeping
| import os | |
| import sys | |
| import glob | |
| from PIL import Image | |
| import torch | |
| from basicsr.archs.rrdbnet_arch import RRDBNet | |
| from realesrgan import RealESRGANer | |
| def setup_upscaler(): | |
| """Configure l'upscaler Real-ESRGAN.""" | |
| model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4) | |
| model_path = 'RealESRGAN_x4plus.pth' | |
| # Téléchargement automatique du modèle si absent (géré par RealESRGANer ou wget) | |
| if not os.path.exists(model_path): | |
| import subprocess | |
| print("📥 Téléchargement du modèle Real-ESRGAN...") | |
| subprocess.run(["wget", "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth"]) | |
| upsampler = RealESRGANer( | |
| scale=4, | |
| model_path=model_path, | |
| model=model, | |
| tile=0, | |
| tile_pad=10, | |
| pre_pad=0, | |
| half=True # Utilise FP16 pour les GPUs modernes (A10, A100, etc.) | |
| ) | |
| return upsampler | |
| def main(): | |
| input_dir = "input_images" | |
| output_dir = "output_images" | |
| os.makedirs(output_dir, exist_ok=True) | |
| print("🚀 Initialisation de l'upscaler AI...") | |
| upsampler = setup_upscaler() | |
| images = glob.glob(os.path.join(input_dir, "*")) | |
| print(f"📂 {len(images)} images à traiter.") | |
| for img_path in images: | |
| name = os.path.basename(img_path) | |
| print(f"🪄 Processing {name}...") | |
| try: | |
| import cv2 | |
| img = cv2.imread(img_path, cv2.IMREAD_UNCHANGED) | |
| output, _ = upsampler.enhance(img, outscale=4) | |
| cv2.imwrite(os.path.join(output_dir, name), output) | |
| print(f"✅ Terminé : {name}") | |
| except Exception as e: | |
| print(f"❌ Erreur sur {name} : {e}") | |
| if __name__ == "__main__": | |
| # Installation des dépendances si nécessaire (sur l'instance Lambda) | |
| import subprocess | |
| print("📦 Installation des dépendances AI...") | |
| subprocess.run(["pip", "install", "basicsr", "realesrgan", "opencv-python"]) | |
| main() | |