acherkrau/super-resolution-api-test

1import pathlib2from dataclasses import dataclass3 4import cv25import numpy as np6import redis7from logger import logger8 9from config import settings10 11try:12    logger.info(13        f"Connecting to Redis: {settings.get('redis_url', 'redis://localhost:6379')}"14    )15    redis_client = redis.from_url(settings.get("redis_url", "redis://localhost:6379"))16    logger.info(f"Redis ping: {redis_client.ping()}")17except Exception as e:18    logger.error(f"Failed to connect to Redis: {e}")19    exit(1)20 21 22@dataclass23class ModelInfo:24    name: str = ""25    path: str = ""26    scale: int = 427    algo: str = ""28 29 30BASE_STREAM_NAME = (31    "super_resolution_api_queue"32    if not settings.get("worker_id")33    else f"super_resolution_api_queue_{settings.get('worker_id')}"34)35WORKER_KEY_PREFIX = "super_resolution_api_worker_"36DISTRIBUTED_STREAM_NAME = "super_resolution_api_distributed_queue"37RESULT_KEY_PREFIX = (38    "super_resolution_api_result_"39    if not settings.get("worker_id")40    else f"super_resolution_api_result_{settings.get('worker_id')}_"41)42PROGRESS_TIMEOUT = settings.get("timeout", 30)43MAX_ALLOWED_TIMEOUT = settings.get("max_timeout", 300)44MAX_THREAD = settings.get("max_thread", 8)45MODEL_NAME_DEFAULT = "x4_Anime_6B-Official"46MODEL_NAME_X4_JP_ILLUSTRATION_FIX1 = "x4_JP_Illustration-fix1"47MODEL_NAME_X4_JP_ILLUSTRATION_FIX2 = "x4_JP_Illustration-fix2"48MODEL_NAME_X4_JP_ILLUSTRATION_FIX1_D = "x4_JP_Illustration-fix1-d"49MODEL_NAME_X4_ANIME_6B_OFFICIAL = "x4_Anime_6B-Official"50 51 52model_Anime_Official = ModelInfo(53    MODEL_NAME_X4_ANIME_6B_OFFICIAL,54    "models/x4_Anime_6B-Official.onnx",55    4,56    "real-esrgan",57)58 59model_JP_Illustration_fix1 = ModelInfo(60    MODEL_NAME_X4_JP_ILLUSTRATION_FIX1,61    "models/x4_jp_Illustration-fix1.onnx",62    4,63    "real-hatgan",64)65 66 67model_JP_Illustration_fix2 = ModelInfo(68    MODEL_NAME_X4_JP_ILLUSTRATION_FIX2,69    "models/x4_jp_Illustration-fix2.onnx",70    4,71    "real-esrgan",72)73 74model_JP_Illustration_fix1_d = ModelInfo(75    MODEL_NAME_X4_JP_ILLUSTRATION_FIX1_D,76    "models/x4_jp_Illustration-fix1-d.onnx",77    4,78    "real-esrgan",79)80 81models = {82    MODEL_NAME_X4_ANIME_6B_OFFICIAL: model_Anime_Official,83    MODEL_NAME_X4_JP_ILLUSTRATION_FIX1: model_JP_Illustration_fix1,84    MODEL_NAME_X4_JP_ILLUSTRATION_FIX2: model_JP_Illustration_fix2,85    MODEL_NAME_X4_JP_ILLUSTRATION_FIX1_D: model_JP_Illustration_fix1_d,86}87 88 89def get_image_size(image_path: pathlib.Path) -> tuple[int, int]:90    """91    return: (width, height)92    """93    img = cv2.imread(str(image_path))94    if img is None:95        raise Exception(f"Failed to load image: {image_path}")96    return img.shape[1], img.shape[0]97 98 99@dataclass100class TileInfo:101    x: int102    y: int103    filpath: pathlib.Path104 105 106def split_image(107    img_path: pathlib.Path,108    save_dir: pathlib.Path,109    grid_size: tuple[int, int],110    overlap: int = 16,111) -> list[TileInfo]:112    save_path = pathlib.Path(save_dir)113    save_path.mkdir(parents=True, exist_ok=True)114 115    img = cv2.imread(str(img_path))116    if img is None:117        raise Exception(f"Failed to load image: {img_path}")118 119    height, width = img.shape[:2]120    rows, cols = grid_size121 122    base_h = height // rows123    base_w = width // cols124 125    tiles_info = []126 127    for row in range(rows):128        for col in range(cols):129            x1 = max(0, col * base_w - overlap)130            y1 = max(0, row * base_h - overlap)131            x2 = min(width, (col + 1) * base_w + overlap)132            y2 = min(height, (row + 1) * base_h + overlap)133 134            tile = img[y1:y2, x1:x2]135            tile_name = f"{img_path.stem}_tile_{row}_{col}.png"136            tile_path = save_path / tile_name137            cv2.imwrite(str(tile_path), tile)138            tiles_info.append(TileInfo(col, row, tile_path))139 140    return tiles_info141 142 143def merge_sr_tiles(144    tiles: list[TileInfo],145    output: pathlib.Path,146    original_size: tuple[int, int],147    scale: int,148    overlap: int = 16,149):150    """151    合并超分辨率后的图块152 153    tiles: 超分辨率后的图块信息列表, 需要根据 filepath 读取图块, 根据 x, y 位置信息进行拼接154    output: 合并后的图片保存路径155    original_size: 原始图片的尺寸156    overlap 为原始图片切割时设定的重叠像素数157    scale 为超分辨率倍数158    """159    # Calculate output dimensions160    logger.debug(161        f"正在合并 {len(tiles)} 张超分辨率图块, 原尺寸: {original_size}, 缩放倍数: {scale}"162    )163 164    width, height = original_size165    out_width = width * scale166    out_height = height * scale167    output_img = np.zeros((out_height, out_width, 3), dtype=np.uint8)168 169    # Calculate base tile sizes170    rows = max([t.y for t in tiles]) + 1171    cols = max([t.x for t in tiles]) + 1172    base_h = height // rows173    base_w = width // cols174 175    # Scale dimensions176    scaled_base_h = base_h * scale177    scaled_base_w = base_w * scale178    scaled_overlap = overlap * scale179 180    for tile_info in tiles:181        # Read tile182        tile = cv2.imread(str(tile_info.filpath))183        if tile is None:184            raise Exception(f"Failed to load tile: {tile_info.filpath}")185 186        # Calculate positions187        x1 = max(0, tile_info.x * scaled_base_w - scaled_overlap)188        y1 = max(0, tile_info.y * scaled_base_h - scaled_overlap)189        x2 = min(out_width, (tile_info.x + 1) * scaled_base_w + scaled_overlap)190        y2 = min(out_height, (tile_info.y + 1) * scaled_base_h + scaled_overlap)191 192        # Calculate blend mask for overlapping regions193        h, w = y2 - y1, x2 - x1194        blend_mask = np.ones((h, w, 1), dtype=np.float32)195 196        # Apply feathering at edges197        if tile_info.x > 0:  # Left edge198            blend_mask[:, :scaled_overlap] = np.linspace(0, 1, scaled_overlap).reshape(199                1, -1, 1200            )201        if tile_info.x < cols - 1:  # Right edge202            blend_mask[:, -scaled_overlap:] = np.linspace(1, 0, scaled_overlap).reshape(203                1, -1, 1204            )205        if tile_info.y > 0:  # Top edge206            blend_mask[:scaled_overlap, :] *= np.linspace(0, 1, scaled_overlap).reshape(207                -1, 1, 1208            )209        if tile_info.y < rows - 1:  # Bottom edge210            blend_mask[-scaled_overlap:, :] *= np.linspace(211                1, 0, scaled_overlap212            ).reshape(-1, 1, 1)213 214        # Blend tiles215        output_img[y1:y2, x1:x2] = (216            output_img[y1:y2, x1:x2] * (1 - blend_mask)217            + tile[: y2 - y1, : x2 - x1] * blend_mask218        ).astype(np.uint8)219 220    cv2.imwrite(output.as_posix(), output_img)221 222 223def calculate_grid(image_width, image_height, workers):224    if workers <= 0:225        raise ValueError("Worker count must be positive")226 227    best_rows, best_cols = 1, workers228    min_aspect_diff = float("inf")229 230    for rows in range(1, workers + 1):231        if workers % rows == 0:232            cols = workers // rows233            tile_width = image_width / cols234            tile_height = image_height / rows235            aspect_ratio = max(tile_width, tile_height) / min(tile_width, tile_height)236            aspect_diff = aspect_ratio - 1237 238            if aspect_diff < min_aspect_diff:239                best_rows, best_cols = rows, cols240                min_aspect_diff = aspect_diff241    logger.debug(f"calculate_grid: {best_rows}x{best_cols}")242    return best_rows, best_cols243