def CalculateMetrics(fground_truth, mground_truth, predicted_map):
# import ipdb; ipdb.set_trace()
predicted_map = normalize_map(predicted_map)
predicted_map = postprocess_prediction(
predicted_map, (predicted_map.shape[0], predicted_map.shape[1]))
predicted_map = normalize_map(predicted_map)
predicted_map *= 255
fground_truth = cv2.resize(fground_truth, (0, 0), fx=0.5, fy=0.5)
predicted_map = cv2.resize(predicted_map, (0, 0), fx=0.5, fy=0.5)
mground_truth = cv2.resize(mground_truth, (0, 0), fx=0.5, fy=0.5)
fground_truth = fground_truth.astype(np.float32) / 255
predicted_map = predicted_map.astype(np.float32)
mground_truth = mground_truth.astype(np.float32)
AUC_judd_answer = AUC_Judd(predicted_map, fground_truth)
AUC_Borji_answer = AUC_Borji(predicted_map, fground_truth)
nss_answer = NSS(predicted_map, fground_truth)
cc_answer = CC(predicted_map, mground_truth)
sim_answer = SIM(predicted_map, mground_truth)
return AUC_judd_answer, AUC_Borji_answer, nss_answer, cc_answer, sim_answer
def inference(y):
predictions = []
files = sorted(os.listdir(os.path.join(INPUT_PATH, '{:03d}'.format(y))),
key=lambda x: int(x.split(".")[0]))
for file in files:
filename = os.path.join(INPUT_PATH, '{:03d}'.format(y), file)
image_tensor, image_size = load_image(filename)
if DEPTH:
num_image = int(file.split('.')[0])
if num_image < 10:
ima_name = '0010.png'
else:
ima_name = '{:04d}.png'.format(int(num_image / 10) * 10)
depth = cv2.imread(os.path.join(DEPTH_PATH, str(y), ima_name), 0)
depth = cv2.resize(depth, (256, 192), interpolation=cv2.INTER_AREA)
depth = depth.astype(np.float32)
depth = np.expand_dims(depth, axis=0)
depth = torch.FloatTensor(depth)
image_tensor = torch.cat([image_tensor, depth], 0)
if COORD:
c1 = np.empty(shape=(192, 256))
c2 = np.empty(shape=(192, 256))
for i in range(0, 256):
c1[:, i] = np.linspace(-1, 1, 192)
for i in range(0, 192):
c2[i, :] = np.linspace(-1, 1, 256)
c1 = c1.astype(np.float32)
c2 = c2.astype(np.float32)
c1 = np.expand_dims(c1, axis=0)
c2 = np.expand_dims(c2, axis=0)
c1 = torch.FloatTensor(c1)
c2 = torch.FloatTensor(c2)
image_tensor = torch.cat([image_tensor, c1], 0)
image_tensor = torch.cat([image_tensor, c2], 0)
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
predictions.append(saliency)
return predictions
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")
# convert to torch Tensor
image = torch.FloatTensor(image)
image2 = torch.FloatTensor(image2)
# swap channel dimensions
image = image.permute(2, 0, 1)
image2 = image2.permute(2, 0, 1)
if USE_GPU:
image = image.cuda()
image2 = image2.cuda()
for img, name in zip([image, image2], ["nocrop.png", "crop.png"]):
# run model inference
prediction = model.forward(img[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, (320, 640))
saliency = normalize_map(saliency)
saliency *= 255
cv2.imwrite(os.path.join("/home/saliency_maps/", name), saliency)