Browse Source

Added "gpu-info" subcommand

master
wisp101 3 years ago
parent
commit
50fdb856f8
  1. 28
      gpu_info.py
  2. 78
      main.py

28
gpu_info.py

@ -0,0 +1,28 @@ @@ -0,0 +1,28 @@
import torch
import json
def get_info():
has_cuda = torch.cuda.is_available()
devices = []
if has_cuda:
count = torch.cuda.device_count()
for i in range(count):
devices.append(torch.cuda.get_device_name(i))
return {
"has_cuda": has_cuda,
"devices": devices,
}
def main(args):
info = get_info()
if args.json:
data = json.dumps(info)
print(data)
else:
print("Has Cuda: {}".format(info["has_cuda"]))
for (i, device) in enumerate(info["devices"]):
print("GPU {}: {}".format(i, device))

78
main.py

@ -10,6 +10,7 @@ import cv2 @@ -10,6 +10,7 @@ import cv2
import time
import os
import imageio
import sentry_sdk
import rook
@ -21,10 +22,14 @@ from multiprocessing import freeze_support @@ -21,10 +22,14 @@ from multiprocessing import freeze_support
from multiprocessing.pool import ThreadPool
from dotenv import load_dotenv
import gpu_info
#
load_dotenv()
parser = argparse.ArgumentParser()
subparsers = parser.add_subparsers()
parser.add_argument(
"-i", "--input", default="input.png", help="path of the photo to transform"
)
@ -88,6 +93,7 @@ parser.add_argument( @@ -88,6 +93,7 @@ parser.add_argument(
parser.add_argument(
"--n_cores", type=int, default=4, help="number of cpu cores to use (default: 4)",
)
scale_mod = parser.add_mutually_exclusive_group()
scale_mod.add_argument(
"--auto-resize",
@ -108,6 +114,16 @@ scale_mod.add_argument( @@ -108,6 +114,16 @@ scale_mod.add_argument(
help="Scale image to 512x512",
)
gpu_info_parser = subparsers.add_parser('gpu-info')
gpu_info_parser.add_argument(
"-j",
"--json",
default=False,
action="store_true",
help="Print GPU info as JSON"
)
def check_crops_coord():
def type_func(a):
@ -130,13 +146,15 @@ scale_mod.add_argument( @@ -130,13 +146,15 @@ scale_mod.add_argument(
def check_json_args_file():
def type_func(a):
if not os.path.isfile(a):
raise argparse.ArgumentTypeError("Arguments json file {} not found.".format(a))
raise argparse.ArgumentTypeError(
"Arguments json file {} not found.".format(a))
with open(a) as f:
data = {}
try:
data = json.load(f)
except JSONDecodeError:
raise argparse.ArgumentTypeError("Arguments json file {} is not in valid JSON format.".format(a))
raise argparse.ArgumentTypeError(
"Arguments json file {} is not in valid JSON format.".format(a))
l = []
for k, v in data.items():
if not isinstance(v, bool):
@ -156,20 +174,6 @@ parser.add_argument( @@ -156,20 +174,6 @@ parser.add_argument(
"If a command line argument is also provide the json value will be ignore for this argument.",
)
args = parser.parse_args()
# Handle special cases for ignoring arguments in json file if provided in command line
if args.json_args:
l = args.json_args
if "--cpu" in sys.argv[1:] or "--gpu" in sys.argv[1:]:
l = list(filter(lambda a: a not in ("--cpu", "--gpu"), l))
if "--auto-resize" in sys.argv[1:] or "--auto-resize-crop" in sys.argv[1:] \
or "--auto-rescale" in sys.argv[1:] or "--overlay" in sys.argv[1:]:
l = list(filter(lambda a: a not in ("--auto-resize", "--auto-resize-crop", "--auto-rescale", "--overlay"), l))
args = parser.parse_args(l + sys.argv[1:])
"""
main.py
@ -178,11 +182,13 @@ main.py @@ -178,11 +182,13 @@ main.py
"""
# ------------------------------------------------- main()
def main():
def main(args):
if not os.path.isfile(args.input):
print("Error : {} file doesn't exist".format(args.input), file=sys.stderr)
print("Error : {} file doesn't exist".format(
args.input), file=sys.stderr)
exit(1)
start = time.time()
@ -210,13 +216,15 @@ def main(): @@ -210,13 +216,15 @@ def main():
# See if image loaded correctly
if image is None:
print("Error : {} file is not valid".format(args.input), file=sys.stderr)
print("Error : {} file is not valid".format(
args.input), file=sys.stderr)
exit(1)
# Preprocess
if args.overlay:
original_image = image.copy()
image = utils.crop_input(image, args.overlay[0], args.overlay[1], args.overlay[2], args.overlay[3])
image = utils.crop_input(
image, args.overlay[0], args.overlay[1], args.overlay[2], args.overlay[3])
elif args.auto_resize:
image = utils.resize_input(image)
elif args.auto_resize_crop:
@ -226,7 +234,8 @@ def main(): @@ -226,7 +234,8 @@ def main():
# See if image has the correct shape after preprocessing
if image.shape != (512, 512, 3):
print("Error : image is not 512 x 512, got shape: {}".format(image.shape), file=sys.stderr)
print("Error : image is not 512 x 512, got shape: {}".format(
image.shape), file=sys.stderr)
exit(1)
# Process
@ -239,7 +248,8 @@ def main(): @@ -239,7 +248,8 @@ def main():
cv2.imwrite(args.output, result)
else:
base_output_filename = utils.strip_file_extension(args.output, ".png")
base_output_filename = utils.strip_file_extension(
args.output, ".png")
def process_one_image(i):
result = process(image, gpu_ids, args.enablepubes)
@ -291,6 +301,26 @@ def main(): @@ -291,6 +301,26 @@ def main():
sys.exit()
# Register Command Handlers
parser.set_defaults(func=main)
gpu_info_parser.set_defaults(func=gpu_info.main)
args = parser.parse_args()
# Handle special cases for ignoring arguments in json file if provided in command line
if args.json_args:
l = args.json_args
if "--cpu" in sys.argv[1:] or "--gpu" in sys.argv[1:]:
l = list(filter(lambda a: a not in ("--cpu", "--gpu"), l))
if "--auto-resize" in sys.argv[1:] or "--auto-resize-crop" in sys.argv[1:] \
or "--auto-rescale" in sys.argv[1:] or "--overlay" in sys.argv[1:]:
l = list(filter(lambda a: a not in ("--auto-resize",
"--auto-resize-crop", "--auto-rescale", "--overlay"), l))
args = parser.parse_args(l + sys.argv[1:])
def process_gif_wrapper(gif_imgs, filename, gpu_ids, enablepubes, n_cores):
tmp_dir = tempfile.mkdtemp()
process_gif(gif_imgs, gpu_ids, enablepubes, tmp_dir, n_cores)
@ -323,4 +353,4 @@ if __name__ == "__main__": @@ -323,4 +353,4 @@ if __name__ == "__main__":
freeze_support()
start_sentry()
start_rook()
main()
args.func(args)

Loading…
Cancel
Save