hawk_eye.core.classifier

Overview

Image classification aims to assign a class to an image. For UAV, we are focued on determining if an image is either “background” or “target.” A background image does not contain one of the AUVSI SUAS targets that we are interested. A image classified as a target should have one of these targets. Note, in this model, we don’t actually determine the specific target, that is left to a detector. In this way, the classifier works as a filter to feed only taget-containing images into the object detector.

UAV Austin uses a fast classifier to speed up processing time, saving more room for the slower object detector. As technology improves, it might be possible to remove the classifier from the step completely and solely rely on the detector to process each image.

Module

A classifier model which wraps around a backbone. This setup allows for easy interchangeability during experimentation and a reliable way to load saved models.

class hawk_eye.core.classifier.Classifier(num_classes: Optional[int] = 2, timestamp: Optional[str] = None, backbone: Optional[str] = None, half_precision: Optional[bool] = False)[source]
Parameters

  • num_classes – The number of classes to predict.

  • timestamp – The timestamp of the model to download from GCloud.

  • backbone – A string designating which model to load.

  • half_precision – Whether to use half precision. This should be False for training but True during inference.

Raises

ValueError – Error if neither a timestamp or backbone arg is supplied.

Examples

>>> classifier = Classifier(2, backbone="vovnet-19")
>>> with torch.no_grad():
...    predictions = classifier.classify(torch.randn(1, 3, 64, 64), True)
>>> predictions.shape
torch.Size([1, 2])
classify(x: torch.Tensor, probability: bool = False) → torch.Tensor[source]

Take in an image batch and return the class for each image. If specified, softmax will be applied to the predictions.

Parameters

  • x – Input tensor of size (batch, height, width, channels).

  • probability – Whether or not to apply softmax.

Returns

The output tensor.

forward(x: torch.Tensor) → torch.Tensor[source]

Forward pass through classifier.

Parameters

x – input tensor.

Returns

the output tensor.

load_backbone(backbone: str) → torch.nn.modules.module.Module[source]

Load the supplied backbone. See this function for the list of potential backbones that can be loaded.

Parameters

backbone – The backbone type to load.

Returns

The loaded model.

Raises

ValueError – If improper backbone is supplied.