File size: 3,032 Bytes
50f870e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 |
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)))
# --- Original ---
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')
# --- Corrupted ---
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')
# Shared colorbar
cbar = fig.colorbar(scatter2, ax=[ax1, ax2], fraction=0.046, pad=0.01, shrink=0.85)
cbar.ax.tick_params(labelsize=10)
return fig |