Merge pull request #1459 from asimard1/master
Add progress bar for SwinIR upscale in cmd
This commit is contained in:
commit
46588c582d
|
@ -5,6 +5,7 @@ import numpy as np
|
||||||
import torch
|
import torch
|
||||||
from PIL import Image
|
from PIL import Image
|
||||||
from basicsr.utils.download_util import load_file_from_url
|
from basicsr.utils.download_util import load_file_from_url
|
||||||
|
from tqdm import tqdm
|
||||||
|
|
||||||
from modules import modelloader
|
from modules import modelloader
|
||||||
from modules.paths import models_path
|
from modules.paths import models_path
|
||||||
|
@ -122,18 +123,20 @@ def inference(img, model, tile, tile_overlap, window_size, scale):
|
||||||
E = torch.zeros(b, c, h * sf, w * sf, dtype=torch.half, device=device).type_as(img)
|
E = torch.zeros(b, c, h * sf, w * sf, dtype=torch.half, device=device).type_as(img)
|
||||||
W = torch.zeros_like(E, dtype=torch.half, device=device)
|
W = torch.zeros_like(E, dtype=torch.half, device=device)
|
||||||
|
|
||||||
for h_idx in h_idx_list:
|
with tqdm(total=len(h_idx_list) * len(w_idx_list), desc="SwinIR tiles") as pbar:
|
||||||
for w_idx in w_idx_list:
|
for h_idx in h_idx_list:
|
||||||
in_patch = img[..., h_idx: h_idx + tile, w_idx: w_idx + tile]
|
for w_idx in w_idx_list:
|
||||||
out_patch = model(in_patch)
|
in_patch = img[..., h_idx: h_idx + tile, w_idx: w_idx + tile]
|
||||||
out_patch_mask = torch.ones_like(out_patch)
|
out_patch = model(in_patch)
|
||||||
|
out_patch_mask = torch.ones_like(out_patch)
|
||||||
|
|
||||||
E[
|
E[
|
||||||
..., h_idx * sf: (h_idx + tile) * sf, w_idx * sf: (w_idx + tile) * sf
|
..., h_idx * sf: (h_idx + tile) * sf, w_idx * sf: (w_idx + tile) * sf
|
||||||
].add_(out_patch)
|
].add_(out_patch)
|
||||||
W[
|
W[
|
||||||
..., h_idx * sf: (h_idx + tile) * sf, w_idx * sf: (w_idx + tile) * sf
|
..., h_idx * sf: (h_idx + tile) * sf, w_idx * sf: (w_idx + tile) * sf
|
||||||
].add_(out_patch_mask)
|
].add_(out_patch_mask)
|
||||||
|
pbar.update(1)
|
||||||
output = E.div_(W)
|
output = E.div_(W)
|
||||||
|
|
||||||
return output
|
return output
|
||||||
|
|
Loading…
Reference in New Issue
Block a user