YOLO helpers

class flat_bug.yolo_helpers.ResultsWithTiles(tiles: List[int] = None, polygons=None, *args, **kwargs)
new() ResultsWithTiles

Creates a new Results object with the same image, path, names, and speed attributes.

Returns:

A new Results object with copied attributes from the original instance.

Return type:

(Results)

Examples

>>> results = model("path/to/image.jpg")
>>> new_result = results[0].new()
flat_bug.yolo_helpers.clip_boxes(boxes: Tensor | ndarray, shape: Tuple[int, int])

Takes a list of bounding boxes and a shape (height, width) and clips the bounding boxes to the shape.

Parameters:
  • boxes (Union[torch.Tensor, np.ndarray]) – the bounding boxes to clip

  • shape (tuple) – The maximum x and y values for the bounding boxes.

Returns:

Clipped boxes

Return type:

out (Union[torch.Tensor, np.ndarray])

flat_bug.yolo_helpers.crop_mask(masks: Tensor, boxes: Tensor) Tensor

It takes a mask and a bounding box, and returns a mask that is cropped to the bounding box.

Parameters:
  • masks (torch.Tensor) – [n, h, w] tensor of masks

  • boxes (torch.Tensor) – [n, 4] tensor of bbox coordinates in relative point form

Returns:

The masks are being cropped to the bounding box.

Return type:

out (torch.Tensor)

flat_bug.yolo_helpers.expand_bottom_right(mask: Tensor)

Add an extra pixel above next to bottom/right edges of the region of 1s.

Parameters:

mask (torch.Tensor) – A binary mask tensor of shape [h, w].

Returns:

A binary mask tensor of shape [h, w], where an extra pixel is added above next to left/top edges of the region of 1s.

Return type:

out (torch.Tensor)

flat_bug.yolo_helpers.merge_tile_results(results=typing.List[ultralytics.engine.results.Results], orig_img: Tensor | None = None, box_offsetters: Tensor | None = None, mask_offsetters: Tensor | None = None, new_shape: Tuple[int, int] | List[int] = None, clamp_boxes: Tuple[int, int] | List[int] = (None, None), max_mask_size: int = 700, exclude_masks: bool = False) ResultsWithTiles

Merges results from multiple images into a single Results object, possibly with a new image.

flat_bug.yolo_helpers.postprocess(preds, imgs: List[Tensor], max_det: int = 300, min_confidence: float = 0, iou_threshold: float = 0.1, nms: int = 0, valid_size_range: Tuple[int, int] | List[int] | None = None, edge_margin: int | None = None) List[Results]

Postprocesses the predictions of the model.

Parameters:
  • preds (list) – A list of predictions from the model.

  • imgs (List[torch.Tensor]) – A list of images that were passed to the model.

  • max_det (int, optional) – The maximum number of detections to return. Defaults to 300.

  • min_confidence (float, optional) – The minimum confidence of the predictions to return. Defaults to 0.

  • iou_threshold (float, optional) – The IoU threshold for non-maximum suppression. Defaults to 0.1.

  • nms (int, optional) – The type of non-maximum suppression to use. Defaults to 0. 0 is no NMS, 1 is standard NMS, 2 is fancy NMS and 3 is mask NMS.

  • valid_size_range (tuple, optional) – The range of valid sizes for the bounding boxes in pixels. Defaults to None (no valid size range).

  • edge_margin (int, optional) – The minimum gap between the edge of the image and the bounding box in pixels for a prediction to be considered valid. Defaults to None (no edge margin).

Returns:

A list of postprocessed predictions.

Return type:

out (List[ultralytics.engine.results.Results])

flat_bug.yolo_helpers.process_mask(protos: Tensor, masks_in: Tensor, bboxes: Tensor, shape: Tuple[int, int] | List[int], upsample: bool = False) Tensor

Apply masks to bounding boxes using the output of the mask head.

Parameters:
  • protos (torch.Tensor) – A tensor of shape [mask_dim, mask_h, mask_w].

  • masks_in (torch.Tensor) – A tensor of shape [n, mask_dim], where n is the number of masks after NMS.

  • bboxes (torch.Tensor) – A tensor of shape [n, 4], where n is the number of masks after NMS.

  • shape (tuple) – A tuple of integers representing the size of the input image in the format (h, w).

  • upsample (bool, optional) – A flag to indicate whether to upsample the mask to the original image size. Default is False.

Returns:

A binary mask tensor of shape [n, h, w], where n is the number of masks after NMS, and h and w

are the height and width of the input image. The mask is applied to the bounding boxes.

Return type:

out (torch.Tensor)

flat_bug.yolo_helpers.scale_boxes(img1_shape: Tuple[int, int], boxes: Tensor, img0_shape: Tuple[int, int], ratio_pad=None, padding: bool = True, xywh: bool = False) Tensor

Rescales bounding boxes (in the format of xyxy by default) from the shape of the image they were originally specified in (img1_shape) to the shape of a different image (img0_shape).

Parameters:
  • img1_shape (tuple) – The shape of the image that the bounding boxes are for, in the format of (height, width).

  • boxes (torch.Tensor) – the bounding boxes of the objects in the image, in the format of (x1, y1, x2, y2)

  • img0_shape (tuple) – the shape of the target image, in the format of (height, width).

  • ratio_pad (Optional[Tuple[float, Tuple[int, int]]], optional) – a tuple of (ratio, pad) for scaling the boxes. If None, the ratio and pad will be calculated based on the size difference between the two images. Defaults to None.

  • padding (bool, optional) – If True, assuming the boxes is based on image augmented by yolo style. If False then do regular rescaling. Defaults to True.

  • xywh (bool, optional) – The box format is xywh or not. Defaults to False.

Returns:

The scaled bounding boxes, in the format of (x1, y1, x2, y2)

Return type:

boxes (torch.Tensor)

flat_bug.yolo_helpers.stack_masks(masks: List[Masks], orig_shape: Tuple[int, int] | List[int] | None = None) Masks

Stacks a list of ultralytics.engine.results.Masks objects (or torch.Tensor) into a single ultralytics.engine.results.Masks object.

If the masks are not all the same size, they are resized to the largest size in the list.

Parameters:
  • masks (list) – A list of ultralytics.engine.results.Masks objects (or torch.Tensor).

  • orig_shape (tuple, optional) – The original shape of the image. Defaults to None. If None, the original shape is inferred from the first Masks object in the list if there is one, otherwise the original shape None.

  • antialias (bool, optional) – A flag to indicate whether to use antialiasing when resizing the masks. Defaults to False.

Returns:

A Masks object containing the stacked masks.

Return type:

out (ultralytics.engine.results.Masks)