File size: 4,466 Bytes
c85d0ce 20b04f8 c85d0ce 20b04f8 c85d0ce 20b04f8 c85d0ce 20b04f8 c85d0ce a01cc5b c85d0ce 20b04f8 c85d0ce 20b04f8 c85d0ce 20b04f8 c85d0ce 20b04f8 c85d0ce 20b04f8 c85d0ce 20b04f8 c85d0ce |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 |
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# The file source is from the [ESRGAN](https://github.com/xinntao/ESRGAN) project
# forked by authors [joeyballentine](https://github.com/joeyballentine/ESRGAN) and [BlueAmulet](https://github.com/BlueAmulet/ESRGAN).
import gc
import numpy as np
import torch
def bgr_to_rgb(image: torch.Tensor) -> torch.Tensor:
# flip image channels
# https://github.com/pytorch/pytorch/issues/229
out: torch.Tensor = image.flip(-3)
# out: torch.Tensor = image[[2, 1, 0], :, :] #RGB to BGR #may be faster
return out
def rgb_to_bgr(image: torch.Tensor) -> torch.Tensor:
# same operation as bgr_to_rgb(), flip image channels
return bgr_to_rgb(image)
def bgra_to_rgba(image: torch.Tensor) -> torch.Tensor:
out: torch.Tensor = image[[2, 1, 0, 3], :, :]
return out
def rgba_to_bgra(image: torch.Tensor) -> torch.Tensor:
# same operation as bgra_to_rgba(), flip image channels
return bgra_to_rgba(image)
def auto_split_upscale(
lr_img: np.ndarray,
upscale_function,
scale: int = 4,
overlap: int = 32,
max_depth: int = None,
current_depth: int = 1,
current_tile: int = 1, # Tracks the current tile being processed
total_tiles: int = 1, # Total number of tiles at this depth level
):
# Attempt to upscale if unknown depth or if reached known max depth
if max_depth is None or max_depth == current_depth:
try:
print(f"auto_split_upscale depth: {current_depth}", end=" ", flush=True)
result, _ = upscale_function(lr_img, scale)
print(f"progress: {current_tile}/{total_tiles}")
return result, current_depth
except RuntimeError as e:
# Check to see if its actually the CUDA out of memory error
if "CUDA" in str(e):
print("RuntimeError: CUDA out of memory...")
# Re-raise the exception if not an OOM error
else:
raise RuntimeError(e)
# Collect garbage (clear VRAM)
torch.cuda.empty_cache()
gc.collect()
input_h, input_w, input_c = lr_img.shape
# Split the image into 4 quadrants with some overlap
top_left = lr_img[: input_h // 2 + overlap, : input_w // 2 + overlap, :]
top_right = lr_img[: input_h // 2 + overlap, input_w // 2 - overlap :, :]
bottom_left = lr_img[input_h // 2 - overlap :, : input_w // 2 + overlap, :]
bottom_right = lr_img[input_h // 2 - overlap :, input_w // 2 - overlap :, :]
current_depth = current_depth + 1
current_tile = (current_tile - 1) * 4
total_tiles = total_tiles * 4
# Recursively upscale each quadrant and track the current tile number
# After we go through the top left quadrant, we know the maximum depth and no longer need to test for out-of-memory
top_left_rlt, depth = auto_split_upscale(
top_left, upscale_function, scale=scale, overlap=overlap, max_depth=max_depth,
current_depth=current_depth, current_tile=current_tile + 1, total_tiles=total_tiles,
)
top_right_rlt, _ = auto_split_upscale(
top_right, upscale_function, scale=scale, overlap=overlap, max_depth=depth,
current_depth=current_depth, current_tile=current_tile + 2, total_tiles=total_tiles,
)
bottom_left_rlt, _ = auto_split_upscale(
bottom_left, upscale_function, scale=scale, overlap=overlap, max_depth=depth,
current_depth=current_depth, current_tile=current_tile + 3, total_tiles=total_tiles,
)
bottom_right_rlt, _ = auto_split_upscale(
bottom_right, upscale_function, scale=scale, overlap=overlap, max_depth=depth,
current_depth=current_depth, current_tile=current_tile + 4, total_tiles=total_tiles,
)
# Define the output image size
out_h = input_h * scale
out_w = input_w * scale
# Create an empty output image
output_img = np.zeros((out_h, out_w, input_c), np.uint8)
# Fill the output image with the upscaled quadrants, removing overlap regions
output_img[: out_h // 2, : out_w // 2, :] = top_left_rlt[: out_h // 2, : out_w // 2, :]
output_img[: out_h // 2, -out_w // 2 :, :] = top_right_rlt[: out_h // 2, -out_w // 2 :, :]
output_img[-out_h // 2 :, : out_w // 2, :] = bottom_left_rlt[-out_h // 2 :, : out_w // 2, :]
output_img[-out_h // 2 :, -out_w // 2 :, :] = bottom_right_rlt[-out_h // 2 :, -out_w // 2 :, :]
return output_img, depth
|