|
|
import numpy as np |
|
|
import cv2 |
|
|
import matplotlib.pyplot as plt |
|
|
from mpl_toolkits.axes_grid1 import make_axes_locatable |
|
|
|
|
|
def getDenseMask(landmarks, h, w): |
|
|
RL, LL, H = landmarks[:44], landmarks[44:94], landmarks[94:] |
|
|
img = np.zeros([h, w], dtype='uint8') |
|
|
RL = RL.reshape(-1, 1, 2).astype('int') |
|
|
LL = LL.reshape(-1, 1, 2).astype('int') |
|
|
H = H.reshape(-1, 1, 2).astype('int') |
|
|
img = cv2.drawContours(img, [RL], -1, 1, -1) |
|
|
img = cv2.drawContours(img, [LL], -1, 1, -1) |
|
|
img = cv2.drawContours(img, [H], -1, 2, -1) |
|
|
return img |
|
|
|
|
|
def drawOnTop(img, landmarks, original_shape): |
|
|
h, w = original_shape |
|
|
output = getDenseMask(landmarks, h, w) |
|
|
image = np.zeros([h,w,3]) |
|
|
image[:,:,0] = img + 0.3*(output==1).astype('float') - 0.1*(output==2).astype('float') |
|
|
image[:,:,1] = img + 0.3*(output==2).astype('float') - 0.1*(output==1).astype('float') |
|
|
image[:,:,2] = img - 0.1*(output==1).astype('float') - 0.2*(output==2).astype('float') |
|
|
image = np.clip(image,0,1) |
|
|
RL, LL, H = landmarks[:44], landmarks[44:94], landmarks[94:] |
|
|
for l in RL: image = cv2.circle(image,(int(l[0]),int(l[1])),5,(1,0,1),-1) |
|
|
for l in LL: image = cv2.circle(image,(int(l[0]),int(l[1])),5,(1,0,1),-1) |
|
|
for l in H: image = cv2.circle(image,(int(l[0]),int(l[1])),5,(1,1,0),-1) |
|
|
return image |
|
|
|
|
|
def create_overlay(img, landmarks): |
|
|
h, w = img.shape[:2] |
|
|
dense_mask = getDenseMask(landmarks, h, w) |
|
|
overlay = np.zeros([h, w, 3]) |
|
|
|
|
|
overlay[:,:,0] = img + 0.3 * (dense_mask == 1).astype('float') - 0.1 * (dense_mask == 2).astype('float') |
|
|
overlay[:,:,1] = img + 0.3 * (dense_mask == 2).astype('float') - 0.1 * (dense_mask == 1).astype('float') |
|
|
overlay[:,:,2] = img - 0.1 * (dense_mask == 1).astype('float') - 0.2 * (dense_mask == 2).astype('float') |
|
|
overlay = np.clip(overlay, 0, 1) |
|
|
|
|
|
return overlay |
|
|
|
|
|
def plot_side_by_side_comparison(img_orig, means_orig, uncertainty_orig, img_corr, means_corr, uncertainty_corr): |
|
|
|
|
|
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 7)) |
|
|
|
|
|
fig.set_constrained_layout(True) |
|
|
|
|
|
vmax = max(np.max(np.mean(uncertainty_orig, axis=1)), np.max(np.mean(uncertainty_corr, axis=1))) |
|
|
|
|
|
|
|
|
overlay_orig = create_overlay(img_orig, means_orig) |
|
|
ax1.imshow(overlay_orig) |
|
|
scatter1 = ax1.scatter( |
|
|
means_orig[:, 0], means_orig[:, 1], |
|
|
c=np.mean(uncertainty_orig, axis=1), |
|
|
cmap='hot', s=50, vmin=0, vmax=vmax |
|
|
) |
|
|
ax1.set_title("Original", fontsize=16, pad=10) |
|
|
ax1.axis('off') |
|
|
|
|
|
|
|
|
overlay_corr = create_overlay(img_corr, means_corr) |
|
|
ax2.imshow(overlay_corr) |
|
|
scatter2 = ax2.scatter( |
|
|
means_corr[:, 0], means_corr[:, 1], |
|
|
c=np.mean(uncertainty_corr, axis=1), |
|
|
cmap='hot', s=50, vmin=0, vmax=vmax |
|
|
) |
|
|
ax2.set_title("Corrupted", fontsize=16, pad=10) |
|
|
ax2.axis('off') |
|
|
|
|
|
|
|
|
cbar = fig.colorbar(scatter2, ax=[ax1, ax2], fraction=0.046, pad=0.01, shrink=0.85) |
|
|
cbar.ax.tick_params(labelsize=10) |
|
|
|
|
|
return fig |
