deoldify / demo.py
ali-ghamdan's picture
publish
31a913a
#NOTE: This must be the first call in order to work properly!
from deoldify import device
from deoldify.device_id import DeviceId
from deoldify.visualize import *
import argparse
parser = argparse.ArgumentParser()
parser.add_argument("-i", "--input", type=str, required=True, help="the image path")
parser.add_argument("-o", "--output", type=str, required=True, help="the output path")
parser.add_argument("--artistic", action="store_true", default=True, help="enable artistic mode")
parser.add_argument("--compare", action="store_true", default=True, help="enable compare mode")
parser.add_argument("-r", "--render_factor", type=int, default=35, help="max is 45 for 11GB video cards. default is 35")
args = parser.parse_args()
device.set(device=DeviceId.GPU0)
colorizer = get_image_colorizer(artistic=args.artistic)
result_path = colorizer.plot_transformed_image(path=args.input, render_factor=args.render_factor, compare=args.compare, results_dir=args.output)
print(result_path)