| from kornia.contrib import connected_components | |
| import torch | |
| import pdb | |
| import matplotlib.pyplot as plt | |
| import time | |
| # def reorder_int_labels(x): | |
| # _, y = torch.unique(x, return_inverse=True) | |
| # y -= y.min() | |
| # return y | |
| # def label_connected_component(labels, max_area=500, min_area=20, max_ccs=128, num_iterations=500): | |
| # assert len(labels.size()) == 2 | |
| # # per-label binary mask | |
| # unique_labels = torch.unique(labels).reshape(-1, 1, 1) # [?, 1, 1] | |
| # binary_masks = (labels.unsqueeze(0) == unique_labels).float() # [?, H, W] | |
| # # label connected components | |
| # cc = connected_components(binary_masks.unsqueeze(1), num_iterations=num_iterations) # [?, 1, H, W] | |
| # cc = reorder_int_labels(cc) | |
| # bincount = torch.bincount(cc.long().flatten()) | |
| # # find all connected components (id, mask, area, valid) | |
| # # cc_id = torch.nonzero(bincount) # [num_cc] | |
| # cc_id = torch.argsort(bincount)[-max_ccs:] | |
| # cc_mask = (cc == cc_id.reshape(1, -1, 1, 1)).sum(0) # [num_cc, H, W] | |
| # cc_area = bincount[cc_id] # [num_cc] | |
| # valid = (cc_area >= min_area) & (cc_area <= max_area) # [num_cc] | |
| # valid = valid.reshape(-1, 1, 1) # [num_cc, 1, 1] | |
| # # final labels for connected component | |
| # out = valid * cc_mask | |
| # out = out.argmax(0) | |
| # return out | |
| def reorder_int_labels(x): | |
| _, y = torch.unique(x, return_inverse=True) | |
| y -= y.min() | |
| return y | |
| def label_connected_component(labels, min_area=20, topk=256): | |
| size = labels.size() | |
| assert len(size) == 2 | |
| max_area = size[0] * size[1] - 1 | |
| # per-label binary mask | |
| unique_labels = torch.unique(labels).reshape(-1, 1, 1) # [?, 1, 1], where ? is the number of unique id | |
| binary_masks = (labels.unsqueeze(0) == unique_labels).float() # [?, H, W] | |
| # label connected components | |
| # cc is an integer tensor, each unique id represents a single connected component | |
| cc = connected_components(binary_masks.unsqueeze(1), num_iterations=500) # [?, 1, H, W] | |
| # reorder indices in cc so that cc_area tensor below is a smaller | |
| cc = reorder_int_labels(cc) | |
| # area of each connected components | |
| cc_area = torch.bincount(cc.long().flatten().cpu()).cuda() # bincount on GPU is much slower | |
| num_cc = cc_area.shape[0] | |
| valid = (cc_area >= min_area) & (cc_area <= max_area) # [num_cc] | |
| if num_cc < topk: | |
| selected_cc = torch.arange(num_cc).cuda() | |
| else: | |
| _, selected_cc = torch.topk(cc_area, k=topk) | |
| valid = valid[selected_cc] | |
| # collapse the 0th dimension, since there is only matched one connected component (across 0th dimension) | |
| cc_mask = (cc == selected_cc.reshape(1, -1, 1, 1)).sum(0) # [num_cc, H, W] | |
| cc_mask = cc_mask * valid.reshape(-1, 1, 1) | |
| out = cc_mask.argmax(0) | |
| return out | |
| # def reorder_int_labels(x): | |
| # _, y = torch.unique(x, return_inverse=True) | |
| # y -= y.min() | |
| # return y | |
| # def label_connected_component(labels, max_area=500, min_area=20): | |
| # assert len(labels.size()) == 2 | |
| # # per-label binary mask | |
| # unique_labels = torch.unique(labels).reshape(-1, 1, 1) # [?, 1, 1] | |
| # binary_masks = (labels.unsqueeze(0) == unique_labels).float() # [?, H, W] | |
| # # label connected components | |
| # cc = connected_components(binary_masks.unsqueeze(1)) # [?, 1, H, W] | |
| # cc = reorder_int_labels(cc) | |
| # bincount = torch.bincount(cc.long().flatten()) | |
| # # find all connected components (id, mask, area, valid) | |
| # cc_id = torch.nonzero(bincount) # [num_cc] | |
| # cc_mask = (cc == cc_id.reshape(1, -1, 1, 1)).sum(0) # [num_cc, H, W] | |
| # cc_area = bincount[cc_id] # [num_cc] | |
| # valid = (cc_area >= min_area) & (cc_area <= max_area) # [num_cc] | |
| # valid = valid.reshape(-1, 1, 1) # [num_cc, 1, 1] | |
| # # final labels for connected component | |
| # out = valid * cc_mask | |
| # out = out.argmax(0) | |