| 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="Куда сохранить результат.",
|
| )
|
|
|
|
|
| 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() |
