|
|
"""
|
|
|
Vis utilities. Code adapted from LOST: https://github.com/valeoai/LOST
|
|
|
"""
|
|
|
import cv2
|
|
|
import torch
|
|
|
import skimage.io
|
|
|
import numpy as np
|
|
|
import torch.nn as nn
|
|
|
from PIL import Image
|
|
|
import scipy
|
|
|
|
|
|
import matplotlib.pyplot as plt
|
|
|
|
|
|
def visualize_img(image, vis_folder, im_name):
|
|
|
pltname = f"{vis_folder}/{im_name}"
|
|
|
Image.fromarray(image).save(pltname)
|
|
|
print(f"Original image saved at {pltname}.")
|
|
|
|
|
|
def visualize_predictions(img, pred, vis_folder, im_name, save=True):
|
|
|
"""
|
|
|
Visualization of the predicted box and the corresponding seed patch.
|
|
|
"""
|
|
|
image = np.copy(img)
|
|
|
|
|
|
cv2.rectangle(
|
|
|
image,
|
|
|
(int(pred[0]), int(pred[1])),
|
|
|
(int(pred[2]), int(pred[3])),
|
|
|
(255, 0, 0), 3,
|
|
|
)
|
|
|
if save:
|
|
|
pltname = f"{vis_folder}/{im_name}_TokenCut_pred.jpg"
|
|
|
Image.fromarray(image).save(pltname)
|
|
|
print(f"Predictions saved at {pltname}.")
|
|
|
return image
|
|
|
|
|
|
def visualize_predictions_gt(img, pred, gt, vis_folder, im_name, dim, scales, save=True):
|
|
|
"""
|
|
|
Visualization of the predicted box and the corresponding seed patch.
|
|
|
"""
|
|
|
image = np.copy(img)
|
|
|
|
|
|
cv2.rectangle(
|
|
|
image,
|
|
|
(int(pred[0]), int(pred[1])),
|
|
|
(int(pred[2]), int(pred[3])),
|
|
|
(255, 0, 0), 3,
|
|
|
)
|
|
|
|
|
|
if len(gt>1):
|
|
|
for i in range(len(gt)):
|
|
|
cv2.rectangle(
|
|
|
image,
|
|
|
(int(gt[i][0]), int(gt[i][1])),
|
|
|
(int(gt[i][2]), int(gt[i][3])),
|
|
|
(0, 0, 255), 3,
|
|
|
)
|
|
|
if save:
|
|
|
pltname = f"{vis_folder}/{im_name}_TokenCut_BBOX.jpg"
|
|
|
Image.fromarray(image).save(pltname)
|
|
|
|
|
|
return image
|
|
|
|
|
|
def visualize_eigvec(eigvec, vis_folder, im_name, dim, scales, save=True):
|
|
|
"""
|
|
|
Visualization of the second smallest eigvector
|
|
|
"""
|
|
|
eigvec = scipy.ndimage.zoom(eigvec, scales, order=0, mode='nearest')
|
|
|
if save:
|
|
|
pltname = f"{vis_folder}/{im_name}_TokenCut_attn.jpg"
|
|
|
plt.imsave(fname=pltname, arr=eigvec, cmap='cividis')
|
|
|
print(f"Eigen attention saved at {pltname}.")
|
|
|
|
