data-processing-kit 0.1.1

A kit for managing data in a pytorch training/inference pipeline

Data Processing Kit

A kit for managing data in a PyTorch training/inference pipeline.

Features

• ImageProcessor: Configure and apply image transformation pipelines for training and inference with separate transforms for training (with augmentation) and evaluation
• TransformDataset: Generic dataset wrapper that applies transformations with type safety and automatic class information extraction
• CheckpointManager: Safe model checkpointing using safetensors format to avoid pickle security issues

Installation

Prerequisites

• Python 3.14 or higher
• uv (recommended) or pip

Using uv (recommended)

uv add data-processing-kit

Using pip

pip install data-processing-kit 

Usage

ImageProcessor

import torch
import torchvision.transforms.v2 as v2
from PIL import Image
from data_processing_kit import ImageProcessor

processor = ImageProcessor(
    train_transforms=[
        v2.Resize((224, 224)),
        v2.RandomHorizontalFlip(),
        v2.ToImage(),
        v2.ToDtype(torch.float32, scale=True),
        v2.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    ],
    eval_transforms=[
        v2.Resize((224, 224)),
        v2.ToImage(),
        v2.ToDtype(torch.float32, scale=True),
        v2.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    ]
)

# Process single image during inference
image = Image.open("example.jpg")
tensor = processor.inference_process_single(image, device=torch.device("cpu"))

# Process batch or stream
images = [Image.open(f"example-{i}.jpg") for i in range(10)] # Open 10 images
tensor_batch = processor.inference_process_batch(images, device=torch.device("cuda"))
tensor_stream = processor.inference_process_stream(images, device=torch.device("cuda"))

# Save processor to a file (or use CheckpointManager, see below)
processor.save("processor.pt")

TransformDataset

from data_processing_kit import TransformDataset
from torchvision.datasets import ImageFolder

dataset = ImageFolder("path/to/data")
transform_dataset = TransformDataset(
    dataset=dataset,
    transform=processor.get_train_compose(),
    classes=["cat", "dog"]
)

print(f"Number of classes: {transform_dataset.num_classes}")
print(f"Class mapping: {transform_dataset.class_to_idx}")

CheckpointManager

import torch
from data_processing_kit import CheckpointManager

manager = CheckpointManager()

# Save model and processor
model = YourModel()
manager.save(model.state_dict(), processor, save_directory="checkpoints/my_model")

# Load checkpoint
model_state, loaded_processor = manager.load("checkpoints/my_model")
model.load_state_dict(model_state)

License

MIT License - see LICENSE for details.