Metadata-Version: 2.3
Name: data-processing-kit
Version: 0.1.1
Summary: A kit for managing data in a pytorch training/inference pipeline
Requires-Dist: packaging>=26.2
Requires-Dist: pillow>=12.2.0
Requires-Dist: pydantic>=2.13.3
Requires-Dist: safetensors>=0.7.0
Requires-Dist: torch>=2.11.0
Requires-Dist: torchvision>=0.26.0
Requires-Python: >=3.14
Description-Content-Type: text/markdown

# 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](https://github.com/astral-sh/uv) (recommended) or pip

### Using uv (recommended)
```bash
uv add data-processing-kit
```

### Using pip
```bash
pip install data-processing-kit 
```

## Usage

### ImageProcessor

```python
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

```python
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

```python
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](LICENSE) for details.
