def load_model(config_file, device):
config_folder = os.path.join('config_folder', config_file)
option_file = os.path.join(config_folder, 'options.pickle')
checkpoint_file = os.path.join(config_folder, 'checkpoint.pyt')
checkpoint = torch.load(checkpoint_file, map_location='cpu')
_, hidden_config, noise_config = utils.load_options(option_file)
noiser = Noiser(noise_config, device)
model = Hidden(hidden_config, device, noiser, tb_logger=None)
utils.model_from_checkpoint(model, checkpoint)
return model
def system_setup(run_data):
"""Setup the system."""
run_data["options"] = utils.load_options()
if run_data["options"].verbosity > 0:
print("Options set as:")
print(run_data["options"])
if 'cuda' in (dev_name := run_data["options"].model_params['device']):
p = not torch.cuda.is_available()
c = int(dev_name[dev_name.find(':') + 1:]) + 1
q = c > torch.cuda.device_count()
if p or q:
raise utils.errors.MisconfigurationError(
f"Device '{dev_name}' is not available on this machine"
)
'--device',
'-d',
type=str,
required=False,
default="cameras",
help='Device use to collect data. It can be either "webcam" or "cameras"')
args = parser.parse_args()
# pass args
exp = args.exp
subject_id = args.subject
label_id = args.label
device = args.device
# load camera info --> options.json
op, cam = utils.load_options(device)
# load experiment labels --> labels.json
labels = utils.load_labels(exp)
log = Logger(name="Capture")
# get folder to store the collected data
if exp == "emotions":
folder = op.folder_emotions
elif exp == "signals":
folder = op.folder_signals
elif exp == "gestures":
folder = op.folder_gestures
elif exp == "adl":
folder = op.folder_adl
elif exp == "falls":
def main():
device = torch.device(
'cuda') if torch.cuda.is_available() else torch.device('cpu')
parent_parser = argparse.ArgumentParser(
description='Training of HiDDeN nets')
subparsers = parent_parser.add_subparsers(dest='command',
help='Sub-parser for commands')
new_run_parser = subparsers.add_parser('new', help='starts a new run')
new_run_parser.add_argument('--data-dir',
'-d',
required=True,
type=str,
help='The directory where the data is stored.')
# Anno dir
new_run_parser.add_argument(
'--anno-dir',
'-a',
type=str,
help=
'The directory where the annotations are stored. Specify only if you have annotations in a different folder.'
)
new_run_parser.add_argument('--batch-size',
'-b',
required=True,
type=int,
help='The batch size.')
new_run_parser.add_argument('--epochs',
'-e',
default=300,
type=int,
help='Number of epochs to run the simulation.')
new_run_parser.add_argument('--name',
required=True,
type=str,
help='The name of the experiment.')
new_run_parser.add_argument(
'--size',
'-s',
default=128,
type=int,
help=
'The size of the images (images are square so this is height and width).'
)
new_run_parser.add_argument('--message',
'-m',
default=256,
type=int,
help='The length in bits of the watermark.')
new_run_parser.add_argument(
'--continue-from-folder',
'-c',
default='',
type=str,
help=
'The folder from where to continue a previous run. Leave blank if you are starting a new experiment.'
)
# parser.add_argument('--tensorboard', dest='tensorboard', action='store_true',
# help='If specified, use adds a Tensorboard log. On by default')
new_run_parser.add_argument('--tensorboard',
action='store_true',
help='Use to switch on Tensorboard logging.')
new_run_parser.add_argument('--enable-fp16',
dest='enable_fp16',
action='store_true',
help='Enable mixed-precision training.')
new_run_parser.add_argument(
'--noise',
nargs='*',
action=NoiseArgParser,
help=
"Noise layers configuration. Use quotes when specifying configuration, e.g. 'cropout((0.55, 0.6), (0.55, 0.6))'"
)
new_run_parser.set_defaults(tensorboard=False)
new_run_parser.set_defaults(enable_fp16=False)
new_run_parser.add_argument('--vocab-path',
'-v',
type=str,
default='./data/vocab.pkl',
help='load the vocab')
continue_parser = subparsers.add_parser('continue',
help='Continue a previous run')
continue_parser.add_argument(
'--folder',
'-f',
required=True,
type=str,
help='Continue from the last checkpoint in this folder.')
continue_parser.add_argument(
'--data-dir',
'-d',
required=False,
type=str,
help=
'The directory where the data is stored. Specify a value only if you want to override the previous value.'
)
# Anno dir
continue_parser.add_argument(
'--anno-dir',
'-a',
required=False,
type=str,
help=
'The directory where the annotations are stored. Specify a value only if you want to override the previous value.'
)
continue_parser.add_argument(
'--epochs',
'-e',
required=False,
type=int,
help=
'Number of epochs to run the simulation. Specify a value only if you want to override the previous value.'
)
args = parent_parser.parse_args()
checkpoint = None
loaded_checkpoint_file_name = None
with open(args.vocab_path, 'rb') as f:
vocab = pickle.load(f)
if args.command == 'continue':
this_run_folder = args.folder
options_file = os.path.join(this_run_folder,
'options-and-config.pickle')
train_options, hidden_config, noise_config = utils.load_options(
options_file)
checkpoint, loaded_checkpoint_file_name = utils.load_last_checkpoint(
os.path.join(this_run_folder, 'checkpoints'))
train_options.start_epoch = checkpoint['epoch'] + 1
if args.data_dir is not None:
train_options.train_folder = os.path.join(args.data_dir, 'train')
train_options.validation_folder = os.path.join(
args.data_dir, 'val')
if args.epochs is not None:
if train_options.start_epoch < args.epochs:
train_options.number_of_epochs = args.epochs
else:
print(
f'Command-line specifies of number of epochs = {args.epochs}, but folder={args.folder} '
f'already contains checkpoint for epoch = {train_options.start_epoch}.'
)
exit(1)
else:
assert args.command == 'new'
start_epoch = 1
train_options = TrainingOptions(
batch_size=args.batch_size,
number_of_epochs=args.epochs,
train_folder=os.path.join(args.data_dir, 'train'),
validation_folder=os.path.join(args.data_dir, 'val'),
ann_train=os.path.join(args.data_dir, 'ann_train.json'),
ann_val=os.path.join(args.data_dir, 'ann_val.json'),
runs_folder=os.path.join('.', 'runs'),
start_epoch=start_epoch,
experiment_name=args.name)
noise_config = args.noise if args.noise is not None else []
hidden_config = HiDDenConfiguration(H=args.size,
W=args.size,
message_length=args.message,
encoder_blocks=4,
encoder_channels=64,
decoder_blocks=7,
decoder_channels=64,
use_discriminator=True,
use_vgg=False,
discriminator_blocks=3,
discriminator_channels=64,
decoder_loss=1,
encoder_loss=0.7,
adversarial_loss=1e-3,
vocab_size=len(vocab),
enable_fp16=args.enable_fp16)
this_run_folder = utils.create_folder_for_run(
train_options.runs_folder, args.name)
with open(os.path.join(this_run_folder, 'options-and-config.pickle'),
'wb+') as f:
pickle.dump(train_options, f)
pickle.dump(noise_config, f)
pickle.dump(hidden_config, f)
logging.basicConfig(level=logging.INFO,
format='%(message)s',
handlers=[
logging.FileHandler(
os.path.join(
this_run_folder,
f'{train_options.experiment_name}.log')),
logging.StreamHandler(sys.stdout)
])
if (args.command == 'new' and args.tensorboard) or \
(args.command == 'continue' and os.path.isdir(os.path.join(this_run_folder, 'tb-logs'))):
logging.info('Tensorboard is enabled. Creating logger.')
from tensorboard_logger import TensorBoardLogger
tb_logger = TensorBoardLogger(os.path.join(this_run_folder, 'tb-logs'))
else:
tb_logger = None
noiser = Noiser(noise_config, device)
model = Hidden(hidden_config, device, noiser, tb_logger)
if args.command == 'continue':
# if we are continuing, we have to load the model params
assert checkpoint is not None
logging.info(
f'Loading checkpoint from file {loaded_checkpoint_file_name}')
utils.model_from_checkpoint(model, checkpoint)
logging.info('HiDDeN model: {}\n'.format(model.to_stirng()))
logging.info('Model Configuration:\n')
logging.info(pprint.pformat(vars(hidden_config)))
logging.info('\nNoise configuration:\n')
logging.info(pprint.pformat(str(noise_config)))
logging.info('\nTraining train_options:\n')
logging.info(pprint.pformat(vars(train_options)))
train(model, device, hidden_config, train_options, this_run_folder,
tb_logger, vocab)
def main():
service_name = 'GestureRecognizer.Recognition'
log = Logger(name=service_name)
op = load_options()
channel = Channel(op.broker_uri)
log.info('Connected to broker {}', op.broker_uri)
exporter = create_exporter(service_name=service_name, uri=op.zipkin_uri)
subscription = Subscription(channel=channel, name=service_name)
for group_id in list(op.group_ids):
subscription.subscribe(
'SkeletonsGrouper.{}.Localization'.format(group_id))
model = GestureRecognizer("model_gesture1_72.00.pth")
log.info('Initialize the model')
unc = Gauge('uncertainty_total', "Uncertainty about predict")
unc.set(0.0)
start_http_server(8000)
buffer = list()
predict_flag = False
mean = lambda x: (sum(x) / len(x))
while True:
msg = channel.consume()
tracer = Tracer(exporter, span_context=msg.extract_tracing())
span = tracer.start_span(name='detection_and_info')
annotations = msg.unpack(ObjectAnnotations)
skeleton = select_skeletons(annotations=annotations,
min_keypoints=op.skeletons.min_keypoints,
x_range=op.skeletons.x_range,
y_range=op.skeletons.y_range)
if skeleton is None:
tracer.end_span()
continue
skl = Skeleton(skeleton)
skl_normalized = skl.normalize()
pred, prob, uncertainty = model.predict(skl_normalized)
if pred == 0 and predict_flag is False:
pass
elif pred != 0 and predict_flag is False:
initial_time = time.time()
predict_flag = True
buffer.append(uncertainty)
elif pred != 0 and predict_flag is True:
buffer.append(uncertainty)
elif pred == 0 and predict_flag is True:
predict_flag = False
exec_time = time.time() - initial_time
if exec_time >= op.exec_time:
unc.set(mean(buffer))
log.info("execution_ms: {}, buffer_mean: {}",
(exec_time * 1000), mean(buffer))
buffer = []
tracer.end_span()
info = {
'prediction': pred,
'probability': prob,
'uncertainty': uncertainty,
'took_ms': {
'service': round(span_duration_ms(span), 2)
}
}
log.info('{}', str(info).replace("'", '"'))
def main():
device = torch.device(
'cuda') if torch.cuda.is_available() else torch.device('cpu')
parser = argparse.ArgumentParser(description='Training of HiDDeN nets')
parser.add_argument('--size', '-s', default=128, type=int)
parser.add_argument('--data-dir', '-d', required=True, type=str)
parser.add_argument('--runs-folder',
'-sf',
default=os.path.join('.', 'runs'),
type=str)
parser.add_argument('--message', '-m', default=30, type=int)
parser.add_argument('--epochs', '-e', default=400, type=int)
parser.add_argument('--batch-size', '-b', required=True, type=int)
parser.add_argument('--continue-from-folder', '-c', default='', type=str)
parser.add_argument('--tensorboard',
dest='tensorboard',
action='store_true')
parser.add_argument('--no-tensorboard',
dest='tensorboard',
action='store_false')
parser.set_defaults(tensorboard=True)
args = parser.parse_args()
checkpoint = None
if args.continue_from_folder != '':
this_run_folder = args.continue_from_folder
train_options, hidden_config, noise_config = utils.load_options(
this_run_folder)
checkpoint = utils.load_last_checkpoint(
os.path.join(this_run_folder, 'checkpoints'))
train_options.start_epoch = checkpoint['epoch']
else:
start_epoch = 1
train_options = TrainingOptions(
batch_size=args.batch_size,
number_of_epochs=args.epochs,
train_folder=os.path.join(args.data_dir, 'train'),
validation_folder=os.path.join(args.data_dir, 'val'),
runs_folder=os.path.join('.', 'runs'),
start_epoch=start_epoch)
# noise_config = [
# {
# 'type': 'resize',
# 'resize_ratio': 0.4
# }]
noise_config = []
hidden_config = HiDDenConfiguration(H=args.size,
W=args.size,
message_length=args.message,
encoder_blocks=4,
encoder_channels=64,
decoder_blocks=7,
decoder_channels=64,
use_discriminator=True,
use_vgg=False,
discriminator_blocks=3,
discriminator_channels=64,
decoder_loss=1,
encoder_loss=0.7,
adversarial_loss=1e-3)
this_run_folder = utils.create_folder_for_run(train_options)
with open(os.path.join(this_run_folder, 'options-and-config.pickle'),
'wb+') as f:
pickle.dump(train_options, f)
pickle.dump(noise_config, f)
pickle.dump(hidden_config, f)
noiser = Noiser(noise_config, device)
if args.tensorboard:
print('Tensorboard is enabled. Creating logger.')
from tensorboard_logger import TensorBoardLogger
tb_logger = TensorBoardLogger(os.path.join(this_run_folder, 'tb-logs'))
else:
tb_logger = None
model = Hidden(hidden_config, device, noiser, tb_logger)
if args.continue_from_folder != '':
# if we are continuing, we have to load the model params
assert checkpoint is not None
utils.model_from_checkpoint(model, checkpoint)
print('HiDDeN model: {}\n'.format(model.to_stirng()))
print('Model Configuration:\n')
pprint.pprint(vars(hidden_config))
print('\nNoise configuration:\n')
pprint.pprint(str(noise_config))
print('\nTraining train_options:\n')
pprint.pprint(vars(train_options))
print()
train(model, device, hidden_config, train_options, this_run_folder,
tb_logger)
return resize(im, config)
def show_image(dataset, emotion, config=None):
"""
Display an image from an emotion given a dataset
"""
im = get_image(dataset, emotion, config)
if im is not None:
destroyWindows()
cv2.imshow(emotion.upper(), im)
dt, config = utils.load_options()
keys = utils.load_keymap()
# show command screen
cv2.imshow("COMMANDS", command_screen(keys))
while True:
k = cv2.waitKey(1)
if k == ord(keys.neutral):
show_image(choice(dt), "neutral", config)
elif k == ord(keys.angry):
show_image(choice(dt), "angry", config)
elif k == ord(keys.contempt):
g = cv2.merge([zero_array, g, zero_array])
r = cv2.merge([zero_array, zero_array, r])
resultado = None
if (k <= int(num_frames / 6)):
resultado = r
elif (k > int(num_frames / 6) and (k <= int(2 * num_frames / 6))):
resultado = g
else:
resultado = b
return resultado
log = Logger(name='WatchVideos')
with open('keymap.json', 'r') as f:
keymap = json.load(f)
options = load_options(print_options=False)
#categorias=['TIME_UP_GO_GRAYSCALE']
#categorias=['CAMINHADA_EM_CIRCULO_GRAYSCALE','IDA_E_VOLTA__DUAS_VEZES_GRAYSCALE']
categorias = [
'TIME_UP_GO_GRAYSCALE', 'CAMINHADA_EM_CIRCULO_GRAYSCALE',
'IDA_E_VOLTA__DUAS_VEZES_GRAYSCALE',
'ELEVACAO_EXCESSIVA_DO_CALCANHAR_GRASCALE', 'CIRCUNDACAO_DO_PE_GRAYSCALE',
'ASSIMETRIA_DO_COMPRIMENTO_DE_PASSO_GRAYSCALE'
]
for categoria in categorias:
t = 0
for l in range(0, 13):
aux_diferenca_frame = [0]
aux_soma_frame = [0]
frame_anterior = [0]
def main():
# device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
device = torch.device('cpu')
parser = argparse.ArgumentParser(description='Training of HiDDeN nets')
# parser.add_argument('--size', '-s', default=128, type=int, help='The size of the images (images are square so this is height and width).')
parser.add_argument('--data-dir',
'-d',
required=True,
type=str,
help='The directory where the data is stored.')
parser.add_argument(
'--runs_root',
'-r',
default=os.path.join('.', 'experiments'),
type=str,
help='The root folder where data about experiments are stored.')
args = parser.parse_args()
print_each = 25
completed_runs = [
o for o in os.listdir(args.runs_root)
if os.path.isdir(os.path.join(args.runs_root, o))
and o != 'no-noise-defaults'
]
print(completed_runs)
write_csv_header = True
for run_name in completed_runs:
current_run = os.path.join(args.runs_root, run_name)
print(f'Run folder: {current_run}')
options_file = os.path.join(current_run, 'options-and-config.pickle')
train_options, hidden_config, noise_config = utils.load_options(
options_file)
train_options.train_folder = os.path.join(args.data_dir, 'val')
train_options.validation_folder = os.path.join(args.data_dir, 'val')
train_options.batch_size = 4
checkpoint = utils.load_last_checkpoint(
os.path.join(current_run, 'checkpoints'))
noiser = Noiser(noise_config, device)
model = Hidden(hidden_config, device, noiser, tb_logger=None)
utils.model_from_checkpoint(model, checkpoint)
print('Model loaded successfully. Starting validation run...')
_, val_data = utils.get_data_loaders(hidden_config, train_options)
file_count = len(val_data.dataset)
if file_count % train_options.batch_size == 0:
steps_in_epoch = file_count // train_options.batch_size
else:
steps_in_epoch = file_count // train_options.batch_size + 1
losses_accu = {}
step = 0
for image, _ in val_data:
step += 1
image = image.to(device)
message = torch.Tensor(
np.random.choice(
[0, 1],
(image.shape[0], hidden_config.message_length))).to(device)
losses, (encoded_images, noised_images,
decoded_messages) = model.validate_on_batch(
[image, message])
if not losses_accu: # dict is empty, initialize
for name in losses:
losses_accu[name] = []
for name, loss in losses.items():
losses_accu[name].append(loss)
if step % print_each == 0:
print(f'Step {step}/{steps_in_epoch}')
utils.print_progress(losses_accu)
print('-' * 40)
utils.print_progress(losses_accu)
write_validation_loss(os.path.join(args.runs_root,
'validation_run.csv'),
losses_accu,
run_name,
checkpoint['epoch'],
write_header=write_csv_header)
write_csv_header = False
title = "No Attack"
title = f"Performance of {options.fit_fun} under {title}"
for k in stats.keys():
if filter_fn(k):
ax.plot(epochs, stats[k], label=k.replace('_', ' '))
plt.xlabel("Epochs")
plt.ylabel("Rate")
plt.title(title.title(), fontdict={'fontsize': 7})
plt.legend(loc=1, fontsize=5, framealpha=0.4)
plt.savefig(img_name, dpi=320, metadata={'comment': str(options)})
print(f"Done. Saved plot as {img_name}")
if __name__ == '__main__':
print("Calculating statistics and generating plots...")
options = utils.load_options()
stats = gen_stats(options)
img_name = "{}_{}_{:.1f}_{}_{}_{}{}".format(
options.dataset, options.fit_fun,
options.adversaries['percent_adv'] * 100, options.adversaries['type'],
'optimized' if options.adversaries['optimized'] else f"{tt}" if
(tt := options.adversaries['toggle_times']) else 'no_toggle',
f"{d}_delay" if (d := options.adversaries['delay']) is not None
and d > 0 else 'no_delay',
"_scaled" if options.adversaries['scale_up'] else '').replace(
' ', '_')
match_acc = lambda k: re.match('accuracy_\d', k)
save_plot(options, stats, lambda k: match_acc(k) is None,
f"{img_name}.png")
save_plot(options, stats, lambda k: match_acc(k) is not None,
f"{img_name}_accuracies.png")
def main():
device = torch.device(
'cuda') if torch.cuda.is_available() else torch.device('cpu')
parser = argparse.ArgumentParser(description='Training of HiDDeN nets')
parser.add_argument('--data-dir',
'-d',
required=True,
type=str,
help='The directory where the data is stored.')
parser.add_argument('--batch-size',
'-b',
required=True,
type=int,
help='The batch size.')
parser.add_argument('--epochs',
'-e',
default=400,
type=int,
help='Number of epochs to run the simulation.')
parser.add_argument('--name',
required=True,
type=str,
help='The name of the experiment.')
parser.add_argument(
'--runs-folder',
'-sf',
default=os.path.join('.', 'runs'),
type=str,
help='The root folder where data about experiments are stored.')
parser.add_argument(
'--size',
'-s',
default=128,
type=int,
help=
'The size of the images (images are square so this is height and width).'
)
parser.add_argument('--message',
'-m',
default=30,
type=int,
help='The length in bits of the watermark.')
parser.add_argument(
'--continue-from-folder',
'-c',
default='',
type=str,
help=
'The folder from where to continue a previous run. Leave blank if you are starting a new experiment.'
)
parser.add_argument(
'--tensorboard',
dest='tensorboard',
action='store_true',
help='If specified, use adds a Tensorboard log. On by default')
parser.add_argument('--no-tensorboard',
dest='tensorboard',
action='store_false',
help='Use to switch off Tensorboard logging.')
parser.add_argument(
'--noise',
nargs='*',
action=NoiseArgParser,
help=
"Noise layers configuration. Use quotes when specifying configuration, e.g. 'cropout((0.55, 0.6), (0.55, 0.6))'"
)
parser.set_defaults(tensorboard=True)
args = parser.parse_args()
checkpoint = None
if args.continue_from_folder != '':
this_run_folder = args.continue_from_folder
options_file = os.path.join(this_run_folder,
'options-and-config.pickle')
train_options, hidden_config, noise_config = utils.load_options(
options_file)
checkpoint = utils.load_last_checkpoint(
os.path.join(this_run_folder, 'checkpoints'))
train_options.start_epoch = checkpoint['epoch'] + 1
else:
start_epoch = 1
train_options = TrainingOptions(
batch_size=args.batch_size,
number_of_epochs=args.epochs,
train_folder=os.path.join(args.data_dir, 'train'),
validation_folder=os.path.join(args.data_dir, 'val'),
runs_folder=os.path.join('.', 'runs'),
start_epoch=start_epoch,
experiment_name=args.name)
noise_config = args.noise if args.noise is not None else []
hidden_config = HiDDenConfiguration(H=args.size,
W=args.size,
message_length=args.message,
encoder_blocks=4,
encoder_channels=64,
decoder_blocks=7,
decoder_channels=64,
use_discriminator=True,
use_vgg=False,
discriminator_blocks=3,
discriminator_channels=64,
decoder_loss=1,
encoder_loss=0.7,
adversarial_loss=1e-3)
this_run_folder = utils.create_folder_for_run(
train_options.runs_folder, args.name)
with open(os.path.join(this_run_folder, 'options-and-config.pickle'),
'wb+') as f:
pickle.dump(train_options, f)
pickle.dump(noise_config, f)
pickle.dump(hidden_config, f)
logging.basicConfig(level=logging.INFO,
format='%(message)s',
handlers=[
logging.FileHandler(
os.path.join(this_run_folder,
f'{args.name}.log')),
logging.StreamHandler(sys.stdout)
])
noiser = Noiser(noise_config, device)
if args.tensorboard:
logging.info('Tensorboard is enabled. Creating logger.')
from tensorboard_logger import TensorBoardLogger
tb_logger = TensorBoardLogger(os.path.join(this_run_folder, 'tb-logs'))
else:
tb_logger = None
model = Hidden(hidden_config, device, noiser, tb_logger)
if args.continue_from_folder != '':
# if we are continuing, we have to load the model params
assert checkpoint is not None
utils.model_from_checkpoint(model, checkpoint)
logging.info('HiDDeN model: {}\n'.format(model.to_stirng()))
logging.info('Model Configuration:\n')
logging.info(pprint.pformat(vars(hidden_config)))
logging.info('\nNoise configuration:\n')
logging.info(pprint.pformat(str(noise_config)))
logging.info('\nTraining train_options:\n')
logging.info(pprint.pformat(vars(train_options)))
train(model, device, hidden_config, train_options, this_run_folder,
tb_logger)
def main(options=None):
args = get_args()
if options is not None:
args = utils.load_options(args, options)
seed = 1 # Do NOT modify the seed. The captions have been generated from images generated from this seed.
torch.cuda.manual_seed_all(seed)
torch.manual_seed(seed)
np.random.seed(seed)
random.seed(seed)
# -------------------------------- INSTANTIATE MAIN ACTORS ----------------------------- #
# --------------- Create dataset ---------------- #
print('Creating dataset', flush=True)
image_normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
transform = transforms.Compose([
transforms.Resize(128), # Smaller edge will be matched to this number
transforms.CenterCrop((128, 128)),
transforms.ToTensor(),
image_normalize,
])
train_dataset = dataset.ImageAudioDataset(args.folder_dataset +
args.name_dataset,
split='train',
random_sampling=True,
transform=transform,
loading_image=args.loading_image)
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.workers,
pin_memory=True,
sampler=None)
val_dataset = dataset.ImageAudioDataset(args.folder_dataset +
args.name_dataset,
split='val',
transform=transform,
loading_image=args.loading_image)
val_loader = torch.utils.data.DataLoader(val_dataset,
batch_size=args.batch_size,
shuffle=False,
num_workers=args.workers,
pin_memory=True)
test_dataset = dataset.ImageAudioDataset(args.folder_dataset +
args.name_dataset,
split='test',
transform=transform,
loading_image=args.loading_image)
test_loader = torch.utils.data.DataLoader(test_dataset,
batch_size=args.batch_size,
shuffle=False,
num_workers=args.workers,
pin_memory=True)
# -------------- Create model --------------- #
print('Creating model', flush=True)
module = __import__('models')
model_class = getattr(module, args.model)
model = model_class(args)
model = torch.nn.DataParallel(model).cuda()
# Print model information
utils.print_model_report(model)
optimizer = torch.optim.SGD(model.parameters(),
args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
# Load model
resume_epoch = 0
if args.seed:
if args.seed == 'EXPDIR':
if args.name_checkpoint == '':
name = args.model + '_' + args.name_dataset
else:
name = args.name_checkpoint
path_load = args.expdir + 'model_best_' + name + '.pth.tar'
else:
path_load = args.seed
if args.resume:
utils.load_from_checkpoint(model,
path_load,
submodels_load=args.submodels_load,
optimizer=None)
checkpoint = torch.load(path_load)
resume_epoch = checkpoint['epoch']
else:
utils.load_from_checkpoint(model,
path_load,
submodels_load=args.submodels_load,
optimizer=None)
# --------------- Instantiate trainer --------------- #
print('Instantiating trainer', flush=True)
all_loaders = {
'val': val_loader,
'train': train_loader,
'test': test_loader
}
trainer = Trainer(model,
optimizer,
all_loaders,
args,
resume_epoch=resume_epoch)
# ------------------------- Others ----------------------- #
current_time = datetime.now().strftime('%b%d_%H-%M-%S')
log_dir = os.path.join(
args.results, 'runs',
args.name_checkpoint + '_' + current_time + '_' + socket.gethostname())
args.writer = SummaryWriter(log_dir=log_dir)
# ----------------------------------- TRAIN ------------------------------------------ #
if args.experiment:
print("Running experiment", flush=True)
experiments.experiment(args.experiment_name, trainer)
elif args.evaluate:
print("Performing evaluation epoch", flush=True)
trainer.eval()
elif args.generate_active_learning:
print("Generating active learning samples", flush=True)
active_learning.generate_active_learning(trainer)
else:
print("Beginning training", flush=True)
trainer.train()
def main():
if torch.cuda.is_available():
device = torch.device('cuda')
else:
device = torch.device('cpu')
parser = argparse.ArgumentParser(description='Test trained models')
parser.add_argument(
'--options-file',
'-o',
default='options-and-config.pickle',
type=str,
help='The file where the simulation options are stored.')
parser.add_argument('--checkpoint-file',
'-c',
required=True,
type=str,
help='Model checkpoint file')
parser.add_argument('--batch-size',
'-b',
default=12,
type=int,
help='The batch size.')
parser.add_argument('--source-image',
'-s',
required=True,
type=str,
help='The image to watermark')
parser.add_argument('--source-text',
'-t',
required=True,
type=str,
help='The text to watermark',
default='data/val_captions.txt')
parser.add_argument('--vocab-path',
'-v',
required=True,
type=str,
help='The path of vocabulary',
default='data/vocab.pkl')
# parser.add_argument('--times', '-t', default=10, type=int,
# help='Number iterations (insert watermark->extract).')
args = parser.parse_args()
train_options, hidden_config, noise_config = utils.load_options(
args.options_file)
noiser = Noiser(noise_config, device)
checkpoint = torch.load(args.checkpoint_file)
hidden_net = Hidden(hidden_config, device, noiser, None)
utils.model_from_checkpoint(hidden_net, checkpoint)
image_pil = Image.open(args.source_image)
image = randomCrop(np.array(image_pil), hidden_config.H, hidden_config.W)
image_tensor = TF.to_tensor(image).to(device)
image_tensor = image_tensor * 2 - 1 # transform from [0, 1] to [-1, 1]
images = torch.stack([image_tensor for _ in range(args.batch_size)], 0)
# for t in range(args.times):
with open(args.vocab_path, 'rb') as f:
vocab = pickle.load(f)
with open(args.source_text) as f:
captions = f.readlines()
captions = random.sample(captions, args.batch_size)
targets = []
for i, caption in enumerate(captions):
tokens = nltk.tokenize.word_tokenize(str(caption).lower())
caption = []
caption.append(vocab('<start>'))
caption.extend([vocab(token) for token in tokens])
caption.append(vocab('<end>'))
target = torch.Tensor(caption)
targets.append(target)
targets.sort(key=lambda x: len(x), reverse=True)
lengths = [len(cap) for cap in targets]
captions = torch.zeros(len(targets), max(lengths)).long()
for i, target in enumerate(targets):
end = lengths[i]
captions[i, :end] = target[:end]
captions = captions.to(device)
keys = np.random.permutation(512)
ekeys = torch.Tensor(np.eye(512)[keys])
dkeys = torch.Tensor(np.transpose(ekeys))
ekeys = torch.stack([ekeys for _ in range(args.batch_size)]).to(device)
dkeys = torch.stack([dkeys for _ in range(args.batch_size)]).to(device)
#print(f'sizes: images-{len(images)}, ekeys-{len(ekeys)}, dkeys-{len(dkeys)}, captions-{len(captions)}, lengths-{lengths}')
losses, (encoded_images, noised_images, decoded_messages, predicted_sents) = \
hidden_net.validate_on_batch([images, ekeys, dkeys, captions, lengths])
predicted_sents = predicted_sents.cpu().numpy()
for i in range(args.batch_size):
try:
print("predict : " + "".join(
[vocab.idx2word[int(idx)] + ' '
for idx in predicted_sents[i]]))
print("ground truth: " + "".join(
[vocab.idx2word[int(idx)] + ' ' for idx in captions[i]]))
except IndexError:
print(f'{i}th batch does not have enough length.')
noise_images = images - encoded_images
utils.save_images_with_noise(images.cpu(),
encoded_images.cpu(),
noise_images.cpu(),
'test_%d' % i,
'.',
resize_to=(256, 256))
def main():
device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
parent_parser = argparse.ArgumentParser(description='Training of HiDDeN nets')
subparsers = parent_parser.add_subparsers(dest='command', help='Sub-parser for commands')
new_run_parser = subparsers.add_parser('new', help='starts a new run')
new_run_parser.add_argument('--data-dir', '-d', required=True, type=str,
help='The directory where the data is stored.')
new_run_parser.add_argument('--batch-size', '-b', default=30, type=int, help='The batch size.')
new_run_parser.add_argument('--epochs', '-e', default=300, type=int, help='Number of epochs to run the simulation.')
new_run_parser.add_argument('--name', required=True, type=str, help='The name of the experiment.')
new_run_parser.add_argument('--size', '-s', default=128, type=int, help='The size of the images (images are square so this is height and width).')
new_run_parser.add_argument('--in_channels', default=3, type=int, help='input channel size')
new_run_parser.add_argument('--message', '-m', default=32, type=int, help='The length in bits of the watermark.')
new_run_parser.add_argument('--ratio', default=0.2, type=float, help='ratio of dataset.')
new_run_parser.add_argument('--continue-from-folder', '-c', default='', type=str,
help='The folder from where to continue a previous run. Leave blank if you are starting a new experiment.')
new_run_parser.add_argument('--enable-fp16', dest='enable_fp16', action='store_true',
help='Enable mixed-precision training.')
new_run_parser.add_argument('--noise', nargs='*', action=NoiseArgParser,
help="Noise layers configuration. Use quotes when specifying configuration, e.g. 'cropout((0.55, 0.6), (0.55, 0.6))'")
new_run_parser.set_defaults(enable_fp16=False)
continue_parser = subparsers.add_parser('continue', help='Continue a previous run')
continue_parser.add_argument('--folder', '-f', required=True, type=str,
help='Continue from the last checkpoint in this folder.')
continue_parser.add_argument('--data-dir', '-d', required=False, type=str,
help='The directory where the data is stored. Specify a value only if you want to override the previous value.')
continue_parser.add_argument('--epochs', '-e', required=False, type=int,
help='Number of epochs to run the simulation. Specify a value only if you want to override the previous value.')
args = parent_parser.parse_args()
checkpoint = None
loaded_checkpoint_file_name = None
if args.command == 'continue':
options_file = os.path.join(args.folder, 'options-and-config.pickle')
train_options, hidden_config, noise_config = utils.load_options(options_file)
checkpoint, loaded_checkpoint_file_name = utils.load_last_checkpoint(os.path.join(args.folder, 'checkpoints'))
train_options.start_epoch = checkpoint['epoch'] + 1
train_options.best_epoch = checkpoint['best_epoch']
train_options.best_cond = checkpoint['best_cond']
if args.epochs is not None:
if train_options.start_epoch < args.epochs:
train_options.number_of_epochs = args.epochs
else:
print(f'Command-line specifies of number of epochs = {args.epochs}, but folder={args.folder} '
f'already contains checkpoint for epoch = {train_options.start_epoch}.')
exit(1)
else:
assert args.command == 'new'
start_epoch = 1
train_options = TrainingOptions(
batch_size=args.batch_size,
number_of_epochs=args.epochs, data_ratio=args.ratio,
data_dir=args.data_dir,
runs_folder='./runs', tb_logger_folder='./logger',
start_epoch=start_epoch, experiment_name=f'{args.name}_r{int(100*args.ratio):03d}b{args.size}ch{args.in_channels}m{args.message}')
noise_config = args.noise if args.noise is not None else []
hidden_config = HiDDenConfiguration(H=args.size, W=args.size,input_channels=args.in_channels,
message_length=args.message,
encoder_blocks=4, encoder_channels=64,
decoder_blocks=7, decoder_channels=64,
use_discriminator=True,
use_vgg=False,
discriminator_blocks=3, discriminator_channels=64,
decoder_loss=1,
encoder_loss=0.7,
adversarial_loss=1e-3,
enable_fp16=args.enable_fp16
)
utils.create_folder_for_run(train_options)
with open(train_options.options_file, 'wb+') as f:
pickle.dump(train_options, f)
pickle.dump(noise_config, f)
pickle.dump(hidden_config, f)
logging.basicConfig(level=logging.INFO,
format='%(message)s',
handlers=[
logging.FileHandler(os.path.join(train_options.this_run_folder, f'{train_options.experiment_name}.log')),
logging.StreamHandler(sys.stdout)
])
logging.info(f'Tensorboard is enabled. Creating logger at {train_options.tb_logger_dir}')
tb_logger = TensorBoardLogger(train_options.tb_logger_dir)
noiser = Noiser(noise_config, device)
model = Hidden(hidden_config, device, noiser, tb_logger)
if args.command == 'continue':
# if we are continuing, we have to load the model params
assert checkpoint is not None
logging.info(f'Loading checkpoint from file {loaded_checkpoint_file_name}')
utils.model_from_checkpoint(model, checkpoint)
logging.info('HiDDeN model: {}\n'.format(model.to_stirng()))
logging.info('Model Configuration:\n')
logging.info(pprint.pformat(vars(hidden_config)))
logging.info('\nNoise configuration:\n')
logging.info(pprint.pformat(str(noise_config)))
logging.info('\nTraining train_options:\n')
logging.info(pprint.pformat(vars(train_options)))
train(model, device, hidden_config, train_options, train_options.this_run_folder, tb_logger)
def main():
# device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
device = torch.device(
'cuda') if torch.cuda.is_available() else torch.device('cpu')
parser = argparse.ArgumentParser(description='Training of HiDDeN nets')
parser.add_argument('--hostname',
default=socket.gethostname(),
help='the host name of the running server')
# parser.add_argument('--size', '-s', default=128, type=int, help='The size of the images (images are square so this is height and width).')
parser.add_argument('--data-dir',
'-d',
required=True,
type=str,
help='The directory where the data is stored.')
parser.add_argument(
'--runs_root',
'-r',
default=os.path.join('.', 'experiments'),
type=str,
help='The root folder where data about experiments are stored.')
parser.add_argument('--batch-size',
'-b',
default=1,
type=int,
help='Validation batch size.')
args = parser.parse_args()
if args.hostname == 'ee898-System-Product-Name':
args.data_dir = '/home/ee898/Desktop/chaoning/ImageNet'
args.hostname = 'ee898'
elif args.hostname == 'DL178':
args.data_dir = '/media/user/SSD1TB-2/ImageNet'
else:
args.data_dir = '/workspace/data_local/imagenet_pytorch'
assert args.data_dir
print_each = 25
completed_runs = [
o for o in os.listdir(args.runs_root)
if os.path.isdir(os.path.join(args.runs_root, o))
and o != 'no-noise-defaults'
]
print(completed_runs)
write_csv_header = True
current_run = args.runs_root
print(f'Run folder: {current_run}')
options_file = os.path.join(current_run, 'options-and-config.pickle')
train_options, hidden_config, noise_config = utils.load_options(
options_file)
train_options.train_folder = os.path.join(args.data_dir, 'val')
train_options.validation_folder = os.path.join(args.data_dir, 'val')
train_options.batch_size = args.batch_size
checkpoint, chpt_file_name = utils.load_last_checkpoint(
os.path.join(current_run, 'checkpoints'))
print(f'Loaded checkpoint from file {chpt_file_name}')
noiser = Noiser(noise_config, device, 'jpeg')
model = Hidden(hidden_config, device, noiser, tb_logger=None)
utils.model_from_checkpoint(model, checkpoint)
print('Model loaded successfully. Starting validation run...')
_, val_data = utils.get_data_loaders(hidden_config, train_options)
file_count = len(val_data.dataset)
if file_count % train_options.batch_size == 0:
steps_in_epoch = file_count // train_options.batch_size
else:
steps_in_epoch = file_count // train_options.batch_size + 1
with torch.no_grad():
noises = ['webp_10', 'webp_25', 'webp_50', 'webp_75', 'webp_90']
for noise in noises:
losses_accu = {}
step = 0
for image, _ in val_data:
step += 1
image = image.to(device)
message = torch.Tensor(
np.random.choice(
[0, 1], (image.shape[0],
hidden_config.message_length))).to(device)
losses, (
encoded_images, noised_images,
decoded_messages) = model.validate_on_batch_specific_noise(
[image, message], noise=noise)
if not losses_accu: # dict is empty, initialize
for name in losses:
losses_accu[name] = AverageMeter()
for name, loss in losses.items():
losses_accu[name].update(loss)
if step % print_each == 0 or step == steps_in_epoch:
print(f'Step {step}/{steps_in_epoch}')
utils.print_progress(losses_accu)
print('-' * 40)
# utils.print_progress(losses_accu)
write_validation_loss(os.path.join(args.runs_root,
'validation_run.csv'),
losses_accu,
noise,
checkpoint['epoch'],
write_header=write_csv_header)
write_csv_header = False
def main():
service_name = "CameraGateway"
log = Logger(service_name)
options = load_options()
camera = CameraGateway(fps=options["fps"])
publish_channel = Channel(options['broker_uri'])
rpc_channel = Channel(options['broker_uri'])
server = ServiceProvider(rpc_channel)
logging = LogInterceptor()
server.add_interceptor(logging)
server.delegate(topic=service_name + ".*.GetConfig",
request_type=FieldSelector,
reply_type=CameraConfig,
function=camera.get_config)
server.delegate(topic=service_name + ".*.SetConfig",
request_type=CameraConfig,
reply_type=Empty,
function=camera.set_config)
exporter = create_exporter(service_name=service_name,
uri=options["zipkin_uri"])
while True:
# iterate through videos listed
for video in options['videos']:
# id of the first sequence of videos
person_id = video['person_id']
gesture_id = video['gesture_id']
# getting the path of the 4 videos
video_files = {
cam_id: os.path.join(
options['folder'],
'p{:03d}g{:02d}c{:02d}.mp4'.format(person_id, gesture_id,
cam_id))
for cam_id in options["cameras_id"]
}
for iteration in range(video['iterations']):
info = {
"person": person_id,
"gesture": gesture_id,
"iteration": iteration
}
log.info('{}', str(info).replace("'", '"'))
# object that let get images from multiples videos files
video_loader = FramesLoader(video_files)
# iterate through all samples on video
while True:
time_initial = time.time()
# listen server for messages about change
try:
message = rpc_channel.consume(timeout=0)
if server.should_serve(message):
server.serve(message)
except socket.timeout:
pass
frame_id, frames = video_loader.read()
for cam in sorted(frames.keys()):
tracer = Tracer(exporter)
span = tracer.start_span(name='frame')
pb_image = to_pb_image(frames[cam])
msg = Message(content=pb_image)
msg.inject_tracing(span)
topic = 'CameraGateway.{}.Frame'.format(cam)
publish_channel.publish(msg, topic=topic)
tracer.end_span()
took_ms = (time.time() - time_initial) * 1000
dt = (1 / camera.fps) - (took_ms / 1000)
if dt > 0:
time.sleep(dt)
info = {
"sample": frame_id,
"took_ms": took_ms,
"wait_ms": dt * 1000
}
log.info('{}', str(info).replace("'", '"'))
if frame_id >= (video_loader.num_samples - 1):
video_loader.release()
del video_loader
gc.collect()
break
if options['loop'] is False:
break
def main():
device = torch.device(
'cuda') if torch.cuda.is_available() else torch.device('cpu')
parent_parser = argparse.ArgumentParser(
description='Training of HiDDeN nets')
subparsers = parent_parser.add_subparsers(dest='command',
help='Sub-parser for commands')
new_run_parser = subparsers.add_parser('new', help='starts a new run')
new_run_parser.add_argument('--data-dir',
'-d',
required=True,
type=str,
help='The directory where the data is stored.')
new_run_parser.add_argument('--batch-size',
'-b',
required=True,
type=int,
help='The batch size.')
new_run_parser.add_argument('--epochs',
'-e',
default=300,
type=int,
help='Number of epochs to run the simulation.')
new_run_parser.add_argument('--name',
required=True,
type=str,
help='The name of the experiment.')
new_run_parser.add_argument('--adv_loss',
default=0,
required=False,
type=float,
help='Coefficient of the adversarial loss.')
new_run_parser.add_argument('--residual',
default=0,
required=False,
type=int,
help='If to use residual or not.')
new_run_parser.add_argument('--video_dataset',
default=0,
required=False,
type=int,
help='If to use video dataset or not.')
new_run_parser.add_argument(
'--save-dir',
'-sd',
default='runs',
required=True,
type=str,
help='The save directory where the result is stored.')
new_run_parser.add_argument(
'--size',
'-s',
default=128,
type=int,
help=
'The size of the images (images are square so this is height and width).'
)
new_run_parser.add_argument('--message',
'-m',
default=30,
type=int,
help='The length in bits of the watermark.')
new_run_parser.add_argument(
'--continue-from-folder',
'-c',
default='',
type=str,
help=
'The folder from where to continue a previous run. Leave blank if you are starting a new experiment.'
)
# parser.add_argument('--tensorboard', dest='tensorboard', action='store_true',
# help='If specified, use adds a Tensorboard log. On by default')
new_run_parser.add_argument('--tensorboard',
action='store_true',
help='Use to switch on Tensorboard logging.')
new_run_parser.add_argument('--enable-fp16',
dest='enable_fp16',
action='store_true',
help='Enable mixed-precision training.')
new_run_parser.add_argument(
'--noise',
nargs='*',
action=NoiseArgParser,
help=
"Noise layers configuration. Use quotes when specifying configuration, e.g. 'cropout((0.55, 0.6), (0.55, 0.6))'"
)
new_run_parser.add_argument('--hostname',
default=socket.gethostname(),
help='the host name of the running server')
new_run_parser.add_argument(
'--cover-dependent',
default=1,
required=False,
type=int,
help='If to use cover dependent architecture or not.')
new_run_parser.add_argument('--jpeg_type',
'-j',
required=False,
type=str,
default='jpeg',
help='Jpeg type used in the combined2 noise.')
new_run_parser.set_defaults(tensorboard=False)
new_run_parser.set_defaults(enable_fp16=False)
continue_parser = subparsers.add_parser('continue',
help='Continue a previous run')
continue_parser.add_argument(
'--folder',
'-f',
required=True,
type=str,
help='Continue from the last checkpoint in this folder.')
continue_parser.add_argument(
'--data-dir',
'-d',
required=False,
type=str,
help=
'The directory where the data is stored. Specify a value only if you want to override the previous value.'
)
continue_parser.add_argument(
'--epochs',
'-e',
required=False,
type=int,
help=
'Number of epochs to run the simulation. Specify a value only if you want to override the previous value.'
)
# continue_parser.add_argument('--tensorboard', action='store_true',
# help='Override the previous setting regarding tensorboard logging.')
# Setting up a seed for debug
seed = 123
torch.manual_seed(seed)
np.random.seed(seed)
args = parent_parser.parse_args()
checkpoint = None
loaded_checkpoint_file_name = None
print(args.cover_dependent)
if not args.video_dataset:
if args.hostname == 'ee898-System-Product-Name':
args.data_dir = '/home/ee898/Desktop/chaoning/ImageNet'
args.hostname = 'ee898'
elif args.hostname == 'DL178':
args.data_dir = '/media/user/SSD1TB-2/ImageNet'
else:
args.data_dir = '/workspace/data_local/imagenet_pytorch'
else:
if args.hostname == 'ee898-System-Product-Name':
args.data_dir = '/home/ee898/Desktop/chaoning/ImageNet'
args.hostname = 'ee898'
elif args.hostname == 'DL178':
args.data_dir = '/media/user/SSD1TB-2/ImageNet'
else:
args.data_dir = './oops_dataset/oops_video'
assert args.data_dir
if args.command == 'continue':
this_run_folder = args.folder
options_file = os.path.join(this_run_folder,
'options-and-config.pickle')
train_options, hidden_config, noise_config = utils.load_options(
options_file)
checkpoint, loaded_checkpoint_file_name = utils.load_last_checkpoint(
os.path.join(this_run_folder, 'checkpoints'))
train_options.start_epoch = checkpoint['epoch'] + 1
if args.data_dir is not None:
train_options.train_folder = os.path.join(args.data_dir, 'train')
train_options.validation_folder = os.path.join(
args.data_dir, 'val')
if args.epochs is not None:
if train_options.start_epoch < args.epochs:
train_options.number_of_epochs = args.epochs
else:
print(
f'Command-line specifies of number of epochs = {args.epochs}, but folder={args.folder} '
f'already contains checkpoint for epoch = {train_options.start_epoch}.'
)
exit(1)
else:
assert args.command == 'new'
start_epoch = 1
train_options = TrainingOptions(
batch_size=args.batch_size,
number_of_epochs=args.epochs,
train_folder=os.path.join(args.data_dir, 'train'),
validation_folder=os.path.join(args.data_dir, 'val'),
runs_folder=os.path.join('.', args.save_dir),
start_epoch=start_epoch,
experiment_name=args.name,
video_dataset=args.video_dataset)
noise_config = args.noise if args.noise is not None else []
hidden_config = HiDDenConfiguration(
H=args.size,
W=args.size,
message_length=args.message,
encoder_blocks=4,
encoder_channels=64,
decoder_blocks=7,
decoder_channels=64,
use_discriminator=True,
use_vgg=False,
discriminator_blocks=3,
discriminator_channels=64,
decoder_loss=1,
encoder_loss=0.7,
adversarial_loss=args.adv_loss,
cover_dependent=args.cover_dependent,
residual=args.residual,
enable_fp16=args.enable_fp16)
this_run_folder = utils.create_folder_for_run(
train_options.runs_folder, args.name)
with open(os.path.join(this_run_folder, 'options-and-config.pickle'),
'wb+') as f:
pickle.dump(train_options, f)
pickle.dump(noise_config, f)
pickle.dump(hidden_config, f)
logging.basicConfig(level=logging.INFO,
format='%(message)s',
handlers=[
logging.FileHandler(
os.path.join(
this_run_folder,
f'{train_options.experiment_name}.log')),
logging.StreamHandler(sys.stdout)
])
if (args.command == 'new' and args.tensorboard) or \
(args.command == 'continue' and os.path.isdir(os.path.join(this_run_folder, 'tb-logs'))):
logging.info('Tensorboard is enabled. Creating logger.')
from tensorboard_logger import TensorBoardLogger
tb_logger = TensorBoardLogger(os.path.join(this_run_folder, 'tb-logs'))
else:
tb_logger = None
noiser = Noiser(noise_config, device, args.jpeg_type)
model = Hidden(hidden_config, device, noiser, tb_logger)
if args.command == 'continue':
# if we are continuing, we have to load the model params
assert checkpoint is not None
logging.info(
f'Loading checkpoint from file {loaded_checkpoint_file_name}')
utils.model_from_checkpoint(model, checkpoint)
logging.info('HiDDeN model: {}\n'.format(model.to_stirng()))
logging.info('Model Configuration:\n')
logging.info(pprint.pformat(vars(hidden_config)))
logging.info('\nNoise configuration:\n')
logging.info(pprint.pformat(str(noise_config)))
logging.info('\nTraining train_options:\n')
logging.info(pprint.pformat(vars(train_options)))
# train(model, device, hidden_config, train_options, this_run_folder, tb_logger)
# train_other_noises(model, device, hidden_config, train_options, this_run_folder, tb_logger)
if str(args.noise[0]) == "WebP()":
noise = 'webp'
elif str(args.noise[0]) == "JpegCompression2000()":
noise = 'jpeg2000'
elif str(args.noise[0]) == "MPEG4()":
noise = 'mpeg4'
elif str(args.noise[0]) == "H264()":
noise = 'h264'
elif str(args.noise[0]) == "XVID()":
noise = 'xvid'
elif str(args.noise[0]) == "DiffQFJpegCompression2()":
noise = 'diff_qf_jpeg2'
elif str(args.noise[0]) == "DiffCorruptions()":
noise = 'diff_corruptions'
else:
noise = 'jpeg'
train_own_noise(model, device, hidden_config, train_options,
this_run_folder, tb_logger, noise)
from is_wire.core import Channel, Message, Subscription, Logger
from is_wire.core import Tracer, ZipkinExporter, BackgroundThreadTransport
from is_wire.core.utils import now
from utils import load_options
from heatmap import SkeletonsHeatmap
from builtins import super
class MyChannel(Channel):
def consume_until(self, deadline):
timeout = max([deadline - now(), 0.0])
return super().consume(timeout=timeout)
service_name = 'Skeletons.Heatmap'
log = Logger(name=service_name)
ops = load_options()
channel = MyChannel(ops.broker_uri)
subscription = Subscription(channel)
exporter = ZipkinExporter(
service_name=service_name,
host_name=ops.zipkin_host,
port=ops.zipkin_port,
transport=BackgroundThreadTransport(max_batch_size=20),
)
subscription.subscribe('Skeletons.Localization')
sks_hm = SkeletonsHeatmap(ops)
period = ops.period_ms / 1000.0
def main():
if torch.cuda.is_available():
device = torch.device('cuda')
else:
device = torch.device('cpu')
parser = argparse.ArgumentParser(description='Test trained models')
parser.add_argument(
'--options-file',
'-o',
default='options-and-config.pickle',
type=str,
help='The file where the simulation options are stored.')
parser.add_argument('--checkpoint-file',
'-c',
required=True,
type=str,
help='Model checkpoint file')
parser.add_argument('--batch-size',
'-b',
default=12,
type=int,
help='The batch size.')
parser.add_argument('--source-image',
'-s',
required=True,
type=str,
help='The image to watermark')
# parser.add_argument('--times', '-t', default=10, type=int,
# help='Number iterations (insert watermark->extract).')
args = parser.parse_args()
train_options, hidden_config, noise_config = utils.load_options(
args.options_file)
noiser = Noiser(noise_config, device)
checkpoint = torch.load(args.checkpoint_file)
hidden_net = Hidden(hidden_config, device, noiser, None)
utils.model_from_checkpoint(hidden_net, checkpoint)
image_pil = Image.open(args.source_image)
image = randomCrop(np.array(image_pil), hidden_config.H, hidden_config.W)
image_tensor = TF.to_tensor(image).to(device)
image_tensor = image_tensor * 2 - 1 # transform from [0, 1] to [-1, 1]
image_tensor.unsqueeze_(0)
# for t in range(args.times):
message = torch.Tensor(
np.random.choice(
[0, 1],
(image_tensor.shape[0], hidden_config.message_length))).to(device)
losses, (encoded_images, noised_images,
decoded_messages) = hidden_net.validate_on_batch(
[image_tensor, message])
decoded_rounded = decoded_messages.detach().cpu().numpy().round().clip(
0, 1)
message_detached = message.detach().cpu().numpy()
print('original: {}'.format(message_detached))
print('decoded : {}'.format(decoded_rounded))
print('error : {:.3f}'.format(
np.mean(np.abs(decoded_rounded - message_detached))))
utils.save_images(image_tensor.cpu(),
encoded_images.cpu(),
'test',
'.',
resize_to=(256, 256))
