torchwebio.models.image package

Submodules

torchwebio.models.image.imageclassificationmodel module

class torchwebio.models.image.imageclassificationmodel.ImageClassificationModel(model_name: Optional[str] = '', number_of_results: Optional[int] = 5, model: Optional[Module] = None, category_list: Optional[str] = None)[source]

Bases: object

An adapter class to expose all Image classification functionality for models that follow the interface of pytorch-image-models (timm).

In a nutshell, this class exposes the following functionality : pre-process, forward and post-process. The process method does all three at a go, and is pretty much the only method that needs to be called externally.

Additionally, there are some helper static methods that are grouped here only for encapsulation purposes

forward(img_tensor: Tensor) → Tensor[source]

Forward pass of the pytorch model

Parameters

img_tensor (Tensor) – input image after pre-processing.

Returns

Output probabilities from the network

Return type

Tensor

post_process(probabilities: Tensor) → List[Tuple[str, float]][source]

Post process the soft-max output of the model. Extracts top-K results from the last layer,

and maps them to the corresponding categories.

Parameters

probabilities (Tensor) – Tensor of probabilities (output of the network)

Returns

List of pairs of Category label and category score

Return type

List[Tuple[str, float]]

pre_process(img: <module 'PIL.Image' from '/home/docs/checkouts/readthedocs.org/user_builds/torchwebio/envs/latest/lib/python3.8/site-packages/PIL/Image.py'>) → Tensor[source]

Preprocess the image. Basically processes the image through the transformations. Assumes batch size = 1

Parameters

img (Image) – Image to pre-process.

Returns

pre-processed image as a tensor

Return type

Tensor

process_img(img) → List[Tuple[str, float]][source]

Processes a single image. Calls pre_process, forward and post_process in succession

Parameters

img (PIL.Image) – Input image as a PIL object

Returns

Output scores and categories

Return type

List[Tuple[str, float]]

Module contents