text: import cv2
  import numpy as np
  import torch
  from lanms import merge_quadrangle_n9 as la_nms
  from mmcv.ops import RoIAlignRotated
  
  from mmocr.models.textdet.postprocess.wrapper import fill_hole
  from .utils import (euclidean_distance_matrix, feature_embedding,
                      normalize_adjacent_matrix)
  
  
  class ProposalLocalGraphs:
      """Propose text components and generate local graphs for GCN to classify
      the k-nearest neighbors of a pivot in DRRG: Deep Relational Reasoning Graph
      Network for Arbitrary Shape Text Detection.
  
      [https://arxiv.org/abs/2003.07493]. This code was partially adapted from
      https://github.com/GXYM/DRRG licensed under the MIT license.
  
      Args:
          k_at_hops (tuple(int)): The number of i-hop neighbors, i = 1, 2.
          num_adjacent_linkages (int): The number of linkages when constructing
              adjacent matrix.
          node_geo_feat_len (int): The length of embedded geometric feature
              vector of a text component.
          pooling_scale (float): The spatial scale of rotated RoI-Align.
          pooling_output_size (tuple(int)): The output size of rotated RoI-Align.
          nms_thr (float): The locality-aware NMS threshold for text components.
          min_width (float): The minimum width of text components.
          max_width (float): The maximum width of text components.
          comp_shrink_ratio (float): The shrink ratio of text components.
          comp_w_h_ratio (float): The width to height ratio of text components.
          comp_score_thr (float): The score threshold of text component.
          text_region_thr (float): The threshold for text region probability map.
          center_region_thr (float): The threshold for text center region
              probability map.
          center_region_area_thr (int): The threshold for filtering small-sized
              text center region.
      """
  
      def __init__(self, k_at_hops, num_adjacent_linkages, node_geo_feat_len,
                   pooling_scale, pooling_output_size, nms_thr, min_width,
                   max_width, comp_shrink_ratio, comp_w_h_ratio, comp_score_thr,
                   text_region_thr, center_region_thr, center_region_area_thr):
  
          assert len(k_at_hops) == 2
          assert isinstance(k_at_hops, tuple)
          assert isinstance(num_adjacent_linkages, int)
          assert isinstance(node_geo_feat_len, int)
          assert isinstance(pooling_scale, float)
          assert isinstance(pooling_output_size, tuple)
          assert isinstance(nms_thr, float)
          assert isinstance(min_width, float)
          assert isinstance(max_width, float)
          assert isinstance(comp_shrink_ratio, float)
          assert isinstance(comp_w_h_ratio, float)
          assert isinstance(comp_score_thr, float)
          assert isinstance(text_region_thr, float)
          assert isinstance(center_region_thr, float)
          assert isinstance(center_region_area_thr, int)
  
          self.k_at_hops = k_at_hops
          self.active_connection = num_adjacent_linkages
          self.local_graph_depth = len(self.k_at_hops)
          self.node_geo_feat_dim = node_geo_feat_len
          self.pooling = RoIAlignRotated(pooling_output_size, pooling_scale)
          self.nms_thr = nms_thr
          self.min_width = min_width
          self.max_width = max_width
          self.comp_shrink_ratio = comp_shrink_ratio
          self.comp_w_h_ratio = comp_w_h_ratio
          self.comp_score_thr = comp_score_thr
          self.text_region_thr = text_region_thr
          self.center_region_thr = center_region_thr
          self.center_region_area_thr = center_region_area_thr
  
      def propose_comps(self, score_map, top_height_map, bot_height_map, sin_map,
                        cos_map, comp_score_thr, min_width, max_width,
                        comp_shrink_ratio, comp_w_h_ratio):
          """Propose text components.
  
          Args:
              score_map (ndarray): The score map for NMS.
              top_height_map (ndarray): The predicted text height map from each
                  pixel in text center region to top sideline.
              bot_height_map (ndarray): The predicted text height map from each
                  pixel in text center region to bottom sideline.
              sin_map (ndarray): The predicted sin(theta) map.
              cos_map (ndarray): The predicted cos(theta) map.
              comp_score_thr (float): The score threshold of text component.
              min_width (float): The minimum width of text components.
              max_width (float): The maximum width of text components.
              comp_shrink_ratio (float): The shrink ratio of text components.
              comp_w_h_ratio (float): The width to height ratio of text
                  components.
  
          Returns:
              text_comps (ndarray): The text components.
          """
  
          comp_centers = np.argwhere(score_map > comp_score_thr)
          comp_centers = comp_centers[np.argsort(comp_centers[:, 0])]
          y = comp_centers[:, 0]
          x = comp_centers[:, 1]
  
          top_height = top_height_map[y, x].reshape((-1, 1)) * comp_shrink_ratio
          bot_height = bot_height_map[y, x].reshape((-1, 1)) * comp_shrink_ratio
          sin = sin_map[y, x].reshape((-1, 1))
          cos = cos_map[y, x].reshape((-1, 1))
  
          top_mid_pts = comp_centers + np.hstack(
              [top_height * sin, top_height * cos])
          bot_mid_pts = comp_centers - np.hstack(
              [bot_height * sin, bot_height * cos])
  
          width = (top_height + bot_height) * comp_w_h_ratio
          width = np.clip(width, min_width, max_width)
          r = width / 2
  
          tl = top_mid_pts[:, ::-1] - np.hstack([-r * sin, r * cos])
          tr = top_mid_pts[:, ::-1] + np.hstack([-r * sin, r * cos])
          br = bot_mid_pts[:, ::-1] + np.hstack([-r * sin, r * cos])
          bl = bot_mid_pts[:, ::-1] - np.hstack([-r * sin, r * cos])
          text_comps = np.hstack([tl, tr, br, bl]).astype(np.float32)
  
          score = score_map[y, x].reshape((-1, 1))
          text_comps = np.hstack([text_comps, score])
  
          return text_comps
  
      def propose_comps_and_attribs(self, text_region_map, center_region_map,
                                    top_height_map, bot_height_map, sin_map,
                                    cos_map):
          """Generate text components and attributes.
  
          Args:
              text_region_map (ndarray): The predicted text region probability
                  map.
              center_region_map (ndarray): The predicted text center region
                  probability map.
              top_height_map (ndarray): The predicted text height map from each
                  pixel in text center region to top sideline.
              bot_height_map (ndarray): The predicted text height map from each
                  pixel in text center region to bottom sideline.
              sin_map (ndarray): The predicted sin(theta) map.
              cos_map (ndarray): The predicted cos(theta) map.
  
          Returns:
              comp_attribs (ndarray): The text component attributes.
              text_comps (ndarray): The text components.
          """
  
          assert (text_region_map.shape == center_region_map.shape ==
                  top_height_map.shape == bot_height_map.shape == sin_map.shape
                  == cos_map.shape)
          text_mask = text_region_map > self.text_region_thr
          center_region_mask = (center_region_map >
                                self.center_region_thr) * text_mask
  
          scale = np.sqrt(1.0 / (sin_map**2 + cos_map**2 + 1e-8))
          sin_map, cos_map = sin_map * scale, cos_map * scale
  
          center_region_mask = fill_hole(center_region_mask)
          center_region_contours, _ = cv2.findContours(
              center_region_mask.astype(np.uint8), cv2.RETR_TREE,
              cv2.CHAIN_APPROX_SIMPLE)
  
          mask_sz = center_region_map.shape
          comp_list = []
          for contour in center_region_contours:
              current_center_mask = np.zeros(mask_sz)
              cv2.drawContours(current_center_mask, [contour], -1, 1, -1)
              if current_center_mask.sum() <= self.center_region_area_thr:
                  continue
              score_map = text_region_map * current_center_mask
  
              text_comps = self.propose_comps(score_map, top_height_map,
                                              bot_height_map, sin_map, cos_map,
                                              self.comp_score_thr,
                                              self.min_width, self.max_width,
                                              self.comp_shrink_ratio,
                                              self.comp_w_h_ratio)
  
              text_comps = la_nms(text_comps, self.nms_thr)
              text_comp_mask = np.zeros(mask_sz)
              text_comp_boxes = text_comps[:, :8].reshape(
                  (-1, 4, 2)).astype(np.int32)
  
              cv2.drawContours(text_comp_mask, text_comp_boxes, -1, 1, -1)
              if (text_comp_mask * text_mask).sum() < text_comp_mask.sum() * 0.5:
                  continue
              if text_comps.shape[-1] > 0:
                  comp_list.append(text_comps)
  
          if len(comp_list) <= 0:
              return None, None
  
          text_comps = np.vstack(comp_list)
          text_comp_boxes = text_comps[:, :8].reshape((-1, 4, 2))
          centers = np.mean(text_comp_boxes, axis=1).astype(np.int32)
          x = centers[:, 0]
          y = centers[:, 1]
  
          scores = []
          for text_comp_box in text_comp_boxes:
              text_comp_box[:, 0] = np.clip(text_comp_box[:, 0], 0,
                                            mask_sz[1] - 1)
              text_comp_box[:, 1] = np.clip(text_comp_box[:, 1], 0,
                                            mask_sz[0] - 1)
              min_coord = np.min(text_comp_box, axis=0).astype(np.int32)
              max_coord = np.max(text_comp_box, axis=0).astype(np.int32)
              text_comp_box = text_comp_box - min_coord
              box_sz = (max_coord - min_coord + 1)
              temp_comp_mask = np.zeros((box_sz[1], box_sz[0]), dtype=np.uint8)
              cv2.fillPoly(temp_comp_mask, [text_comp_box.astype(np.int32)], 1)
              temp_region_patch = text_region_map[min_coord[1]:(max_coord[1] +
                                                                1),
                                                  min_coord[0]:(max_coord[0] +
                                                                1)]
              score = cv2.mean(temp_region_patch, temp_comp_mask)[0]
              scores.append(score)
          scores = np.array(scores).reshape((-1, 1))
          text_comps = np.hstack([text_comps[:, :-1], scores])
  
          h = top_height_map[y, x].reshape(
              (-1, 1)) + bot_height_map[y, x].reshape((-1, 1))
          w = np.clip(h * self.comp_w_h_ratio, self.min_width, self.max_width)
          sin = sin_map[y, x].reshape((-1, 1))
          cos = cos_map[y, x].reshape((-1, 1))
  
          x = x.reshape((-1, 1))
          y = y.reshape((-1, 1))
          comp_attribs = np.hstack([x, y, h, w, cos, sin])
  
          return comp_attribs, text_comps
  
      def generate_local_graphs(self, sorted_dist_inds, node_feats):
          """Generate local graphs and graph convolution network input data.
  
          Args:
              sorted_dist_inds (ndarray): The node indices sorted according to
                  the Euclidean distance.
              node_feats (tensor): The features of nodes in graph.
  
          Returns:
              local_graphs_node_feats (tensor): The features of nodes in local
                  graphs.
              adjacent_matrices (tensor): The adjacent matrices.
              pivots_knn_inds (tensor): The k-nearest neighbor indices in
                  local graphs.
              pivots_local_graphs (tensor): The indices of nodes in local
                  graphs.
          """
  
          assert sorted_dist_inds.ndim == 2
          assert (sorted_dist_inds.shape[0] == sorted_dist_inds.shape[1] ==
                  node_feats.shape[0])
  
          knn_graph = sorted_dist_inds[:, 1:self.k_at_hops[0] + 1]
          pivot_local_graphs = []
          pivot_knns = []
          device = node_feats.device
  
          for pivot_ind, knn in enumerate(knn_graph):
  
              local_graph_neighbors = set(knn)
  
              for neighbor_ind in knn:
                  local_graph_neighbors.update(
                      set(sorted_dist_inds[neighbor_ind,
                                           1:self.k_at_hops[1] + 1]))
  
              local_graph_neighbors.discard(pivot_ind)
              pivot_local_graph = list(local_graph_neighbors)
              pivot_local_graph.insert(0, pivot_ind)
              pivot_knn = [pivot_ind] + list(knn)
  
              pivot_local_graphs.append(pivot_local_graph)
              pivot_knns.append(pivot_knn)
  
          num_max_nodes = max([
              len(pivot_local_graph) for pivot_local_graph in pivot_local_graphs
          ])
  
          local_graphs_node_feat = []
          adjacent_matrices = []
          pivots_knn_inds = []
          pivots_local_graphs = []
  
          for graph_ind, pivot_knn in enumerate(pivot_knns):
              pivot_local_graph = pivot_local_graphs[graph_ind]
              num_nodes = len(pivot_local_graph)
              pivot_ind = pivot_local_graph[0]
              node2ind_map = {j: i for i, j in enumerate(pivot_local_graph)}
  
              knn_inds = torch.tensor([node2ind_map[i]
                                       for i in pivot_knn[1:]]).long().to(device)
              pivot_feats = node_feats[pivot_ind]
              normalized_feats = node_feats[pivot_local_graph] - pivot_feats
  
              adjacent_matrix = np.zeros((num_nodes, num_nodes))
              for node in pivot_local_graph:
                  neighbors = sorted_dist_inds[node,
                                               1:self.active_connection + 1]
                  for neighbor in neighbors:
                      if neighbor in pivot_local_graph:
                          adjacent_matrix[node2ind_map[node],
                                          node2ind_map[neighbor]] = 1
                          adjacent_matrix[node2ind_map[neighbor],
                                          node2ind_map[node]] = 1
  
              adjacent_matrix = normalize_adjacent_matrix(adjacent_matrix)
              pad_adjacent_matrix = torch.zeros((num_max_nodes, num_max_nodes),
                                                dtype=torch.float,
                                                device=device)
              pad_adjacent_matrix[:num_nodes, :num_nodes] = torch.from_numpy(
                  adjacent_matrix)
  
              pad_normalized_feats = torch.cat([
                  normalized_feats,
                  torch.zeros(
                      (num_max_nodes - num_nodes, normalized_feats.shape[1]),
                      dtype=torch.float,
                      device=device)
              ],
                                               dim=0)
  
              local_graph_nodes = torch.tensor(pivot_local_graph)
              local_graph_nodes = torch.cat([
                  local_graph_nodes,
                  torch.zeros(num_max_nodes - num_nodes, dtype=torch.long)
              ],
                                            dim=-1)
  
              local_graphs_node_feat.append(pad_normalized_feats)
              adjacent_matrices.append(pad_adjacent_matrix)
              pivots_knn_inds.append(knn_inds)
              pivots_local_graphs.append(local_graph_nodes)
  
          local_graphs_node_feat = torch.stack(local_graphs_node_feat, 0)
          adjacent_matrices = torch.stack(adjacent_matrices, 0)
          pivots_knn_inds = torch.stack(pivots_knn_inds, 0)
          pivots_local_graphs = torch.stack(pivots_local_graphs, 0)
  
          return (local_graphs_node_feat, adjacent_matrices, pivots_knn_inds,
                  pivots_local_graphs)
  
      def __call__(self, preds, feat_maps):
          """Generate local graphs and graph convolutional network input data.
  
          Args:
              preds (tensor): The predicted maps.
              feat_maps (tensor): The feature maps to extract content feature of
                  text components.
  
          Returns:
              none_flag (bool): The flag showing whether the number of proposed
                  text components is 0.
              local_graphs_node_feats (tensor): The features of nodes in local
                  graphs.
              adjacent_matrices (tensor): The adjacent matrices.
              pivots_knn_inds (tensor): The k-nearest neighbor indices in
                  local graphs.
              pivots_local_graphs (tensor): The indices of nodes in local
                  graphs.
              text_comps (ndarray): The predicted text components.
          """
  
          if preds.ndim == 4:
              assert preds.shape[0] == 1
              preds = torch.squeeze(preds)
          pred_text_region = torch.sigmoid(preds[0]).data.cpu().numpy()
          pred_center_region = torch.sigmoid(preds[1]).data.cpu().numpy()
          pred_sin_map = preds[2].data.cpu().numpy()
          pred_cos_map = preds[3].data.cpu().numpy()
          pred_top_height_map = preds[4].data.cpu().numpy()
          pred_bot_height_map = preds[5].data.cpu().numpy()
          device = preds.device
  
          comp_attribs, text_comps = self.propose_comps_and_attribs(
              pred_text_region, pred_center_region, pred_top_height_map,
              pred_bot_height_map, pred_sin_map, pred_cos_map)
  
          if comp_attribs is None or len(comp_attribs) < 2:
              none_flag = True
              return none_flag, (0, 0, 0, 0, 0)
  
          comp_centers = comp_attribs[:, 0:2]
          distance_matrix = euclidean_distance_matrix(comp_centers, comp_centers)
  
          geo_feats = feature_embedding(comp_attribs, self.node_geo_feat_dim)
          geo_feats = torch.from_numpy(geo_feats).to(preds.device)
  
          batch_id = np.zeros((comp_attribs.shape[0], 1), dtype=np.float32)
          comp_attribs = comp_attribs.astype(np.float32)
          angle = np.arccos(comp_attribs[:, -2]) * np.sign(comp_attribs[:, -1])
          angle = angle.reshape((-1, 1))
          rotated_rois = np.hstack([batch_id, comp_attribs[:, :-2], angle])
          rois = torch.from_numpy(rotated_rois).to(device)
  
          content_feats = self.pooling(feat_maps, rois)
          content_feats = content_feats.view(content_feats.shape[0],
                                             -1).to(device)
          node_feats = torch.cat([content_feats, geo_feats], dim=-1)
  
          sorted_dist_inds = np.argsort(distance_matrix, axis=1)
          (local_graphs_node_feat, adjacent_matrices, pivots_knn_inds,
           pivots_local_graphs) = self.generate_local_graphs(
               sorted_dist_inds, node_feats)
  
          none_flag = False
          return none_flag, (local_graphs_node_feat, adjacent_matrices,
                             pivots_knn_inds, pivots_local_graphs, text_comps)
  