Spaces:
Build error
Build error
| 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 |