def compute_refined_masks(data: ImageData):
dir_results = data.output_directory
if not cfg.visualization_dict['no_result_folder']:
clean_directory(dir_results)
refined_masks = []
all_lines = dict()
if cfg.visualization_dict['edges']:
cv2.imwrite(os.path.join(dir_results, "edges.jpg"), data.edge_img)
if cfg.visualization_dict['initial_mask']:
image_masked = draw_masks_on_image(data.edge_img, data.planercnn_masks)
cv2.imwrite(os.path.join(dir_results, "masked_image.jpg"),
image_masked)
for i, mask in enumerate(data.planercnn_masks):
logger.info(f"{'#' * 35} Running on plane {i} {'#' * 35}")
mask, lines = compute_refined_mask_and_lines(data.img, data.edge_img,
mask, dir_results, i)
if mask is not None:
refined_masks.append(mask)
if lines is not None:
all_lines[i] = lines
if cfg.visualization_dict['refined_mask']:
image_refined = draw_masks_on_image(
get_colored_image(data.edge_img.copy()), refined_masks)
cv2.imwrite(os.path.join(dir_results, "refined_masked_image.jpg"),
image_refined)
if cfg.visualization_dict['lines']:
all_lines_image = draw_all_lines_on_image(data.img, all_lines)
cv2.imwrite(
os.path.join(
dir_results, "image_all_lines_{}.png".format(
time.strftime("%Y%m%d_%H%M%S", time.localtime()))),
all_lines_image)
return refined_masks
def get_edge_candidate_clusters_from_mask(image, mask, n_mask, ksize,
output_directory):
from refiner.clustering.dbscan import dbscan_with_masked_image
from refiner.image_processing.draw import draw_masks_on_image
if ksize % 2 == 0:
raise ValueError("Kernel size must be odd")
img_masked = draw_masks_on_image(image, [(mask * 255).astype(np.uint8)])
if cfg.visualization_dict['widened_contour']:
cv2.imwrite(
os.path.join(
output_directory,
"mask_{}_around_edges.jpg".format(str(n_mask).zfill(2))),
img_masked)
image[mask == 0] = 0
image = get_grayscaled_image(image)
clustered_edges = dbscan_with_masked_image(
image, eps=cfg.clustering_eps, min_samples=cfg.clustering_min_sample)
if len(clustered_edges.values()) > 50:
logger.debug("First run not successful (found {} edges)".format(
len(clustered_edges.values())))
clustered_edges = dbscan_with_masked_image(
image,
eps=cfg.clustering_eps * 2,
min_samples=cfg.clustering_min_sample)
clustered_edges = filter_edges(image, clustered_edges)
return clustered_edges
def save_image_with_refined_mask(img_edges, mask, points, i, dir_results):
image_refined_curr = draw_masks_on_image(
get_colored_image(img_edges.copy()), [mask])
image_refined_curr = draw_corners_on_image(points,
image_refined_curr,
radius=4)
if cfg.visualization_dict['refined_mask']:
cv2.imwrite(os.path.join(dir_results, "mask_{}_refined.jpg".format(i)),
image_refined_curr)
def compute_lines_from_edge_candidate_clusters(img_edges, edge_candidates,
mask_plane,
mask_extract_contour,
mask_number, output_directory):
img_edge_candidates = np.copy(img_edges)
img_lines = np.copy(img_edges)
lines = []
for i, (key, val) in enumerate(edge_candidates.items()):
color = cfg.color_pallet[i]
val = get_unique_rows(val)
img_edge_candidates = draw_corners_on_image(val,
get_colored_image(
img_edges.copy()),
color=color,
radius=1)
y = img_edges.shape[0] - val[:, 1].reshape(
-1, 1) # OpenCV CS starts in upper left corner
X = val[:, 0].reshape(-1, 1)
prediction = ransac_linear_regression(X, y, draw=False)
if prediction is not None:
line_masked_image = find_line_masked_image(prediction,
np.copy(img_edges),
mask_number, key,
output_directory)
line = find_final_line(line_masked_image)
if line is None:
logger.debug(
"Use minimal line because other algorithm didn't work")
line = [[
int(x),
int(img_edges.shape[0] - prediction(x.reshape(1, -1)))
] for x in [np.min(X), np.max(X)]]
img_lines = cv2.line(img_lines, tuple(line[0]), tuple(line[1]),
color, 3)
lines.append(line)
if cfg.visualization_dict['lines']:
cv2.imwrite(
os.path.join(output_directory,
"mask_{}_line_image.jpg".format(mask_number)),
img_lines)
output_img = get_colored_image(img_edge_candidates)
if cfg.visualization_dict['keep_edges']:
cv2.imwrite(
os.path.join(output_directory,
"mask_{}_keep_edges.jpg".format(mask_number)),
output_img)
image_masked = draw_masks_on_image(output_img, [mask_plane])
if cfg.visualization_dict['initial_mask']:
cv2.imwrite(
os.path.join(output_directory,
"mask_{}_image.jpg".format(mask_number)),
image_masked)
return lines
def draw_ground_truth_and_estimation_on_image(gt_mask, img_mask, mask, color):
contours = cv2.findContours(gt_mask.copy().astype(np.uint8), cv2.RETR_TREE,
cv2.CHAIN_APPROX_SIMPLE)
img_mask = draw_masks_on_image(img_mask, [mask.astype(np.uint8)],
color) # draw estimation
img_mask = cv2.drawContours(img_mask,
contours[0],
0,
tuple([255, 0, 0]),
thickness=4) # draw ground truth
return img_mask
def compute_iou_of_mask_with_gt(refined_masks, gt_masks, draw=False):
from refiner.image_processing.draw import draw_masks_on_image
from utils.image_modification import scale_image
ious = []
for refined_mask in refined_masks:
tmp_ious = []
for gt_mask in gt_masks:
iou = compute_iou_of_masks(refined_mask, gt_mask)
if draw:
img = cv2.cvtColor((gt_mask.copy().astype(np.uint8)),
cv2.COLOR_GRAY2RGB)
img = draw_masks_on_image(img, [refined_mask])
cv2.imshow(
"", scale_image(np.hstack((gt_mask, refined_mask)), 0.6))
cv2.waitKey(0)
tmp_ious.append(iou)
best_idx = np.argmax(tmp_ious)
ious.append([best_idx, tmp_ious[best_idx]])
return ious