pano_process.py

from pathlib import Path
import sys
import argparse
import torch
import torchvision.transforms.functional as trans_F

sys.path.append(str(Path(sys.path[0]).parent.absolute()))

from utils import img
from utils import math

parser = argparse.ArgumentParser()
parser.add_argument('-o', '--output', type=str)
parser.add_argument('dir', type=str)
args = parser.parse_args()

data_dir = Path(args.dir)
output_dir = Path(args.output)
output_dir.mkdir(parents=True, exist_ok=True)

files = [file for file in data_dir.glob('*') if file.suffix == '.png' or file.suffix == '.jpg']
outfiles = [output_dir / file.name for file in data_dir.glob('*')
            if file.suffix == '.png' or file.suffix == '.jpg']
images = img.load(files)
print(f"{images.size(0)} images loaded.")
out_images = torch.zeros_like(images)
H, W = images.shape[-2:]
for row in range(H):
    phi = math.pi / H * (row + 0.5)
    length = math.ceil(math.sin(phi) * W * 0.5) * 2
    cols = slice((W - length) // 2, (W + length) // 2)
    out_images[..., row:row + 1, cols] = trans_F.resize(images[..., row:row + 1, :], [1, length])
    sys.stdout.write(f'{row + 1} / {H} processed.   \r')
print('')
img.save(out_images, outfiles)
print(f"{images.size(0)} images saved.")