Major redesign: Lumina Synth design system, fix upscaling, add progress animations, disable SSR
d928682 | """Core upscaling engine for 4K Upscaler Pro.""" | |
| import os | |
| import tempfile | |
| import cv2 | |
| import numpy as np | |
| from PIL import Image | |
| try: | |
| from basicsr.archs.rrdbnet_arch import RRDBNet | |
| from realesrgan import RealESRGANer | |
| REALESRGAN_AVAILABLE = True | |
| except ImportError: | |
| REALESRGAN_AVAILABLE = False | |
| # ── Resolution targets ────────────────────────────────────────────── | |
| RESOLUTIONS = { | |
| "4K (3840×2160)": (3840, 2160), | |
| "2K (2560×1440)": (2560, 1440), | |
| "1080p (1920×1080)": (1920, 1080), | |
| "2x Original": None, | |
| } | |
| def get_dimensions(orig_w, orig_h, target_str): | |
| """Calculate target dimensions based on original size and target resolution.""" | |
| dims = RESOLUTIONS.get(target_str) | |
| if dims is None: | |
| return orig_w * 2, orig_h * 2 | |
| max_w, max_h = dims | |
| scale = min(max_w / orig_w, max_h / orig_h) | |
| return int(orig_w * scale), int(orig_h * scale) | |
| # ── Model loader ──────────────────────────────────────────────────── | |
| _cached_upscaler = None | |
| def load_upscaler(scale=4): | |
| """Load Real-ESRGAN model (cached after first call).""" | |
| global _cached_upscaler | |
| if _cached_upscaler is not None: | |
| return _cached_upscaler | |
| if not REALESRGAN_AVAILABLE: | |
| return None | |
| model = RRDBNet( | |
| num_in_ch=3, num_out_ch=3, num_feat=64, | |
| num_block=23, num_grow_ch=32, scale=scale, | |
| ) | |
| upscaler = RealESRGANer( | |
| scale=scale, | |
| model_path=( | |
| f"https://github.com/xinntao/Real-ESRGAN/releases/" | |
| f"download/v0.1.0/RealESRGAN_x{scale}plus.pth" | |
| ), | |
| model=model, | |
| tile=512, | |
| tile_pad=10, | |
| pre_pad=0, | |
| half=False, | |
| ) | |
| _cached_upscaler = upscaler | |
| return upscaler | |
| def opencv_upscale(img_np, target_w, target_h): | |
| """High-quality Lanczos resize.""" | |
| return cv2.resize(img_np, (target_w, target_h), interpolation=cv2.INTER_LANCZOS4) | |
| # ── Image upscaling ──────────────────────────────────────────────── | |
| def upscale_image(image_path, method, target_str, progress=None): | |
| """Upscale a single image. Returns (output_path, info_text, download_path).""" | |
| if image_path is None: | |
| return None, "❌ No image provided.", None | |
| if progress: | |
| progress(0.1, desc="📂 Loading image...") | |
| image = Image.open(image_path).convert("RGB") | |
| orig_w, orig_h = image.size | |
| img_np = np.array(image) | |
| target_w, target_h = get_dimensions(orig_w, orig_h, target_str) | |
| if progress: | |
| progress(0.3, desc=f"🔍 Scaling {orig_w}×{orig_h} → {target_w}×{target_h}...") | |
| used_method = method | |
| if method == "Real-ESRGAN (AI)" and REALESRGAN_AVAILABLE: | |
| try: | |
| upscaler = load_upscaler(scale=4) | |
| if progress: | |
| progress(0.5, desc="🤖 Running Real-ESRGAN AI model...") | |
| if upscaler: | |
| output, _ = upscaler.enhance(img_np, outscale=target_w / orig_w) | |
| output = cv2.resize( | |
| output, (target_w, target_h), interpolation=cv2.INTER_LANCZOS4 | |
| ) | |
| else: | |
| output = opencv_upscale(img_np, target_w, target_h) | |
| used_method = "Lanczos (fallback)" | |
| except Exception as e: | |
| output = opencv_upscale(img_np, target_w, target_h) | |
| used_method = f"Lanczos (fallback: {str(e)[:50]})" | |
| else: | |
| if progress: | |
| progress(0.5, desc="⚡ Applying Lanczos4 interpolation...") | |
| output = opencv_upscale(img_np, target_w, target_h) | |
| used_method = "Lanczos (Fast)" | |
| if progress: | |
| progress(0.8, desc="💾 Saving result...") | |
| result_img = Image.fromarray(output) | |
| out_path = os.path.join(tempfile.gettempdir(), "upscaled_output.png") | |
| result_img.save(out_path, format="PNG", optimize=True) | |
| info = ( | |
| f"✅ Upscaled: {orig_w}×{orig_h} → {target_w}×{target_h}\n" | |
| f"📐 Method: {used_method}\n" | |
| f"📄 Format: PNG | Size: {os.path.getsize(out_path) / 1024 / 1024:.1f} MB" | |
| ) | |
| if progress: | |
| progress(1.0, desc="✨ Done!") | |
| return out_path, info, out_path | |
| # ── Video upscaling ───────────────────────────────────────────────── | |
| def upscale_video(video_path, method, target_str, progress=None): | |
| """Upscale a video file. Returns (output_path, info_text).""" | |
| if not video_path: | |
| return None, "❌ No video provided." | |
| if progress: | |
| progress(0.05, desc="📂 Opening video...") | |
| cap = cv2.VideoCapture(video_path) | |
| if not cap.isOpened(): | |
| return None, "❌ Could not open video." | |
| try: | |
| orig_w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) | |
| orig_h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) | |
| fps = cap.get(cv2.CAP_PROP_FPS) | |
| total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) | |
| target_w, target_h = get_dimensions(orig_w, orig_h, target_str) | |
| out_path = os.path.join(tempfile.gettempdir(), "upscaled_video.mp4") | |
| fourcc = cv2.VideoWriter_fourcc(*"mp4v") | |
| writer = cv2.VideoWriter(out_path, fourcc, fps, (target_w, target_h)) | |
| if not writer.isOpened(): | |
| return None, "❌ Could not create output video writer." | |
| try: | |
| upscaler = None | |
| if method == "Real-ESRGAN (AI)" and REALESRGAN_AVAILABLE: | |
| try: | |
| upscaler = load_upscaler(scale=4) | |
| except Exception: | |
| upscaler = None | |
| frame_idx = 0 | |
| while True: | |
| ret, frame = cap.read() | |
| if not ret: | |
| break | |
| frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) | |
| if upscaler: | |
| try: | |
| out_frame, _ = upscaler.enhance( | |
| frame_rgb, outscale=target_w / orig_w | |
| ) | |
| out_frame = cv2.resize( | |
| out_frame, (target_w, target_h), | |
| interpolation=cv2.INTER_LANCZOS4, | |
| ) | |
| except Exception: | |
| out_frame = opencv_upscale(frame_rgb, target_w, target_h) | |
| else: | |
| out_frame = opencv_upscale(frame_rgb, target_w, target_h) | |
| writer.write(cv2.cvtColor(out_frame, cv2.COLOR_RGB2BGR)) | |
| frame_idx += 1 | |
| if total_frames > 0 and progress: | |
| pct = 0.1 + 0.85 * (frame_idx / total_frames) | |
| progress( | |
| pct, | |
| desc=f"🎬 Processing frame {frame_idx}/{total_frames}", | |
| ) | |
| info = ( | |
| f"✅ Video upscaled: {orig_w}×{orig_h} → {target_w}×{target_h}\n" | |
| f"🎞 {total_frames} frames @ {fps:.1f} fps\n" | |
| f"📄 Size: {os.path.getsize(out_path) / 1024 / 1024:.1f} MB" | |
| ) | |
| finally: | |
| writer.release() | |
| finally: | |
| cap.release() | |
| if progress: | |
| progress(1.0, desc="✨ Done!") | |
| return out_path, info | |