def find_checkpoint(checkpoint_str):
# find checkpoint
steps = 0
if checkpoint_str is not None:
if ':' in checkpoint_str:
prefix, steps = checkpoint_str.split(':')
else:
prefix = checkpoint_str
steps = None
log_file, run_id = path.find_log(prefix)
if steps is None:
checkpoint, steps = path.find_checkpoints(run_id)[-1]
else:
checkpoints = path.find_checkpoints(run_id)
try:
checkpoint, steps = next(
filter(lambda t: t[1] == steps, checkpoints))
except StopIteration:
print('The steps not found in checkpoints', steps, checkpoints)
sys.stdout.flush()
raise StopIteration
steps = int(steps)
if args.clear_steps:
steps = 0
else:
_, exp_info = path.read_log(log_file)
exp_info = exp_info[-1]
for k in args.__dict__:
if k in exp_info and k in ('tag', ):
setattr(args, k, eval(exp_info[k]))
print('{}={}, '.format(k, exp_info[k]), end='')
print()
sys.stdout.flush()
return checkpoint, steps
mkdir(os.path.join('logs', 'val'))
mkdir(os.path.join('logs', 'debug'))
mkdir('weights')
mkdir('flows')
# find checkpoint
import path
import logger
steps = 0
if args.checkpoint is not None:
if ':' in args.checkpoint:
prefix, steps = args.checkpoint.split(':')
else:
prefix = args.checkpoint
steps = None
log_file, run_id = path.find_log(prefix)
if steps is None:
checkpoint, steps = path.find_checkpoints(run_id)[-1]
else:
checkpoints = path.find_checkpoints(run_id)
try:
checkpoint, steps = next(filter(lambda t : t[1] == steps, checkpoints))
except StopIteration:
print('The steps not found in checkpoints', steps, checkpoints)
sys.stdout.flush()
raise StopIteration
steps = int(steps)
if args.clear_steps:
steps = 0
else:
_, exp_info = path.read_log(log_file)
def __init__(self, argv):
# MaskFlownet init with args
repoRoot = r'.'
os.environ['PATH'] = r'/usr/local/cuda/bin' + ';' + os.environ['PATH']
parser = argparse.ArgumentParser(add_help=False)
parser.add_argument('config', type=str, nargs='?', default=None)
parser.add_argument('-g',
'--gpu_device',
type=str,
default='',
help='Specify gpu device(s)')
parser.add_argument(
'-c',
'--checkpoint',
type=str,
default=None,
help='model checkpoint to load; by default, the latest one.'
'You can use checkpoint:steps to load to a specific steps')
parser.add_argument('-n', '--network', type=str, default='MaskFlownet')
parser.add_argument('--resize', type=str, default='')
args = parser.parse_args(
argv[1:]) # argv[1:] to ignore the 1st element (scipt's name)
ctx = [mx.cpu()] if args.gpu_device == '' else [
mx.gpu(gpu_id) for gpu_id in map(int, args.gpu_device.split(','))
]
self.infer_resize = [int(s) for s in args.resize.split(',')
] if args.resize else None
# load network configuration
import network.config
with open(os.path.join(repoRoot, 'network', 'config',
args.config)) as f:
config = network.config.Reader(yaml.load(f))
# find checkpoint
import path
prefix = args.checkpoint
log_file, run_id = path.find_log(prefix)
checkpoint, steps = path.find_checkpoints(run_id)[-1]
# initiate
from network import get_pipeline
self.pipe = get_pipeline(args.network, ctx=ctx, config=config)
# load parameters from given checkpoint
print('Load Checkpoint {}'.format(checkpoint))
sys.stdout.flush()
network_class = getattr(config.network, 'class').get()
print('Load the weight for the network')
self.pipe.load(checkpoint)
if network_class == 'MaskFlownet':
print('Fix the weight for the head network')
self.pipe.fix_head()
sys.stdout.flush()
# ROS topics
self.cameraTopic = "bebop/image_raw"
self.flowTopic = "flow"
self.camera_sub = rospy.Subscriber(self.cameraTopic, Image,
self.camera_cb)
self.flow_pub = rospy.Publisher(self.flowTopic, Image, queue_size=1)
self.image_msg = Image()
self.bridge = CvBridge()
self.prevCvImage = None
self.currCvImage_msg = None
self.imageRcvd_fl = False
self.img1 = []
self.img2 = []