from PIL import Image
import cv2
import numpy as np
 
def compare_upscaling_methods(image_path, scale=4):
    """古典手法とAI超解像の比較"""
    img = Image.open(image_path)
    w, h = img.size
    new_w, new_h = w * scale, h * scale
 
    results = {}
 
    # 1. 最近傍補間 (Nearest Neighbor)
    results["nearest"] = img.resize(
        (new_w, new_h), Image.NEAREST
    )
 
    # 2. バイリニア補間
    results["bilinear"] = img.resize(
        (new_w, new_h), Image.BILINEAR
    )
 
    # 3. バイキュービック補間
    results["bicubic"] = img.resize(
        (new_w, new_h), Image.BICUBIC
    )
 
    # 4. Lanczos補間
    results["lanczos"] = img.resize(
        (new_w, new_h), Image.LANCZOS
    )
 
    return results
 
# 品質メトリクスの計算
def calculate_psnr(original, upscaled):
    """PSNR (ピーク信号対雑音比) を計算"""
    mse = np.mean((np.array(original) - np.array(upscaled)) ** 2)
    if mse == 0:
        return float('inf')
    return 20 * np.log10(255.0 / np.sqrt(mse))
 
def calculate_ssim(original, upscaled):
    """SSIM (構造的類似性) を計算"""
    from skimage.metrics import structural_similarity
    return structural_similarity(
        np.array(original),
        np.array(upscaled),
        channel_axis=2,  # カラー画像
    )
 
def calculate_lpips(original, upscaled):
    """LPIPS (知覚的類似性) を計算 — 人間の知覚に近い指標"""
    import torch
    import lpips
 
    loss_fn = lpips.LPIPS(net='alex')
 
    # PIL → Tensor に変換 ([-1, 1] 範囲)
    def to_tensor(img):
        arr = np.array(img).astype(np.float32) / 255.0
        arr = arr * 2.0 - 1.0
        return torch.from_numpy(arr).permute(2, 0, 1).unsqueeze(0)
 
    t_orig = to_tensor(original)
    t_upsc = to_tensor(upscaled)
 
    with torch.no_grad():
        score = loss_fn(t_orig, t_upsc)
 
    return score.item()  # 低いほど類似度が高い
 
 
def comprehensive_quality_assessment(original_path, upscaled_path):
    """包括的な品質評価レポートを生成"""
    from PIL import Image
    import json
 
    original = Image.open(original_path)
    upscaled = Image.open(upscaled_path)
 
    # 解像度チェック
    orig_w, orig_h = original.size
    upsc_w, upsc_h = upscaled.size
    scale_w = upsc_w / orig_w
    scale_h = upsc_h / orig_h
 
    # メトリクス計算
    # ※PSNRとSSIMは同サイズの画像で比較するため、
    # 元画像を同じサイズにリサイズして比較
    original_resized = original.resize(
        (upsc_w, upsc_h), Image.LANCZOS
    )
 
    report = {
        "resolution": {
            "original": f"{orig_w}x{orig_h}",
            "upscaled": f"{upsc_w}x{upsc_h}",
            "scale_factor": f"{scale_w:.1f}x{scale_h:.1f}",
        },
        "metrics": {
            "psnr_db": round(
                calculate_psnr(original_resized, upscaled), 2
            ),
            "ssim": round(
                calculate_ssim(original_resized, upscaled), 4
            ),
        },
        "file_info": {
            "original_size_kb": round(
                os.path.getsize(original_path) / 1024, 1
            ),
            "upscaled_size_kb": round(
                os.path.getsize(upscaled_path) / 1024, 1
            ),
        },
    }
 
    return report

