Spaces:
Runtime error
Runtime error
| from __future__ import absolute_import, division | |
| import time | |
| import networkx as nx | |
| import numpy as np | |
| from scipy.ndimage import binary_dilation, binary_erosion | |
| from scipy.special import comb | |
| from skimage.filters import rank | |
| from skimage.morphology import dilation, disk, erosion, medial_axis | |
| from sklearn.neighbors import radius_neighbors_graph | |
| from .. import logging | |
| from ..metrics import batched_jaccard | |
| from ..utils.operations import bezier_curve | |
| __all__ = ['InteractiveScribblesRobot'] | |
| class InteractiveScribblesRobot(object): | |
| """ Robot that generates realistic scribbles simulating human interaction. | |
| # Attributes | |
| kernel_size: Float. Fraction of the square root of the area used | |
| to compute the dilation and erosion before computing the | |
| skeleton of the error masks. | |
| max_kernel_radius: Float. Maximum kernel radius when applying | |
| dilation and erosion. Default 16 pixels. | |
| min_nb_nodes: Integer. Number of nodes necessary to keep a connected | |
| graph and convert it into a scribble. | |
| nb_points: Integer. Number of points to sample the bezier curve | |
| when converting the final paths into curves. | |
| """ | |
| def __init__(self, | |
| kernel_size=.15, | |
| max_kernel_radius=16, | |
| min_nb_nodes=4, | |
| nb_points=1000): | |
| """ Robot constructor | |
| """ | |
| if kernel_size >= 1. or kernel_size < 0: | |
| raise ValueError('kernel_size must be a value between [0, 1).') | |
| self.kernel_size = kernel_size | |
| self.max_kernel_radius = max_kernel_radius | |
| self.min_nb_nodes = min_nb_nodes | |
| self.nb_points = nb_points | |
| def _generate_scribble_mask(self, mask): | |
| """ Generate the skeleton from a mask | |
| Given an error mask, the medial axis is computed to obtain the | |
| skeleton of the objects. In order to obtain smoother skeleton and | |
| remove small objects, an erosion and dilation operations are performed. | |
| The kernel size used is proportional the squared of the area. | |
| # Arguments | |
| mask: Numpy Array. Error mask | |
| Returns: | |
| skel: Numpy Array. Skeleton mask | |
| """ | |
| mask = np.asarray(mask, dtype=np.uint8) | |
| side = np.sqrt(np.sum(mask > 0)) | |
| mask_ = mask | |
| # kernel_size = int(self.kernel_size * side) | |
| kernel_radius = self.kernel_size * side * .5 | |
| kernel_radius = min(kernel_radius, self.max_kernel_radius) | |
| logging.verbose( | |
| 'Erosion and dilation with kernel radius: {:.1f}'.format( | |
| kernel_radius), 2) | |
| compute = True | |
| while kernel_radius > 1. and compute: | |
| kernel = disk(kernel_radius) | |
| mask_ = rank.minimum(mask.copy(), kernel) | |
| mask_ = rank.maximum(mask_, kernel) | |
| compute = False | |
| if mask_.astype(np.bool).sum() == 0: | |
| compute = True | |
| prev_kernel_radius = kernel_radius | |
| kernel_radius *= .9 | |
| logging.verbose('Reducing kernel radius from {:.1f} '.format( | |
| prev_kernel_radius) + | |
| 'pixels to {:.1f}'.format(kernel_radius), 1) | |
| mask_ = np.pad( | |
| mask_, ((1, 1), (1, 1)), mode='constant', constant_values=False) | |
| skel = medial_axis(mask_.astype(np.bool)) | |
| skel = skel[1:-1, 1:-1] | |
| return skel | |
| def _mask2graph(self, skeleton_mask): | |
| """ Transforms a skeleton mask into a graph | |
| Args: | |
| skeleton_mask (ndarray): Skeleton mask | |
| Returns: | |
| tuple(nx.Graph, ndarray): Returns a tuple where the first element | |
| is a Graph and the second element is an array of xy coordinates | |
| indicating the coordinates for each Graph node. | |
| If an empty mask is given, None is returned. | |
| """ | |
| mask = np.asarray(skeleton_mask, dtype=np.bool) | |
| if np.sum(mask) == 0: | |
| return None | |
| h, w = mask.shape | |
| x, y = np.arange(w), np.arange(h) | |
| X, Y = np.meshgrid(x, y) | |
| X, Y = X.ravel(), Y.ravel() | |
| M = mask.ravel() | |
| X, Y = X[M], Y[M] | |
| points = np.c_[X, Y] | |
| G = radius_neighbors_graph(points, np.sqrt(2), mode='distance') | |
| T = nx.from_scipy_sparse_matrix(G) | |
| return T, points | |
| def _acyclics_subgraphs(self, G): | |
| """ Divide a graph into connected components subgraphs | |
| Divide a graph into connected components subgraphs and remove its | |
| cycles removing the edge with higher weight inside the cycle. Also | |
| prune the graphs by number of nodes in case the graph has not enought | |
| nodes. | |
| Args: | |
| G (nx.Graph): Graph | |
| Returns: | |
| list(nx.Graph): Returns a list of graphs which are subgraphs of G | |
| with cycles removed. | |
| """ | |
| if not isinstance(G, nx.Graph): | |
| raise TypeError('G must be a nx.Graph instance') | |
| S = [] # List of subgraphs of G | |
| for c in nx.connected_components(G): | |
| g = G.subgraph(c).copy() | |
| # Remove all cycles that we may find | |
| has_cycles = True | |
| while has_cycles: | |
| try: | |
| cycle = nx.find_cycle(g) | |
| weights = np.asarray([G[u][v]['weight'] for u, v in cycle]) | |
| idx = weights.argmax() | |
| # Remove the edge with highest weight at cycle | |
| g.remove_edge(*cycle[idx]) | |
| except nx.NetworkXNoCycle: | |
| has_cycles = False | |
| if len(g) < self.min_nb_nodes: | |
| # Prune small subgraphs | |
| logging.verbose('Remove a small line with {} nodes'.format( | |
| len(g)), 1) | |
| continue | |
| S.append(g) | |
| return S | |
| def _longest_path_in_tree(self, G): | |
| """ Given a tree graph, compute the longest path and return it | |
| Given an undirected tree graph, compute the longest path and return it. | |
| The approach use two shortest path transversals (shortest path in a | |
| tree is the same as longest path). This could be improve but would | |
| require implement it: | |
| https://cs.stackexchange.com/questions/11263/longest-path-in-an-undirected-tree-with-only-one-traversal | |
| Args: | |
| G (nx.Graph): Graph which should be an undirected tree graph | |
| Returns: | |
| list(int): Returns a list of indexes of the nodes belonging to the | |
| longest path. | |
| """ | |
| if not isinstance(G, nx.Graph): | |
| raise TypeError('G must be a nx.Graph instance') | |
| if not nx.is_tree(G): | |
| raise ValueError('Graph G must be a tree (graph without cycles)') | |
| # Compute the furthest node to the random node v | |
| v = list(G.nodes())[0] | |
| distance = nx.single_source_shortest_path_length(G, v) | |
| vp = max(distance.items(), key=lambda x: x[1])[0] | |
| # From this furthest point v' find again the longest path from it | |
| distance = nx.single_source_shortest_path(G, vp) | |
| longest_path = max(distance.values(), key=len) | |
| # Return the longest path | |
| return list(longest_path) | |
| def interact(self, | |
| sequence, | |
| pred_masks, | |
| gt_masks, | |
| nb_objects=None, | |
| frame=None): | |
| """ Interaction of the Scribble robot given a prediction. | |
| Given the sequence and a mask prediction, the robot will return a | |
| scribble in the region that fails the most. | |
| # Arguments | |
| sequence: String. Name of the sequence to interact with. | |
| pred_masks: Numpy Array. Array with the prediction masks. It must | |
| be an integer array with shape (B x H x W), with B being the number | |
| of frames of the sequence. | |
| gt_masks: Numpy Array. Array with the ground truth of the sequence. | |
| It must have the same data type and shape as `pred_masks`. | |
| nb_objects: Integer. Number of objects in the ground truth mask. If | |
| `None` the value will be infered from `y_true`. Setting this | |
| value will speed up the computation. | |
| frame: Integer. Frame to generate the scribble. If not given, the | |
| worst frame given by the jaccard will be used. | |
| # Returns | |
| dict: Return a scribble (default representation). | |
| """ | |
| robot_start = time.time() | |
| predictions = np.asarray(pred_masks, dtype=np.int) | |
| annotations = np.asarray(gt_masks, dtype=np.int) | |
| nb_frames = len(annotations) | |
| if nb_objects is None: | |
| obj_ids = np.unique(annotations) | |
| obj_ids = obj_ids[(obj_ids > 0) & (obj_ids < 255)] | |
| nb_objects = len(obj_ids) | |
| obj_ids = [i for i in range(nb_objects + 1)] | |
| # Infer height and width of the sequence | |
| h, w = annotations.shape[1:3] | |
| img_shape = np.asarray([w, h], dtype=np.float) | |
| if frame is None: | |
| jac = batched_jaccard( | |
| annotations, predictions, nb_objects=nb_objects) | |
| worst_frame = jac.argmin() | |
| logging.verbose( | |
| 'For sequence {} the worst frames is #{} with Jaccard: {:.3f}'. | |
| format(sequence, worst_frame, jac.min()), 2) | |
| else: | |
| worst_frame = frame | |
| pred, gt = predictions[worst_frame], annotations[worst_frame] | |
| scribbles = [[] for _ in range(nb_frames)] | |
| for obj_id in obj_ids: | |
| logging.verbose( | |
| 'Creating scribbles from error mask at object_id={}'.format( | |
| obj_id), 2) | |
| start_time = time.time() | |
| error_mask = (gt == obj_id) & (pred != obj_id) | |
| if error_mask.sum() == 0: | |
| logging.info( | |
| 'Error mask of object ID {} is empty. Skip object ID.'. | |
| format(obj_id)) | |
| continue | |
| # Generate scribbles | |
| skel_mask = self._generate_scribble_mask(error_mask) | |
| skel_time = time.time() - start_time | |
| logging.verbose( | |
| 'Time to compute the skeleton mask: {:.3f} ms'.format( | |
| skel_time * 1000), 2) | |
| if skel_mask.sum() == 0: | |
| continue | |
| G, P = self._mask2graph(skel_mask) | |
| mask2graph_time = time.time() - start_time - skel_time | |
| logging.verbose( | |
| 'Time to transform the skeleton mask into a graph: ' + | |
| '{:.3f} ms'.format(mask2graph_time * 1000), 2) | |
| t_start = time.time() | |
| S = self._acyclics_subgraphs(G) | |
| t = (time.time() - t_start) * 1000 | |
| logging.verbose( | |
| 'Time to split into connected components subgraphs ' + | |
| 'and remove the cycles: {:.3f} ms'.format(t), 2) | |
| t_start = time.time() | |
| longest_paths_idx = [self._longest_path_in_tree(s) for s in S] | |
| longest_paths = [P[idx] for idx in longest_paths_idx] | |
| t = (time.time() - t_start) * 1000 | |
| logging.verbose( | |
| 'Time to compute the longest path on the trees: {:.3f} ms'. | |
| format(t), 2) | |
| t_start = time.time() | |
| scribbles_paths = [ | |
| bezier_curve(p, self.nb_points) for p in longest_paths | |
| ] | |
| t = (time.time() - t_start) * 1000 | |
| logging.verbose( | |
| 'Time to compute the bezier curves: {:.3f} ms'.format(t), 2) | |
| end_time = time.time() | |
| logging.verbose( | |
| 'Generating the scribble for object id {} '.format(obj_id) + | |
| 'took {:.3f} ms'.format((end_time - start_time) * 1000), 2) | |
| # Generate scribbles data file | |
| for p in scribbles_paths: | |
| p /= img_shape | |
| path_data = { | |
| 'path': p.tolist(), | |
| 'object_id': int(obj_id), | |
| 'start_time': start_time, | |
| 'end_time': end_time | |
| } | |
| scribbles[worst_frame].append(path_data) | |
| scribbles_data = {'scribbles': scribbles, 'sequence': sequence} | |
| t = time.time() - robot_start | |
| logging.info(('The robot took {:.3f} s to generate all the ' | |
| 'scribbles for {} objects. Sequence {}.').format( | |
| t, nb_objects, sequence)) | |
| return scribbles_data | |