Commit 5aed2cc0 authored by Yuxin Wu's avatar Yuxin Wu

update train script

parent eedec322
...@@ -16,6 +16,9 @@ from tensorpack.utils.utils import mkdir_p ...@@ -16,6 +16,9 @@ from tensorpack.utils.utils import mkdir_p
parser = argparse.ArgumentParser() parser = argparse.ArgumentParser()
parser.add_argument(dest='config') parser.add_argument(dest='config')
parser.add_argument('-o', '--output', help='output directory to dump dataset image') parser.add_argument('-o', '--output', help='output directory to dump dataset image')
parser.add_argument('--index',
help='index of the image component in datapoint',
default=0, type=int)
parser.add_argument('-n', '--number', help='number of images to dump', parser.add_argument('-n', '--number', help='number of images to dump',
default=10, type=int) default=10, type=int)
args = parser.parse_args() args = parser.parse_args()
...@@ -28,7 +31,7 @@ config = get_config_func() ...@@ -28,7 +31,7 @@ config = get_config_func()
if args.output: if args.output:
mkdir_p(args.output) mkdir_p(args.output)
cnt = 0 cnt = 0
index = 0 # TODO: as an argument? index = args.index # TODO: as an argument?
for dp in config.dataset.get_data(): for dp in config.dataset.get_data():
imgbatch = dp[index] imgbatch = dp[index]
if cnt > args.number: if cnt > args.number:
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment