KIEInferencer¶

class mmocr.apis.inferencers.KIEInferencer(model=None, weights=None, device=None, scope='mmocr')[source]¶

Key Information Extraction Inferencer.

Parameters

model (str, optional) – Path to the config file or the model name defined in metafile. For example, it could be “sdmgr_unet16_60e_wildreceipt” or “configs/kie/sdmgr/sdmgr_unet16_60e_wildreceipt.py”. If model is not specified, user must provide the weights saved by MMEngine which contains the config string. Defaults to None.
weights (str, optional) – Path to the checkpoint. If it is not specified and model is a model name of metafile, the weights will be loaded from metafile. Defaults to None.
device (str, optional) – Device to run inference. If None, the available device will be automatically used. Defaults to None.
scope (str, optional) – The scope of the model. Defaults to “mmocr”.

Return type

static kie_collate(data_batch)[source]¶

A collate function designed for KIE, where the first element (input) is a dict and we only want to keep it as-is instead of batching elements inside.

Returns: Transversed Data in the same format as the data_itement of data_batch.
Return type: Any
Parameters: data_batch (Sequence) –

pred2dict(data_sample)[source]¶

Extract elements necessary to represent a prediction into a dictionary. It’s better to contain only basic data elements such as strings and numbers in order to guarantee it’s json-serializable.

Parameters: data_sample (TextRecogDataSample) – The data sample to be converted.
Returns: The output dictionary.
Return type: dict

visualize(inputs, preds, return_vis=False, show=False, wait_time=0, draw_pred=True, pred_score_thr=0.3, save_vis=False, img_out_dir='')[source]¶

Visualize predictions.

Parameters

inputs (List[Union[str, np.ndarray]]) – Inputs for the inferencer.
preds (List[Dict]) – Predictions of the model.
return_vis (bool) – Whether to return the visualization result. Defaults to False.
show (bool) – Whether to display the image in a popup window. Defaults to False.
wait_time (float) – The interval of show (s). Defaults to 0.
draw_pred (bool) – Whether to draw predicted bounding boxes. Defaults to True.
pred_score_thr (float) – Minimum score of bboxes to draw. Defaults to 0.3.
save_vis (bool) – Whether to save the visualization result. Defaults to False.
img_out_dir (str) – Output directory of visualization results. If left as empty, no file will be saved. Defaults to ‘’.

Returns

Returns visualization results only if applicable.

Return type

List[np.ndarray] or None