reference_detection.py

This module contains the functions necessary to extract head-feet pairs from images, load Detectron2 models, and parse the detectron2 output.

Use get_heads_feet to extract a head-feet pair from a single instance/binary mask.

Use extract_reference to extract a series of head-feet pairs from a single image.

Usage example

# If you plan on using detectron2
predictor, cfg = load_model()  # Basic COCO-trained mask-rcnn, threshold = 0.7
preds = predictor(image)  # Your image here
from detectron2.data import MetadataCatalog
objects = instances_to_dict(preds, MetadataCatalog.get(cfg.DATASETS.TRAIN[0]).get("thing_classes"))
reference = extract_reference(objects)  # These are the head-feet pairs!

Functions

ct_assist.reference_detection.extract_reference(objects: dict, step_size: int = 10, offset: float = 0.1, height_dict: dict = {'car': 1.425, 0.0247, 'person': 1.741, 0.05, 'truck': 3, 1}, always_warn: bool = True) → List[Tuple[numpy.ndarray, float, float]][source]

Extracts references from dictionary filled with predictions.

See instances_to_dict for objects’ format. The output is based on the output for get_heads_feet.

Parameters

  • objects (dict) – key : [mask]

  • step_size (int, optional) – How many pixels to skip, defaults to 10

  • offset (float, optional) – Minimum size relative to median distance between heads and feet, defaults to 0.1

  • height_dict (dictionary) – dictionary mapping object-type to height, defaults to HEIGHT_DICT

  • always_warn – If SkipFieldWarning should always be thrown

Returns

[(reference, height, STD)]

Return type

List[Tuple[np.ndarray, float, float]]

ct_assist.reference_detection.get_heads_feet(mask: None._VariableFunctionsClass.tensor, step_size=5, offset=0.1) → numpy.ndarray[source]

Gets head and feet from torch tensor segmentation mask.

Parameters

  • mask (torch.tensor) – Segmentation mask

  • step_size (int, optional) – Amount of found head feet pairs to skip to avoid overfitting, defaults to 5

  • offset (float, optional) – Permitted distance from median head-feet distance in percentages, defaults to 0.1

Returns

Returns numpy array with [heads, feet]

Return type

np.ndarray

ct_assist.reference_detection.instances_to_dict(preds: dict, thing_classes: list) → collections.defaultdict[source]

Take Detectron2.engine.DefaultPredictor output, and turns it into an easily parsable dictionary.

Have a cfg ready? Use: detectron2.data.MetadataCatalog.get(cfg.DATASETS.TRAIN[0]).get(“thing_classes”)) for your thing_classes.

Parameters

  • preds (dict) – dict to Instances object containing segmentation masks, output from defaultpredictor

  • thing_classes (list) – List mapping integers to string

Returns

object-name : [masks]

Return type

defaultdict

ct_assist.reference_detection.load_model(model_url: str = 'COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml', threshold: float = 0.7, return_cfg: bool = False) → detectron2.engine.defaults.DefaultPredictor[source]

Loads detectron2 model at model dir with threshold, option to return

Parameters

  • model_url (str, optional) – Points to yaml file containing config, defaults to “COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml”

  • threshold (float, optional) – Confidence threshold, defaults to 0.7

  • return_cfg (bool, optional) – If CfgNode obj should be returned, defaults to False

Returns

Detectron2 default predictor

Return type

DefaultPredictor