def run(self) -> None:
model_chain = (self.model_str.split("+")
if "+" in self.model_str else self.model_str.split(">"))
for idx, model in enumerate(model_chain):
interpolations = (model.split("|")
if "|" in self.model_str else model.split("&"))
if len(interpolations) > 1:
for i, interpolation in enumerate(interpolations):
interp_model, interp_amount = (interpolation.split("@")
if "@" in interpolation else
interpolation.split(":"))
interp_model = self.__check_model_path(interp_model)
interpolations[i] = f"{interp_model}@{interp_amount}"
model_chain[idx] = "&".join(interpolations)
else:
model_chain[idx] = self.__check_model_path(model)
if not self.input.exists():
self.log.error(f'Folder "{self.input}" does not exist.')
sys.exit(1)
elif self.input.is_file():
self.log.error(f'Folder "{self.input}" is a file.')
sys.exit(1)
elif self.output.is_file():
self.log.error(f'Folder "{self.output}" is a file.')
sys.exit(1)
elif not self.output.exists():
self.output.mkdir(parents=True)
print('Model{:s}: "{:s}"'.format(
"s" if len(model_chain) > 1 else "",
# ", ".join([Path(x).stem for x in model_chain]),
", ".join([x for x in model_chain]),
))
images: List[Path] = []
for ext in ["png", "jpg", "jpeg", "gif", "bmp", "tiff", "tga"]:
images.extend(self.input.glob(f"**/*.{ext}"))
# Store the maximum split depths for each model in the chain
# TODO: there might be a better way of doing this but it's good enough for now
split_depths = {}
with Progress(
# SpinnerColumn(),
"[progress.description]{task.description}",
BarColumn(),
"[progress.percentage]{task.percentage:>3.0f}%",
TimeRemainingColumn(),
) as progress:
task_upscaling = progress.add_task("Upscaling", total=len(images))
for idx, img_path in enumerate(images, 1):
img_input_path_rel = img_path.relative_to(self.input)
output_dir = self.output.joinpath(img_input_path_rel).parent
img_output_path_rel = output_dir.joinpath(
f"{img_path.stem}.png")
output_dir.mkdir(parents=True, exist_ok=True)
if len(model_chain) == 1:
self.log.info(
f'Processing {str(idx).zfill(len(str(len(images))))}: "{img_input_path_rel}"'
)
if self.skip_existing and img_output_path_rel.is_file():
self.log.warning("Already exists, skipping")
if self.delete_input:
img_path.unlink(missing_ok=True)
progress.advance(task_upscaling)
continue
# read image
img = cv2.imread(str(img_path.absolute()),
cv2.IMREAD_UNCHANGED)
if len(img.shape) < 3:
img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
# Seamless modes
if self.seamless == SeamlessOptions.TILE:
img = cv2.copyMakeBorder(img, 16, 16, 16, 16,
cv2.BORDER_WRAP)
elif self.seamless == SeamlessOptions.MIRROR:
img = cv2.copyMakeBorder(img, 16, 16, 16, 16,
cv2.BORDER_REFLECT_101)
elif self.seamless == SeamlessOptions.REPLICATE:
img = cv2.copyMakeBorder(img, 16, 16, 16, 16,
cv2.BORDER_REPLICATE)
elif self.seamless == SeamlessOptions.ALPHA_PAD:
img = cv2.copyMakeBorder(img,
16,
16,
16,
16,
cv2.BORDER_CONSTANT,
value=[0, 0, 0, 0])
final_scale: int = 1
task_model_chain: TaskID = None
if len(model_chain) > 1:
task_model_chain = progress.add_task(
f'{str(idx).zfill(len(str(len(images))))} - "{img_input_path_rel}"',
total=len(model_chain),
)
for i, model_path in enumerate(model_chain):
img_height, img_width = img.shape[:2]
# Load the model so we can access the scale
self.load_model(model_path)
if self.cache_max_split_depth and len(
split_depths.keys()) > 0:
rlt, depth = ops.auto_split_upscale(
img,
self.upscale,
self.last_scale,
max_depth=split_depths[i],
)
else:
rlt, depth = ops.auto_split_upscale(
img, self.upscale, self.last_scale)
split_depths[i] = depth
final_scale *= self.last_scale
# This is for model chaining
img = rlt.astype("uint8")
if len(model_chain) > 1:
progress.advance(task_model_chain)
if self.seamless:
rlt = self.crop_seamless(rlt, final_scale)
cv2.imwrite(str(img_output_path_rel.absolute()), rlt)
if self.delete_input:
img_path.unlink(missing_ok=True)
progress.advance(task_upscaling)
def image(
self,
img: np.ndarray,
device: torch.device = None,
# progress: Progress = None,
# progress_text: str = "",
multi_gpu_release_device=True,
) -> np.ndarray:
self.in_nc = None
self.out_nc = None
# Store the maximum split depths for each model in the chain
split_depths = {}
if len(img.shape) < 3:
img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
# Seamless modes
if self.seamless == SeamlessOptions.tile:
img = cv2.copyMakeBorder(img, 16, 16, 16, 16, cv2.BORDER_WRAP)
elif self.seamless == SeamlessOptions.mirror:
img = cv2.copyMakeBorder(img, 16, 16, 16, 16,
cv2.BORDER_REFLECT_101)
elif self.seamless == SeamlessOptions.replicate:
img = cv2.copyMakeBorder(img, 16, 16, 16, 16, cv2.BORDER_REPLICATE)
elif self.seamless == SeamlessOptions.alpha_pad:
img = cv2.copyMakeBorder(img,
16,
16,
16,
16,
cv2.BORDER_CONSTANT,
value=[0, 0, 0, 0])
final_scale: int = 1
# task_model_chain: TaskID = None
# if len(self.model_chain) > 1 and progress:
# task_model_chain = progress.add_task(
# progress_text, total=len(self.model_chain)
# )
if device == None:
if self.multi_gpu:
device, _ = self.get_available_device()
else:
device = list(self.devices.keys())[0]
for i, model_path in enumerate(self.model_chain):
# img_height, img_width = img.shape[:2]
# Load the model so we can access the scale
self.load_model(model_path)
if self.cache_max_split_depth and len(split_depths.keys()) > 0:
rlt, depth = ops.auto_split_upscale(
img,
device,
self.upscale,
self.last_scale,
max_depth=split_depths[i],
)
else:
rlt, depth = ops.auto_split_upscale(img, device, self.upscale,
self.last_scale)
split_depths[i] = depth
final_scale *= self.last_scale
# This is for model chaining
img = rlt.astype("uint8")
# if len(self.model_chain) > 1 and progress:
# progress.advance(task_model_chain)
if self.seamless:
img = self.crop_seamless(img, final_scale)
# img = img.astype("uint8")
if self.multi_gpu and multi_gpu_release_device:
self.devices[device][0].release()
return img