| import numpy as np | |
| def add_mask(image, label, vis_labels, colors, alpha=0.5): | |
| if len(image.shape) < 3 or image.shape[2] == 1: | |
| image = image.repeat(3).reshape((image.shape[0], image.shape[1], 3)).astype(np.uint8) | |
| ori_image = image.copy() | |
| for ci, vis_label in enumerate(vis_labels): | |
| color = colors[ci] | |
| mask = label == vis_label | |
| for i in range(3): | |
| image[:, :, i][mask] = (color[i] * alpha + image[:, :, i][mask] * (1 - alpha)).astype(np.uint8) | |
| return ori_image, image | |
| def find_haight(mask): | |
| v_sum = (np.sum(mask, axis=1) > 0).astype(np.uint8) | |
| v_diff = np.diff(np.hstack((0, v_sum, 0))) | |
| v_min_y = np.where(v_diff > 0)[0][0] | |
| v_max_y = np.where(v_diff < 0)[0][-1] - 1 | |
| return v_max_y - v_min_y | |
| def simple_vcdr(mask): | |
| disc_mask = (mask > 0).astype(np.uint8) | |
| disc_height = find_haight(disc_mask) | |
| cup_mask = (mask > 1).astype(np.uint8) | |
| cup_height = find_haight(cup_mask) | |
| vcdr = cup_height / disc_height | |
| return vcdr | |