SDMGRPostProcessor

class mmocr.models.kie.SDMGRPostProcessor(link_type='none', key_node_idx=None, value_node_idx=None)

Postprocessor for SDMGR. It converts the node and edge scores into labels and edge labels. If the link_type is not “none”, it reconstructs the edge labels with different strategies specified by link_type, which is generally known as the “openset” mode. In “openset” mode, only the edges connecting from “key” to “value” nodes will be constructed.

参数

• link_type (str) –

  The type of link to be constructed. Defaults to ‘none’. Options are:

  • ’none’: The simplest link type involving no edge postprocessing. The edge prediction will be returned as-is.

  • ’one-to-one’: One key node can be connected to one value node.

  • ’one-to-many’: One key node can be connected to multiple value nodes.

  • ’many-to-one’: Multiple key nodes can be connected to one value node.

  • ’many-to-many’: No restrictions on the number of edges that a key/value node can have.

• key_node_idx (int, optional) – The label index of the key node. It must be specified if link_type is not “none”. Defaults to None.

• value_node_idx (int, optional) – The index of the value node. It must be specified if link_type is not “none”. Defaults to None.

decode_edges(node_labels, edge_scores, edge_labels)

Reconstruct the edges and update edge scores according to link_type.

参数

• data_sample (KIEDataSample) – A datasample containing prediction results.

• node_labels (torch.Tensor) –

• edge_scores (torch.Tensor) –

• edge_labels (torch.Tensor) –

返回

• edge_scores (Tensor): A float tensor of shape (N, N)

  indicating the confidence scores for edge predictions.

• edge_labels (Tensor): An integer tensor of shape (N, N)

  indicating the connection between nodes. Options are 0, 1.

返回类型

tuple(Tensor, Tensor)