File size: 3,732 Bytes
ef7c54c 2fc2cd7 ef7c54c 2fc2cd7 ef7c54c | 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 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 | import os
os.environ["CUDA_VISIBLE_DEVICES"] = "1"
import argparse
from pathlib import Path
import onnxruntime as ort
import torch
import torchvision.transforms.functional as TF
from torchvision.utils import save_image
from PIL import Image, ImageChops, ImageFilter
import time
import random
import numpy as np
from model import UpscalerTransformer
def parse_args():
parser = argparse.ArgumentParser()
parser.add_argument(
"--image",
type=str,
required=True,
help="Путь к входному изображению.",
)
parser.add_argument(
"--weights",
type=str,
required=True,
help="Путь к чекпоинту best.pth или last.pth.",
)
parser.add_argument(
"--output",
type=str,
default="sr_output.png",
help="Куда сохранить результат.",
)
# Если чекпоинт без config, можно задать вручную.
parser.add_argument("--lr_size", type=int, default=128)
parser.add_argument("--dim", type=int, default=256)
parser.add_argument("--depth", type=int, default=8)
parser.add_argument("--heads", type=int, default=8)
parser.add_argument("--window", type=int, default=8)
parser.add_argument("--resize_type", type=int, default=Image.Resampling.LANCZOS)
return parser.parse_args()
def main():
args = parse_args()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Device: {device}")
ckpt = torch.load(args.weights, map_location="cpu")
if isinstance(ckpt, dict):
config = ckpt.get("config", {})
state_dict = ckpt.get("model", ckpt)
else:
config = {}
state_dict = ckpt
model = UpscalerTransformer(
dim=config.get("dim", args.dim),
depth=config.get("depth", args.depth),
num_heads=config.get("heads", args.heads),
window_size=config.get("window", args.window),
)
model.load_state_dict(state_dict)
model.to(device)
model.eval()
"""dummy_input = torch.randn(1, 3, 128, 128)
torch.onnx.export(
model,
dummy_input,
"last.onnx",
export_params=True,
opset_version=17, # >= 16 для поддержки bicubic
do_constant_folding=True, # Оптимизация констант
input_names=['input'],
output_names=['output'],
)"""
lr_size = config.get("lr_size", args.lr_size)
image = Image.open(args.image).convert("RGB")
image = image.resize((lr_size, lr_size), Image.Resampling.LANCZOS)
image_hr = image.resize((lr_size * 2, lr_size * 2), args.resize_type)
x = TF.to_tensor(image).unsqueeze(0).to(device)
with torch.no_grad():
if device.type == "cuda":
with torch.cuda.amp.autocast():
sr = model(x)
else:
sr = model(x)
sr = torch.clamp(sr[0].float().cpu(), 0.0, 1.0)
image_result = sr.permute(1, 2, 0).numpy()
image_result = (image_result * 255.0).astype(np.uint8)
if image_result.shape[2] == 1:
image_result = image_result.squeeze(axis=2)
output_path = Path(args.output)
output_path.parent.mkdir(parents=True, exist_ok=True)
result = Image.new(image_hr.mode, (512, 256))
result.paste(image_hr, (0, 0))
result.paste(Image.fromarray(image_result), (256, 0))
result.save(output_path)
print(f"Saved: {output_path}")
if __name__ == "__main__":
main() |