def simulate_real_world_degradation(image, scale=4):
    """
    Real-ESRGAN の第2次劣化モデルを再現
 
    実世界の劣化 = ぼかし + ダウンサンプリング + ノイズ + JPEG圧縮
    これを2回繰り返す (second-order degradation)
    """
    import cv2
    import numpy as np
 
    img = np.array(image).astype(np.float32) / 255.0
 
    # === 第1次劣化 ===
    # 1. ぼかし (等方性/異方性ガウシアン)
    kernel_size = np.random.choice([7, 9, 11, 13, 15, 17, 19, 21])
    sigma = np.random.uniform(0.2, 3.0)
    img = cv2.GaussianBlur(img, (kernel_size, kernel_size), sigma)
 
    # 2. ダウンサンプリング (バイキュービック/バイリニア/エリア)
    h, w = img.shape[:2]
    method = np.random.choice([
        cv2.INTER_CUBIC,
        cv2.INTER_LINEAR,
        cv2.INTER_AREA,
    ])
    down_scale = np.random.uniform(1.0, scale)
    img = cv2.resize(
        img, (int(w / down_scale), int(h / down_scale)),
        interpolation=method,
    )
 
    # 3. ノイズ追加 (ガウシアン/ポアソン)
    noise_type = np.random.choice(["gaussian", "poisson"])
    if noise_type == "gaussian":
        sigma_n = np.random.uniform(1, 30) / 255.0
        noise = np.random.normal(0, sigma_n, img.shape)
        img = img + noise
    else:
        vals = 2 ** np.ceil(np.log2(len(np.unique(img))))
        img = np.random.poisson(img * vals) / vals
 
    # 4. JPEG圧縮アーティファクト
    quality = np.random.randint(30, 95)
    encode_param = [int(cv2.IMWRITE_JPEG_QUALITY), quality]
    _, enc = cv2.imencode('.jpg', (img * 255).astype(np.uint8),
                          encode_param)
    img = cv2.imdecode(enc, cv2.IMREAD_COLOR).astype(np.float32) / 255.0
 
    # === 第2次劣化 (繰り返し) ===
    # 同様の処理をもう一度適用（パラメータは異なる範囲で）
    kernel_size2 = np.random.choice([7, 9, 11, 13, 15, 17, 19, 21])
    sigma2 = np.random.uniform(0.2, 1.5)
    img = cv2.GaussianBlur(img, (kernel_size2, kernel_size2), sigma2)
 
    # 最終リサイズ
    img = cv2.resize(
        img, (w // scale, h // scale),
        interpolation=cv2.INTER_LINEAR,
    )
 
    img = np.clip(img, 0, 1)
    return (img * 255).astype(np.uint8)

# pip install realesrgan
from basicsr.archs.rrdbnet_arch import RRDBNet
from realesrgan import RealESRGANer
from PIL import Image
import numpy as np
import cv2
 
def upscale_with_realesrgan(image_path, scale=4, model_name="x4plus"):
    """Real-ESRGAN による超解像"""
 
    # モデル選択
    models = {
        "x4plus": {
            "model": RRDBNet(
                num_in_ch=3, num_out_ch=3, num_feat=64,
                num_block=23, num_grow_ch=32, scale=4,
            ),
            "url": "https://github.com/xinntao/Real-ESRGAN/releases/"
                   "download/v0.1.0/RealESRGAN_x4plus.pth",
            "scale": 4,
        },
        "x4plus_anime": {
            "model": RRDBNet(
                num_in_ch=3, num_out_ch=3, num_feat=64,
                num_block=6, num_grow_ch=32, scale=4,
            ),
            "url": "https://github.com/xinntao/Real-ESRGAN/releases/"
                   "download/v0.2.2.4/RealESRGAN_x4plus_anime_6B.pth",
            "scale": 4,
        },
        "x2plus": {
            "model": RRDBNet(
                num_in_ch=3, num_out_ch=3, num_feat=64,
                num_block=23, num_grow_ch=32, scale=2,
            ),
            "url": "https://github.com/xinntao/Real-ESRGAN/releases/"
                   "download/v0.2.1/RealESRGAN_x2plus.pth",
            "scale": 2,
        },
    }
 
    config = models[model_name]
 
    upsampler = RealESRGANer(
        scale=config["scale"],
        model_path=config["url"],
        model=config["model"],
        tile=0,          # タイル処理 (0=無効, 推奨: 400)
        tile_pad=10,     # タイル間のパディング
        pre_pad=0,
        half=True,       # FP16 で高速化
    )
 
    img = cv2.imread(image_path, cv2.IMREAD_UNCHANGED)
    output, _ = upsampler.enhance(img, outscale=scale)
    cv2.imwrite("upscaled.png", output)
 
    return output
 
# アニメ画像用
upscale_with_realesrgan("anime.png", model_name="x4plus_anime")
 
# 写真用
upscale_with_realesrgan("photo.jpg", model_name="x4plus")

def tiled_upscale(image_path, tile_size=512, overlap=64, scale=4):
    """
    大画像をタイル分割して超解像
 
    メモリ制限のある環境で大画像を処理する手法:
    1. 画像をタイルに分割
    2. 各タイルを個別に超解像
    3. 重複領域をブレンドして結合
    """
    from PIL import Image
    import numpy as np
 
    img = np.array(Image.open(image_path))
    h, w, c = img.shape
    out_h, out_w = h * scale, w * scale
 
    output = np.zeros((out_h, out_w, c), dtype=np.float32)
    weight = np.zeros((out_h, out_w, c), dtype=np.float32)
 
    # タイルの位置を計算
    y_positions = list(range(0, h - tile_size + 1, tile_size - overlap))
    x_positions = list(range(0, w - tile_size + 1, tile_size - overlap))
 
    # 端が含まれない場合は追加
    if y_positions[-1] + tile_size < h:
        y_positions.append(h - tile_size)
    if x_positions[-1] + tile_size < w:
        x_positions.append(w - tile_size)
 
    for y in y_positions:
        for x in x_positions:
            # タイルを切り出し
            tile = img[y:y+tile_size, x:x+tile_size]
 
            # 超解像を適用 (ここでは仮のアップスケール)
            upscaled_tile = upscale_single_tile(tile, scale)
 
            # 重み付けマスク (端を滑らかにブレンド)
            blend_mask = create_blend_mask(
                tile_size * scale, tile_size * scale,
                overlap * scale
            )
 
            # 結果に加算
            oy, ox = y * scale, x * scale
            ts = tile_size * scale
            output[oy:oy+ts, ox:ox+ts] += upscaled_tile * blend_mask
            weight[oy:oy+ts, ox:ox+ts] += blend_mask
 
    # 正規化
    output = output / np.maximum(weight, 1e-8)
    return output.astype(np.uint8)
 
def create_blend_mask(h, w, margin):
    """タイルブレンド用のグラデーションマスク"""
    mask = np.ones((h, w, 1), dtype=np.float32)
    # 四辺にグラデーションを適用
    for i in range(margin):
        alpha = i / margin
        mask[i, :, :] *= alpha      # 上辺
        mask[h-1-i, :, :] *= alpha  # 下辺
        mask[:, i, :] *= alpha      # 左辺
        mask[:, w-1-i, :] *= alpha  # 右辺
    return mask

class RealESRGANPipeline:
    """
    Real-ESRGAN を実務で使うための本格的なパイプライン
 
    機能:
    - 自動モデル選択 (画像種別判定)
    - バッチ処理
    - 品質評価
    - 進捗レポート
    """
 
    def __init__(self, device="cuda", half=True):
        self.device = device
        self.half = half
        self.models = {}
        self._load_models()
 
    def _load_models(self):
        """モデルの遅延ロード準備"""
        self.model_configs = {
            "photo": {
                "arch": RRDBNet(
                    num_in_ch=3, num_out_ch=3, num_feat=64,
                    num_block=23, num_grow_ch=32, scale=4,
                ),
                "path": "weights/RealESRGAN_x4plus.pth",
                "scale": 4,
            },
            "anime": {
                "arch": RRDBNet(
                    num_in_ch=3, num_out_ch=3, num_feat=64,
                    num_block=6, num_grow_ch=32, scale=4,
                ),
                "path": "weights/RealESRGAN_x4plus_anime_6B.pth",
                "scale": 4,
            },
        }
 
    def _get_upsampler(self, model_type, tile=0):
        """必要なモデルだけを遅延ロード"""
        if model_type not in self.models:
            config = self.model_configs[model_type]
            self.models[model_type] = RealESRGANer(
                scale=config["scale"],
                model_path=config["path"],
                model=config["arch"],
                tile=tile,
                tile_pad=10,
                pre_pad=0,
                half=self.half,
            )
        return self.models[model_type]
 
    def detect_image_type(self, image_path):
        """画像の種類を自動判別"""
        img = cv2.imread(image_path)
        if img is None:
            raise ValueError(f"画像読み込み失敗: {image_path}")
 
        # 色のユニーク数と分散で判断
        # アニメ/イラストは色数が少なく、グラデーションも少ない
        h, w = img.shape[:2]
        sample = img[::4, ::4]  # 4ピクセルおきにサンプリング
 
        unique_colors = len(np.unique(
            sample.reshape(-1, 3), axis=0
        ))
        color_ratio = unique_colors / (sample.shape[0] * sample.shape[1])
 
        # エッジの鮮明さ
        gray = cv2.cvtColor(sample, cv2.COLOR_BGR2GRAY)
        edges = cv2.Canny(gray, 50, 150)
        edge_ratio = np.sum(edges > 0) / edges.size
 
        # 判定ロジック
        if color_ratio < 0.3 and edge_ratio > 0.05:
            return "anime"
        else:
            return "photo"
 
    def upscale(self, image_path, output_path=None,
                model_type=None, scale=4, tile=0):
        """
        画像を超解像
 
        Parameters:
            image_path: 入力画像パス
            output_path: 出力画像パス (None=自動命名)
            model_type: "photo" / "anime" / None (自動判定)
            scale: 出力倍率 (2 or 4)
            tile: タイルサイズ (0=無効, 400-512推奨)
 
        Returns:
            dict: 処理結果 (出力パス、メタデータ)
        """
        import time
        start_time = time.time()
 
        # モデル選択
        if model_type is None:
            model_type = self.detect_image_type(image_path)
 
        # VRAM に応じたタイルサイズ自動決定
        if tile == 0:
            img = cv2.imread(image_path)
            h, w = img.shape[:2]
            total_pixels = h * w
            if total_pixels > 2_000_000:  # 2MP以上
                tile = 400
            elif total_pixels > 4_000_000:  # 4MP以上
                tile = 256
 
        upsampler = self._get_upsampler(model_type, tile=tile)
 
        img = cv2.imread(image_path, cv2.IMREAD_UNCHANGED)
        output, _ = upsampler.enhance(img, outscale=scale)
 
        if output_path is None:
            from pathlib import Path
            p = Path(image_path)
            output_path = str(p.parent / f"{p.stem}_upscaled{p.suffix}")
 
        cv2.imwrite(output_path, output)
 
        elapsed = time.time() - start_time
 
        return {
            "input": image_path,
            "output": output_path,
            "model": model_type,
            "scale": scale,
            "input_size": f"{img.shape[1]}x{img.shape[0]}",
            "output_size": f"{output.shape[1]}x{output.shape[0]}",
            "elapsed_seconds": round(elapsed, 2),
        }
 
    def batch_upscale(self, input_dir, output_dir,
                      model_type=None, scale=4, tile=400,
                      extensions=("png", "jpg", "jpeg", "webp")):
        """ディレクトリ内の全画像を一括超解像"""
        from pathlib import Path
        import json
 
        input_path = Path(input_dir)
        output_path = Path(output_dir)
        output_path.mkdir(parents=True, exist_ok=True)
 
        files = []
        for ext in extensions:
            files.extend(input_path.glob(f"*.{ext}"))
            files.extend(input_path.glob(f"*.{ext.upper()}"))
 
        results = []
        total = len(files)
 
        for i, f in enumerate(sorted(files)):
            print(f"[{i+1}/{total}] Processing: {f.name}")
            try:
                result = self.upscale(
                    str(f),
                    str(output_path / f.name),
                    model_type=model_type,
                    scale=scale,
                    tile=tile,
                )
                result["status"] = "success"
            except Exception as e:
                result = {
                    "input": str(f),
                    "status": "error",
                    "error": str(e),
                }
            results.append(result)
 
        # レポート出力
        report_path = output_path / "upscale_report.json"
        with open(report_path, "w", encoding="utf-8") as fout:
            json.dump(results, fout, indent=2, ensure_ascii=False)
 
        success_count = sum(
            1 for r in results if r["status"] == "success"
        )
        print(f"\n完了: {success_count}/{total} 成功")
        return results

from diffusers import StableDiffusionUpscalePipeline
from PIL import Image
import torch
 
# SD x4 Upscaler パイプライン
pipe = StableDiffusionUpscalePipeline.from_pretrained(
    "stabilityai/stable-diffusion-x4-upscaler",
    torch_dtype=torch.float16,
).to("cuda")
 
# 低解像度画像を読み込み
low_res = Image.open("small_image.png")  # 例: 256x256
 
# プロンプトガイド付きアップスケール
upscaled = pipe(
    prompt="high resolution, sharp details, photorealistic, 8K",
    negative_prompt="blurry, noisy, artifacts, low quality",
    image=low_res,
    num_inference_steps=25,
    guidance_scale=4.0,    # 低めが推奨 (高すぎるとアーティファクト)
    noise_level=20,        # ノイズレベル (0-350, 推奨: 20-50)
).images[0]
 
upscaled.save("high_res.png")  # 出力: 1024x1024

"""
StableSR: Stable Diffusion の事前知識を活用した超解像
 
特徴:
- Stable Diffusion の学習済み知識をそのまま利用
- Time-aware Encoder で忠実度と品質のバランスを制御
- CFW (Controllable Feature Wrapping) モジュール
"""
 
from diffusers import (
    StableDiffusionPipeline,
    DDPMScheduler,
)
import torch
 
class StableSRPipeline:
    """StableSR の概念的な実装"""
 
    def __init__(self, sd_model_path, sr_module_path, device="cuda"):
        self.device = device
 
        # Stable Diffusion ベースモデル
        self.sd_pipe = StableDiffusionPipeline.from_pretrained(
            sd_model_path,
            torch_dtype=torch.float16,
        ).to(device)
 
        # 超解像用の追加モジュール
        self.encoder_module = self._load_sr_encoder(sr_module_path)
 
    def upscale(self, image, scale=4,
                positive_prompt="",
                negative_prompt="blurry, artifacts",
                num_steps=50,
                color_fix="wavelet"):
        """
        StableSR によるアップスケール
 
        Parameters:
            image: PIL Image (低解像度)
            scale: 拡大倍率
            positive_prompt: 品質向上プロンプト
            negative_prompt: 抑制プロンプト
            num_steps: 拡散ステップ数
            color_fix: 色補正方式 ("none", "adain", "wavelet")
        """
        # 1. 低解像度画像をエンコード
        lr_features = self.encoder_module.encode(image)
 
        # 2. 拡散過程で高解像度化
        #    Time-aware Encoder が各ステップの特徴を調整
        hr_latent = self.sd_pipe(
            prompt=positive_prompt,
            negative_prompt=negative_prompt,
            num_inference_steps=num_steps,
            # 独自の条件付けメカニズム
            cross_attention_kwargs={
                "lr_features": lr_features
            },
        )
 
        # 3. 色補正
        result = hr_latent.images[0]
        if color_fix == "wavelet":
            result = self._wavelet_color_fix(image, result)
        elif color_fix == "adain":
            result = self._adain_color_fix(image, result)
 
        return result
 
    def _wavelet_color_fix(self, source, target):
        """ウェーブレット変換による色補正"""
        import pywt
 
        # 低周波成分（色情報）を元画像から取得
        # 高周波成分（ディテール）を超解像結果から取得
        source_resized = source.resize(target.size, Image.LANCZOS)
 
        src_arr = np.array(source_resized).astype(np.float32)
        tgt_arr = np.array(target).astype(np.float32)
 
        result = np.zeros_like(tgt_arr)
 
        for ch in range(3):
            # ウェーブレット分解
            src_coeffs = pywt.dwt2(src_arr[:, :, ch], 'haar')
            tgt_coeffs = pywt.dwt2(tgt_arr[:, :, ch], 'haar')
 
            # 低周波は元画像、高周波は超解像結果
            new_coeffs = (src_coeffs[0], tgt_coeffs[1])
            result[:, :, ch] = pywt.idwt2(new_coeffs, 'haar')
 
        return Image.fromarray(
            np.clip(result, 0, 255).astype(np.uint8)
        )
 
    def _adain_color_fix(self, source, target):
        """AdaIN による色補正"""
        source_resized = source.resize(target.size, Image.LANCZOS)
        src = np.array(source_resized).astype(np.float32)
        tgt = np.array(target).astype(np.float32)
 
        for ch in range(3):
            src_mean, src_std = src[:,:,ch].mean(), src[:,:,ch].std()
            tgt_mean, tgt_std = tgt[:,:,ch].mean(), tgt[:,:,ch].std()
            tgt[:,:,ch] = (
                (tgt[:,:,ch] - tgt_mean) / (tgt_std + 1e-8)
            ) * src_std + src_mean
 
        return Image.fromarray(
            np.clip(tgt, 0, 255).astype(np.uint8)
        )

"""
SUPIR (Scaling Up to Excellence: Practicing Model Scaling
       for Photo-Realistic Image Restoration)
 
大規模言語モデル (SDXL) を活用した超解像。
テキストプロンプトで生成するディテールを制御可能。
"""
 
# SUPIR は通常CLIまたはGradioで使用
# pip install git+https://github.com/Fanghua-Yu/SUPIR.git
 
# CLI使用例:
# python inference.py \
#   --input_path input.png \
#   --output_path output.png \
#   --prompt "high quality photograph, sharp details" \
#   --upscale 4 \
#   --model_path SUPIR-v0Q.ckpt
 
# API的な使用 (概念コード):
class SUPIRUpscaler:
    """SUPIR超解像の概念的なラッパー"""
 
    def __init__(self, model_path, device="cuda"):
        self.device = device
        # モデルロード (実際にはSUPIRの設定に従う)
        self.model = self._load_model(model_path)
 
    def upscale(self, image, prompt="", scale=4,
                restoration_strength=0.7):
        """
        テキストガイド付き超解像
 
        restoration_strength:
          0.0-0.3: 忠実度重視 (元画像に近い)
          0.4-0.6: バランス
          0.7-1.0: 品質重視 (ディテール生成多め)
        """
        # 1. 低品質画像のエンコード
        # 2. テキストプロンプトのエンコード
        # 3. 拡散過程で高解像度画像を生成
        # 4. 元画像との整合性を保ちつつ復元
        pass
 
    def batch_upscale(self, image_dir, output_dir, **kwargs):
        """ディレクトリ内の全画像を一括処理"""
        from pathlib import Path
        for img_path in Path(image_dir).glob("*.{png,jpg,jpeg}"):
            img = Image.open(img_path)
            result = self.upscale(img, **kwargs)
            result.save(Path(output_dir) / img_path.name)

from gfpgan import GFPGANer
import cv2
import numpy as np
 
class FaceRestorationPipeline:
    """顔復元 + 超解像の統合パイプライン"""
 
    def __init__(self, device="cuda"):
        self.device = device
 
        # GFPGAN 顔復元モデル
        self.face_restorer = GFPGANer(
            model_path="weights/GFPGANv1.4.pth",
            upscale=4,
            arch="clean",
            channel_multiplier=2,
            bg_upsampler=self._create_bg_upsampler(),
        )
 
    def _create_bg_upsampler(self):
        """背景用の Real-ESRGAN アップサンプラー"""
        from realesrgan import RealESRGANer
        from basicsr.archs.rrdbnet_arch import RRDBNet
 
        model = RRDBNet(
            num_in_ch=3, num_out_ch=3, num_feat=64,
            num_block=23, num_grow_ch=32, scale=4,
        )
        return RealESRGANer(
            scale=4,
            model_path="weights/RealESRGAN_x4plus.pth",
            model=model,
            tile=400,
            tile_pad=10,
            pre_pad=0,
            half=True,
        )
 
    def restore(self, image_path, output_path=None,
                fidelity_weight=0.5, only_center_face=False):
        """
        顔復元 + 背景超解像
 
        Parameters:
            image_path: 入力画像パス
            fidelity_weight: 忠実度の重み (0=品質重視, 1=忠実度重視)
            only_center_face: 最大の顔のみ処理するか
 
        Returns:
            dict: 復元結果と検出情報
        """
        img = cv2.imread(image_path, cv2.IMREAD_COLOR)
 
        # 顔復元実行
        _, _, output = self.face_restorer.enhance(
            img,
            has_aligned=False,
            only_center_face=only_center_face,
            paste_back=True,
            weight=fidelity_weight,
        )
 
        if output_path:
            cv2.imwrite(output_path, output)
 
        return {
            "output": output,
            "input_size": f"{img.shape[1]}x{img.shape[0]}",
            "output_size": f"{output.shape[1]}x{output.shape[0]}",
        }
 
 
    def batch_restore(self, input_dir, output_dir,
                      fidelity_weight=0.5):
        """ディレクトリ内の全画像を顔復元"""
        from pathlib import Path
 
        input_path = Path(input_dir)
        output_path = Path(output_dir)
        output_path.mkdir(parents=True, exist_ok=True)
 
        results = []
        for img_file in sorted(input_path.glob("*")):
            if img_file.suffix.lower() in (
                ".png", ".jpg", ".jpeg", ".webp"
            ):
                try:
                    result = self.restore(
                        str(img_file),
                        str(output_path / img_file.name),
                        fidelity_weight=fidelity_weight,
                    )
                    result["status"] = "success"
                    result["file"] = img_file.name
                except Exception as e:
                    result = {
                        "file": img_file.name,
                        "status": "error",
                        "error": str(e),
                    }
                results.append(result)
 
        return results

"""
CodeFormer: 離散コードブックを用いた顔復元
 
GFPGAN との違い:
- コードブック: VQ-VAE で学習した離散表現を使用
- 忠実度制御: パラメータ w で連続的に制御可能
  w=0: 高品質だが忠実度低い
  w=1: 忠実度高いが品質は低い
  推奨: w=0.5-0.7
"""
 
# pip install codeformer-pip
 
def restore_face_codeformer(
    image_path,
    output_path,
    fidelity_weight=0.5,
    upscale=4,
    detection_model="retinaface_resnet50",
):
    """CodeFormer による顔復元"""
    import subprocess
 
    cmd = [
        "python", "inference_codeformer.py",
        "-i", image_path,
        "-o", output_path,
        "-w", str(fidelity_weight),
        "-s", str(upscale),
        "--detection_model", detection_model,
        "--bg_upsampler", "realesrgan",
        "--face_upsample",
    ]
 
    result = subprocess.run(cmd, capture_output=True, text=True)
 
    if result.returncode != 0:
        raise RuntimeError(f"CodeFormer error: {result.stderr}")
 
    return output_path

"""
SwinIR: Swin Transformer を超解像に適用
 
特徴:
- CNN より広い受容野 (ウィンドウアテンション)
- シフトウィンドウで隣接ウィンドウ間の情報交換
- 忠実度が高い (GAN系ほどディテールを「想像」しない)
"""
 
import torch
from PIL import Image
import numpy as np
 
def upscale_with_swinir(image_path, scale=4, task="real_sr"):
    """
    SwinIR による超解像
 
    task:
      "classical_sr": 古典的超解像 (bicubic劣化想定)
      "real_sr": 実世界超解像 (複合劣化想定)
      "lightweight_sr": 軽量版
      "jpeg_car": JPEG圧縮アーティファクト除去
      "color_dn": カラーノイズ除去
      "gray_dn": グレースケールノイズ除去
    """
    # モデル設定
    model_configs = {
        "classical_sr": {
            "model_path": "weights/001_classicalSR_DF2K_s64w8_"
                         f"SwinIR-M_x{scale}.pth",
            "window_size": 8,
            "img_size": 64,
        },
        "real_sr": {
            "model_path": "weights/003_realSR_BSRGAN_DFOWMFC_"
                         f"s64w8_SwinIR-L_x{scale}_GAN.pth",
            "window_size": 8,
            "img_size": 64,
        },
        "lightweight_sr": {
            "model_path": "weights/002_lightweightSR_DIV2K_"
                         f"s64w8_SwinIR-S_x{scale}.pth",
            "window_size": 8,
            "img_size": 64,
        },
    }
 
    config = model_configs[task]
 
    # 画像読み込みとパディング
    img = np.array(Image.open(image_path)).astype(np.float32) / 255.0
    img = torch.from_numpy(img).permute(2, 0, 1).unsqueeze(0)
 
    # ウィンドウサイズの倍数にパディング
    ws = config["window_size"]
    _, _, h, w = img.shape
    pad_h = (ws - h % ws) % ws
    pad_w = (ws - w % ws) % ws
    img = torch.nn.functional.pad(img, (0, pad_w, 0, pad_h), mode="reflect")
 
    # モデルロードと推論
    model = torch.load(config["model_path"])
    model.eval()
 
    with torch.no_grad():
        output = model(img.to("cuda"))
 
    # パディング除去
    output = output[:, :, :h * scale, :w * scale]
 
    # Tensor → PIL Image
    output = output.squeeze(0).permute(1, 2, 0).cpu().numpy()
    output = np.clip(output * 255, 0, 255).astype(np.uint8)
    return Image.fromarray(output)

import cv2
from pathlib import Path
import time
 
class VideoUpscaler:
    """動画超解像パイプライン"""
 
    def __init__(self, model_type="realesrgan", scale=4, device="cuda"):
        self.scale = scale
        self.model_type = model_type
        self._init_model(device)
 
    def _init_model(self, device):
        """モデル初期化"""
        if self.model_type == "realesrgan":
            from basicsr.archs.rrdbnet_arch import RRDBNet
            from realesrgan import RealESRGANer
 
            model = RRDBNet(
                num_in_ch=3, num_out_ch=3, num_feat=64,
                num_block=23, num_grow_ch=32, scale=self.scale,
            )
            self.upsampler = RealESRGANer(
                scale=self.scale,
                model_path=f"weights/RealESRGAN_x{self.scale}plus.pth",
                model=model,
                tile=400,
                tile_pad=10,
                half=True,
            )
 
    def upscale_video(self, input_path, output_path,
                      codec="libx264", crf=18,
                      audio_copy=True):
        """
        動画を超解像
 
        Parameters:
            input_path: 入力動画パス
            output_path: 出力動画パス
            codec: 出力コーデック
            crf: 品質 (0=ロスレス, 51=最低, 18推奨)
            audio_copy: 音声をコピーするか
        """
        cap = cv2.VideoCapture(input_path)
 
        fps = cap.get(cv2.CAP_PROP_FPS)
        total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
 
        out_w = w * self.scale
        out_h = h * self.scale
 
        # 一時ファイルに映像のみ出力
        tmp_video = output_path + ".tmp.mp4"
        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
        writer = cv2.VideoWriter(
            tmp_video, fourcc, fps, (out_w, out_h)
        )
 
        print(f"入力: {w}x{h} @ {fps}fps, {total_frames}フレーム")
        print(f"出力: {out_w}x{out_h}")
 
        start_time = time.time()
        frame_idx = 0
 
        while True:
            ret, frame = cap.read()
            if not ret:
                break
 
            # 超解像
            output, _ = self.upsampler.enhance(
                frame, outscale=self.scale
            )
            writer.write(output)
 
            frame_idx += 1
            if frame_idx % 10 == 0:
                elapsed = time.time() - start_time
                fps_actual = frame_idx / elapsed
                eta = (total_frames - frame_idx) / fps_actual
                print(
                    f"  [{frame_idx}/{total_frames}] "
                    f"{fps_actual:.1f} fps, "
                    f"ETA: {eta:.0f}s"
                )
 
        cap.release()
        writer.release()
 
        # ffmpeg で音声を結合
        if audio_copy:
            import subprocess
            cmd = [
                "ffmpeg", "-y",
                "-i", tmp_video,
                "-i", input_path,
                "-c:v", codec,
                "-crf", str(crf),
                "-c:a", "copy",
                "-map", "0:v:0",
                "-map", "1:a:0?",
                output_path,
            ]
            subprocess.run(cmd, check=True, capture_output=True)
            Path(tmp_video).unlink()
        else:
            Path(tmp_video).rename(output_path)
 
        total_time = time.time() - start_time
        print(f"\n完了: {total_time:.1f}秒 "
              f"({total_frames / total_time:.1f} fps)")
 
        return {
            "input": input_path,
            "output": output_path,
            "frames": total_frames,
            "elapsed_seconds": round(total_time, 1),
            "avg_fps": round(total_frames / total_time, 2),
        }

def temporal_consistent_upscale(frames, upsampler, flow_model=None):
    """
    時間的一貫性を保った超解像
 
    単純にフレーム毎に超解像すると、フレーム間で
    ディテールが揺れる (temporal flickering) 問題が発生する。
 
    対策:
    1. オプティカルフロー整合化
    2. 前フレームの結果を参照
    3. テンポラルブレンディング
    """
 
    results = []
    prev_output = None
 
    for i, frame in enumerate(frames):
        # 現フレームを超解像
        current_output, _ = upsampler.enhance(
            frame, outscale=4
        )
 
        if prev_output is not None and flow_model is not None:
            # オプティカルフローで前フレームをワープ
            flow = flow_model.estimate(
                frames[i-1], frame
            )
            warped_prev = warp_image(prev_output, flow, scale=4)
 
            # 現フレームとワープ前フレームをブレンド
            # → テンポラルフリッカリングを抑制
            alpha = 0.3  # ブレンド比率
            occlusion_mask = compute_occlusion_mask(flow)
 
            current_output = np.where(
                occlusion_mask[..., None],
                current_output,  # オクルージョン領域は現フレームのみ
                (1 - alpha) * current_output
                + alpha * warped_prev  # その他はブレンド
            ).astype(np.uint8)
 
        results.append(current_output)
        prev_output = current_output
 
    return results

from fastapi import FastAPI, UploadFile, File, HTTPException
from fastapi.responses import FileResponse
import tempfile
import uuid
from pathlib import Path
 
app = FastAPI(title="Super Resolution API")
 
# グローバルパイプライン
pipeline = RealESRGANPipeline(device="cuda", half=True)
 
@app.post("/upscale")
async def upscale_image(
    file: UploadFile = File(...),
    scale: int = 4,
    model: str = "auto",
):
    """
    画像をアップスケールする API エンドポイント
 
    Parameters:
        file: アップロード画像
        scale: 拡大倍率 (2 or 4)
        model: モデル名 ("photo", "anime", "auto")
    """
    # バリデーション
    if scale not in (2, 4):
        raise HTTPException(400, "scale must be 2 or 4")
 
    allowed_types = {"image/png", "image/jpeg", "image/webp"}
    if file.content_type not in allowed_types:
        raise HTTPException(400, f"Unsupported: {file.content_type}")
 
    # ファイルサイズ制限 (20MB)
    contents = await file.read()
    if len(contents) > 20 * 1024 * 1024:
        raise HTTPException(400, "File too large (max 20MB)")
 
    # 一時ファイルに保存
    job_id = str(uuid.uuid4())
    tmp_dir = Path(tempfile.mkdtemp())
    input_path = tmp_dir / f"input_{job_id}.png"
    output_path = tmp_dir / f"output_{job_id}.png"
 
    with open(input_path, "wb") as f:
        f.write(contents)
 
    try:
        model_type = None if model == "auto" else model
        result = pipeline.upscale(
            str(input_path),
            str(output_path),
            model_type=model_type,
            scale=scale,
            tile=400,
        )
    except Exception as e:
        raise HTTPException(500, f"Processing error: {str(e)}")
 
    return FileResponse(
        str(output_path),
        media_type="image/png",
        filename=f"upscaled_{file.filename}",
        headers={
            "X-Input-Size": result["input_size"],
            "X-Output-Size": result["output_size"],
            "X-Model": result["model"],
            "X-Elapsed": str(result["elapsed_seconds"]),
        },
    )
 
 
@app.post("/upscale/batch")
async def batch_upscale(
    files: list[UploadFile] = File(...),
    scale: int = 4,
    model: str = "auto",
):
    """複数画像の一括アップスケール"""
    if len(files) > 50:
        raise HTTPException(400, "Max 50 files per batch")
 
    results = []
    for file in files:
        try:
            # 各ファイルを個別処理
            result = await upscale_image(file, scale, model)
            results.append({
                "filename": file.filename,
                "status": "success",
            })
        except HTTPException as e:
            results.append({
                "filename": file.filename,
                "status": "error",
                "detail": e.detail,
            })
 
    return {"results": results}

import gradio as gr
from PIL import Image
import numpy as np
 
def create_upscaling_ui():
    """Gradio ベースの超解像 GUI"""
 
    pipeline = RealESRGANPipeline(device="cuda")
 
    def upscale_handler(image, model_choice, scale, tile_size):
        """超解像処理ハンドラ"""
        if image is None:
            return None, "画像をアップロードしてください"
 
        # PIL → 一時ファイル → 処理 → PIL
        import tempfile
        with tempfile.NamedTemporaryFile(
            suffix=".png", delete=False
        ) as tmp:
            Image.fromarray(image).save(tmp.name)
            result = pipeline.upscale(
                tmp.name,
                model_type=(
                    None if model_choice == "自動判定"
                    else model_choice.lower()
                ),
                scale=scale,
                tile=tile_size,
            )
 
        output = Image.open(result["output"])
        info = (
            f"入力: {result['input_size']} → "
            f"出力: {result['output_size']}\n"
            f"モデル: {result['model']}\n"
            f"処理時間: {result['elapsed_seconds']}秒"
        )
        return np.array(output), info
 
    with gr.Blocks(title="AI 超解像") as demo:
        gr.Markdown("# AI 超解像ツール")
 
        with gr.Row():
            with gr.Column():
                input_img = gr.Image(label="入力画像")
                model_choice = gr.Radio(
                    ["自動判定", "Photo", "Anime"],
                    label="モデル",
                    value="自動判定",
                )
                scale = gr.Slider(
                    minimum=2, maximum=4, step=2,
                    value=4, label="拡大倍率",
                )
                tile_size = gr.Slider(
                    minimum=0, maximum=800, step=100,
                    value=400, label="タイルサイズ (0=無効)",
                )
                btn = gr.Button("超解像を実行", variant="primary")
 
            with gr.Column():
                output_img = gr.Image(label="出力画像")
                info_text = gr.Textbox(label="処理情報")
 
        btn.click(
            fn=upscale_handler,
            inputs=[input_img, model_choice, scale, tile_size],
            outputs=[output_img, info_text],
        )
 
    return demo
 
# demo = create_upscaling_ui()
# demo.launch(server_name="0.0.0.0", server_port=7860)

import replicate
import requests
from pathlib import Path
 
class CloudUpscaler:
    """クラウド API を使った超解像 (ローカルGPU不要)"""
 
    def __init__(self, provider="replicate"):
        self.provider = provider
 
    def upscale_replicate(self, image_path, scale=4, model="real-esrgan"):
        """Replicate API で超解像"""
        model_versions = {
            "real-esrgan": "xinntao/realesrgan:latest",
            "supir": "cjwbw/supir:latest",
            "swinir": "jingyunliang/swinir:latest",
        }
 
        with open(image_path, "rb") as f:
            output = replicate.run(
                model_versions[model],
                input={
                    "image": f,
                    "scale": scale,
                    "face_enhance": True,
                },
            )
 
        # 結果のダウンロード
        output_path = Path(image_path).stem + "_upscaled.png"
        if isinstance(output, str):
            response = requests.get(output)
            with open(output_path, "wb") as f:
                f.write(response.content)
        elif isinstance(output, list):
            response = requests.get(output[0])
            with open(output_path, "wb") as f:
                f.write(response.content)
 
        return output_path
 
    def upscale_runpod(self, image_path, scale=4):
        """RunPod Serverless で超解像"""
        import base64
 
        with open(image_path, "rb") as f:
            img_base64 = base64.b64encode(f.read()).decode()
 
        response = requests.post(
            "https://api.runpod.ai/v2/<endpoint_id>/run",
            headers={
                "Authorization": f"Bearer {RUNPOD_API_KEY}",
                "Content-Type": "application/json",
            },
            json={
                "input": {
                    "image": img_base64,
                    "scale": scale,
                    "model": "realesrgan",
                },
            },
        )
 
        job_id = response.json()["id"]
 
        # ポーリングで結果を取得
        import time
        while True:
            status = requests.get(
                f"https://api.runpod.ai/v2/<endpoint_id>/status/{job_id}",
                headers={
                    "Authorization": f"Bearer {RUNPOD_API_KEY}",
                },
            ).json()
 
            if status["status"] == "COMPLETED":
                output_b64 = status["output"]["image"]
                output_bytes = base64.b64decode(output_b64)
                output_path = Path(image_path).stem + "_upscaled.png"
                with open(output_path, "wb") as f:
                    f.write(output_bytes)
                return output_path
 
            elif status["status"] == "FAILED":
                raise RuntimeError(f"RunPod error: {status}")
 
            time.sleep(2)

class UpscaleErrorHandler:
    """超解像でよくあるエラーのハンドリング"""
 
    @staticmethod
    def handle_oom(func):
        """OOM エラー時にタイルサイズを縮小して再試行"""
        import functools
 
        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            tile_sizes = [0, 512, 400, 256, 128]
            last_error = None
 
            for tile_size in tile_sizes:
                try:
                    kwargs["tile"] = tile_size
                    return func(*args, **kwargs)
                except RuntimeError as e:
                    if "out of memory" in str(e).lower():
                        last_error = e
                        # GPU メモリをクリア
                        import torch
                        torch.cuda.empty_cache()
                        if tile_size == 0:
                            print(
                                "VRAM不足: タイル処理に切り替え "
                                f"(tile={tile_sizes[1]})"
                            )
                        else:
                            print(
                                f"VRAM不足: タイル縮小 "
                                f"(tile={tile_size}→次のサイズ)"
                            )
                    else:
                        raise
 
            raise RuntimeError(
                f"全タイルサイズで OOM: {last_error}"
            )
 
        return wrapper
 
    @staticmethod
    def validate_input(image_path):
        """入力画像のバリデーション"""
        import cv2
        from pathlib import Path
 
        path = Path(image_path)
 
        # ファイル存在チェック
        if not path.exists():
            raise FileNotFoundError(f"ファイルが見つかりません: {path}")
 
        # 拡張子チェック
        valid_exts = {".png", ".jpg", ".jpeg", ".webp", ".bmp", ".tiff"}
        if path.suffix.lower() not in valid_exts:
            raise ValueError(
                f"非対応フォーマット: {path.suffix}\n"
                f"対応: {valid_exts}"
            )
 
        # 画像読み込みテスト
        img = cv2.imread(str(path))
        if img is None:
            raise ValueError(f"画像読み込み失敗 (破損?): {path}")
 
        h, w = img.shape[:2]
 
        # サイズチェック
        if h < 16 or w < 16:
            raise ValueError(
                f"画像が小さすぎます: {w}x{h} (最小: 16x16)"
            )
 
        if h > 10000 or w > 10000:
            print(
                f"警告: 大画像 ({w}x{h})。"
                "タイル処理を推奨します。"
            )
 
        # メモリ見積もり (4x超解像)
        estimated_vram_gb = (
            w * h * 3 * 4 * 16  # 概算
        ) / (1024 ** 3)
 
        return {
            "path": str(path),
            "size": f"{w}x{h}",
            "channels": img.shape[2] if len(img.shape) > 2 else 1,
            "estimated_vram_gb": round(estimated_vram_gb, 2),
            "recommendation": (
                "タイル処理推奨" if w * h > 2_000_000
                else "タイル処理不要"
            ),
        }

def post_process_color_correction(
    original_path, upscaled_path, output_path,
    method="histogram_matching"
):
    """
    超解像後の色補正
 
    超解像モデルは色味を変える場合がある。
    元画像の色分布を参照して補正する。
    """
    import cv2
    import numpy as np
 
    original = cv2.imread(original_path)
    upscaled = cv2.imread(upscaled_path)
 
    # 元画像を超解像サイズにリサイズ (参照用)
    h, w = upscaled.shape[:2]
    original_resized = cv2.resize(
        original, (w, h), interpolation=cv2.INTER_LANCZOS4
    )
 
    if method == "histogram_matching":
        # ヒストグラムマッチング
        result = np.zeros_like(upscaled)
        for ch in range(3):
            result[:, :, ch] = _match_histogram(
                upscaled[:, :, ch],
                original_resized[:, :, ch],
            )
 
    elif method == "color_transfer":
        # LAB色空間での統計的色転写
        result = _lab_color_transfer(original_resized, upscaled)
 
    elif method == "linear":
        # 線形回帰による色補正
        result = _linear_color_correction(
            original_resized, upscaled
        )
 
    cv2.imwrite(output_path, result)
    return output_path
 
 
def _match_histogram(source, reference):
    """ヒストグラムマッチング (単一チャネル)"""
    src_values, src_unique_indices, src_counts = np.unique(
        source.ravel(), return_inverse=True, return_counts=True
    )
    ref_values, ref_counts = np.unique(
        reference.ravel(), return_counts=True
    )
 
    src_cdf = np.cumsum(src_counts).astype(np.float64)
    src_cdf /= src_cdf[-1]
 
    ref_cdf = np.cumsum(ref_counts).astype(np.float64)
    ref_cdf /= ref_cdf[-1]
 
    interp_values = np.interp(src_cdf, ref_cdf, ref_values)
    return interp_values[src_unique_indices].reshape(source.shape).astype(
        np.uint8
    )
 
 
def _lab_color_transfer(source, target):
    """LAB色空間での色転写 (Reinhard et al.)"""
    source_lab = cv2.cvtColor(source, cv2.COLOR_BGR2LAB).astype(np.float32)
    target_lab = cv2.cvtColor(target, cv2.COLOR_BGR2LAB).astype(np.float32)
 
    for ch in range(3):
        src_mean = source_lab[:, :, ch].mean()
        src_std = source_lab[:, :, ch].std()
        tgt_mean = target_lab[:, :, ch].mean()
        tgt_std = target_lab[:, :, ch].std()
 
        target_lab[:, :, ch] = (
            (target_lab[:, :, ch] - tgt_mean)
            * (src_std / (tgt_std + 1e-8))
            + src_mean
        )
 
    target_lab = np.clip(target_lab, 0, 255).astype(np.uint8)
    return cv2.cvtColor(target_lab, cv2.COLOR_LAB2BGR)

# BasicSR の設定ファイル例
# finetune_realesrgan_x4plus.yml
 
name: finetune_RealESRGANx4plus
model_type: RealESRGANModel
scale: 4
num_gpu: 1
 
datasets:
  train:
    name: custom_dataset
    type: RealESRGANPairedDataset
    dataroot_gt: /data/train/HR  # 高解像度画像
    dataroot_lq: /data/train/LR  # 低解像度画像
    io_backend:
      type: disk
    gt_size: 256
    use_hflip: true
    use_rot: true
 
# ネットワーク設定
network_g:
  type: RRDBNet
  num_in_ch: 3
  num_out_ch: 3
  num_feat: 64
  num_block: 23
  num_grow_ch: 32
  scale: 4
 
# 学習設定
train:
  optim_g:
    type: Adam
    lr: !!float 1e-4
    weight_decay: 0
    betas: [0.9, 0.99]
 
  # 事前学習済みモデルから開始
  path:
    pretrain_network_g: weights/RealESRGAN_x4plus.pth
    strict_load_g: true
 
  total_iter: 50000
  warmup_iter: -1

def ensemble_upscale(image_path, models, weights=None):
    """複数モデルの結果を重み付き平均で統合"""
    results = []
    for model in models:
        result = model.upscale(image_path)
        results.append(np.array(result).astype(np.float32))
 
    if weights is None:
        weights = [1.0 / len(models)] * len(models)
 
    ensemble = np.zeros_like(results[0])
    for result, weight in zip(results, weights):
        ensemble += result * weight
 
    return Image.fromarray(
        np.clip(ensemble, 0, 255).astype(np.uint8)
    )