Skip to content

torch

torch

Attributes

__all__ module-attribute 

__all__ = ['ClassificationPipeline', 'ClassificationMixin']

Classes

ClassificationMixin

Shared classification preprocessing/postprocessing logic.

Attributes:

Name Type Description
image_size tuple[int, int]

Target image size.
mean Tensor

Normalization mean.
std Tensor

Normalization std.
class_names dict[int, str]

Mapping from class ID to class name.
transform Compose

Image transformation pipeline.
Attributes
image_size instance-attribute 
image_size: tuple[int, int]
mean instance-attribute 
mean: Tensor
std instance-attribute 
std: Tensor
class_names instance-attribute 
class_names: dict[int, str]
transform instance-attribute 
transform: Compose

ClassificationPipeline 

ClassificationPipeline(runtime: Runtime[Tensor, Tensor], image_size: tuple[int, int] = (224, 224), mean: tuple[float, float, float] = (0.485, 0.456, 0.406), std: tuple[float, float, float] = (0.229, 0.224, 0.225), class_names: dict[int, str] | None = None, batch_strategy: BatchStrategy[Tensor, Tensor] | None = None)

Bases: ClassificationMixin, Pipeline[Tensor, Tensor, ClassificationOutput]

Image classification pipeline (sync version).

Performs
  • Image decoding and conversion
  • Resizing and normalization
  • Model inference
  • Class prediction with confidence

Parameters:

Name Type Description Default
runtime Runtime[Tensor, Tensor]

Inference runtime.

 required
image_size tuple[int, int]

Target image size (default: (224, 224)).

 (224, 224)
mean tuple[float, float, float]

Normalization mean (default: ImageNet mean).

 (0.485, 0.456, 0.406)
std tuple[float, float, float]

Normalization std (default: ImageNet std).

 (0.229, 0.224, 0.225)
class_names dict[int, str] | None

Optional mapping from class ID to class name.

 None
batch_strategy BatchStrategy[Tensor, Tensor] | None

Optional batching strategy.

 None
Example 
runtime = TorchScriptRuntime(
    model_path="resnet50.pt", device="cuda"
)
pipeline = ClassificationPipeline(
    runtime=runtime,
    class_names={0: "cat", 1: "dog", 2: "bird"},
)

with pipeline.serve():
    result = pipeline(image_bytes)
    print(f"{result.class_name}: {result.confidence:.2%}")
Source code in inferflow/pipeline/classification/torch.py 
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
def __init__(
    self,
    runtime: Runtime[torch.Tensor, torch.Tensor],
    image_size: tuple[int, int] = (224, 224),
    mean: tuple[float, float, float] = (0.485, 0.456, 0.406),
    std: tuple[float, float, float] = (0.229, 0.224, 0.225),
    class_names: dict[int, str] | None = None,
    batch_strategy: BatchStrategy[torch.Tensor, torch.Tensor] | None = None,
):
    super().__init__(runtime=runtime, batch_strategy=batch_strategy)

    self.image_size = image_size
    self.mean = torch.tensor(mean).view(3, 1, 1)
    self.std = torch.tensor(std).view(3, 1, 1)
    self.class_names = class_names or {}

    self.transform = T.Compose([T.Resize(image_size), T.ToTensor()])
Attributes
image_size instance-attribute 
image_size = image_size
mean instance-attribute 
mean = view(3, 1, 1)
std instance-attribute 
std = view(3, 1, 1)
class_names instance-attribute 
class_names = class_names or {}
transform instance-attribute 
transform = Compose([Resize(image_size), ToTensor()])
Functions
preprocess 
preprocess(input: ImageInput) -> Tensor

Preprocess image input.

Parameters:

Name Type Description Default
input ImageInput

Raw image input (bytes, numpy array, PIL Image, or tensor).

 required

Returns:

Type Description
Tensor

Preprocessed tensor ready for model inference.
Source code in inferflow/pipeline/classification/torch.py 
173
174
175
176
177
178
179
180
181
182
183
184
185
186
def preprocess(self, input: ImageInput) -> torch.Tensor:
    """Preprocess image input.

    Args:
        input: Raw image input (bytes, numpy array, PIL Image, or tensor).

    Returns:
        Preprocessed tensor ready for model inference.
    """
    if isinstance(input, torch.Tensor):
        return self._preprocess_tensor(input)

    image = self._convert_to_pil(input)
    return self._preprocess_image(image)
postprocess 
postprocess(raw: Tensor) -> ClassificationOutput

Postprocess model output to classification result.

Parameters:

Name Type Description Default
raw Tensor

Raw output tensor from model inference.

 required

Returns:

Type Description
ClassificationOutput

ClassificationOutput with class ID, confidence, and optional class name.
Source code in inferflow/pipeline/classification/torch.py 
188
189
190
191
192
193
194
195
196
197
def postprocess(self, raw: torch.Tensor) -> ClassificationOutput:
    """Postprocess model output to classification result.

    Args:
        raw: Raw output tensor from model inference.

    Returns:
        ClassificationOutput with class ID, confidence, and optional class name.
    """
    return self._postprocess_raw(raw)