File size: 1,606 Bytes
ec188ed | 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 | import argparse
from pathlib import Path
from starnet_v1_pytorch import StarNet
def main():
input_default = "rgb_test5.tif"
output_default = "rgb_test5_starless.tif"
weights_default = str(Path(__file__).resolve().parent / "models" / "weights")
mode_default = "RGB"
parser = argparse.ArgumentParser(description="PyTorch version of the TF2 StarNet transform demo")
parser.add_argument("input", nargs="?", default=input_default, help="Input TIFF image")
parser.add_argument(
"output",
nargs="?",
default=output_default,
help="Output TIFF image",
)
parser.add_argument(
"--weights",
default=weights_default,
help="Weight prefix or a direct .pth file path",
)
parser.add_argument("--mode", choices=("RGB", "Greyscale"), default=mode_default)
parser.add_argument("--window-size", type=int, default=512)
parser.add_argument("--stride", type=int, default=128)
parser.add_argument(
"--device",
default=None,
help="Torch device such as cuda, cuda:0, or cpu (default: auto-detect)",
)
args = parser.parse_args()
starnet = StarNet(
mode=args.mode,
window_size=args.window_size,
stride=args.stride,
device=args.device,
)
weights_path = starnet.load_model(args.weights)
print("Loaded weights: %s" % weights_path)
output_path = starnet.transform(args.input, args.output)
print("Saved result: %s" % Path(output_path).resolve())
if __name__ == "__main__":
main()
|