def main():
"""
Runs pytorch-SalGAN on a sample images
"""
# create output file
if not os.path.exists(OUTPUT_PATH):
os.makedirs(OUTPUT_PATH)
# init model with pre-trained weights
model = create_model()
model.load_state_dict(torch.load(PATH_PYTORCH_WEIGHTS)['state_dict'])
model.eval()
# if GPU is enabled
if USE_GPU:
model.cuda()
if not os.path.exists(OUTPUT_PATH):
os.makedirs(OUTPUT_PATH)
# load and preprocess images in folder
for img in os.listdir(INPUT_PATH):
if 'val' in img:
filename = os.path.join(INPUT_PATH, img)
image_tensor, image_size = load_image(filename)
if USE_GPU:
image_tensor = image_tensor.cuda()
# run model inference
prediction = model.forward(
image_tensor[None, ...]) # add extra batch dimension
# get result to cpu and squeeze dimensions
if USE_GPU:
prediction = prediction.squeeze().data.cpu().numpy()
else:
prediction = prediction.squeeze().data.numpy()
# postprocess
# first normalize [0,1]
prediction = normalize_map(prediction)
saliency = postprocess_prediction(prediction, image_size)
saliency = normalize_map(saliency)
saliency *= 255
saliency = saliency.astype(np.uint8)
# save saliency
cv2.imwrite(os.path.join(OUTPUT_PATH, img), saliency)
print("Processed image {} ".format(img), end="\r")
sys.stdout.flush()
print("\n")
# Dataloaders
dataloader = {
TRAIN: DataLoader(ds_train, batch_size=batch_size,
shuffle=True, num_workers=2),
VAL: DataLoader(ds_validate, batch_size=batch_size,
shuffle=False, num_workers=2)
}
# model ====================================================================
print("Init model...")
# init model with pre-trained weights
model = create_model()
model.load_state_dict(torch.load('../trained_models/baseline_weights/gen_model.pt'))
model.train()
model.cuda()
cudnn.benchmark = True
# loss =====================================================================
print("BCE criterium...")
bce_loss = BCELoss()
# select only decoder parameters, keep vgg16 with pretrained weights
decoder_parameters = []
for i, (a, p) in enumerate(model.named_parameters()):
if i>25:
print(i, a, p.shape)
decoder_parameters.append(p)
sAUC = AUC_shuffled(saliency_map, fground_truth, other_map)
nss = NSS(saliency_map, fground_truth)
cc = CC(saliency_map, mground_truth)
sim = SIM(saliency_map, mground_truth)
return (AUC_judd, sAUC, nss, cc, sim)
assert (MODEL is not None), "Introduce model in args [--model]"
# create output file
if SAVE:
if not os.path.exists(OUTPUT_PATH):
os.makedirs(OUTPUT_PATH)
# init model with pre-trained weights
if DEPTH and COORD:
model = create_model(6)
elif DEPTH:
model = create_model(4)
elif COORD:
model = create_model(5)
else:
model = create_model(3)
model = add_bn(model)
model.load_state_dict(torch.load(PATH_PYTORCH_WEIGHTS)['state_dict'])
model.eval()
# if GPU is enabled
if USE_GPU:
model.cuda()
start = datetime.datetime.now().replace(microsecond=0)
dataloader = {
TRAIN: DataLoader(ds_train, batch_size=batch_size,
shuffle=True, num_workers=2),
VAL: DataLoader(ds_validate, batch_size=batch_size,
shuffle=False, num_workers=2)
}
torch.cuda.set_device(1)
# model ====================================================================
print("Init model...")
# init model with pre-trained weights
vgg_weights = torch.load('../trained_models/salgan_baseline.pt')['state_dict']
if DEPTH and COORD:
model = create_model(6)
for i in range(0,3):
vgg_weights = add_layer_weights(vgg_weights)
elif DEPTH:
model = create_model(4)
add_layer_weights(vgg_weights)
elif COORD:
model = create_model(5)
for i in range(0,2):
vgg_weights = add_layer_weights(vgg_weights)
else: model = create_model(3)
model.load_state_dict(vgg_weights)
# Add batch normalization to current model
model = add_bn(model)
TRAIN:
DataLoader(ds_train,
batch_size=batch_size,
shuffle=True,
num_workers=4),
VAL:
DataLoader(ds_validate,
batch_size=batch_size,
shuffle=False,
num_workers=4)
}
# model ====================================================================
print("Init model...")
# Init generator model with pretrained weights from SALICON's
modelG = create_model()
modelG.load_state_dict(
torch.load('../trained_models/salgan_salicon_3epochs/models/best.pt')
['state_dict'])
modelG.train()
modelG.cuda()
# Init discriminator model with weights ????????
modelD = Discriminator()
#modelD.load_state_dict(torch.load(????)['state_dict'])
modelD.train()
modelD.cuda()
cudnn.benchmark = True
# loss =====================================================================
print("BCE criterium...")
