dataloader-rs 0.3.0

High-performance DataLoader for Rust with a PyTorch-like interface
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58

//! The `#[pymodule]`: the classes and functions of the `dataloader_rs` Python module.

use pyo3::prelude::*;

use crate::python::{
    dataloader::PyDataloader,
    dataset::{
        PyAsyncDatasetBase, PyDatasetBase, PyFileDataset, PyHttpDataset, PyObjectStoreDataset,
        PyRustDataset,
    },
    image::{self, PyImage, PyImageBatch, PyImageCollator, PyImageTransform},
    iterator::PyDataloaderIter,
    pool,
};

#[pymodule(gil_used = false)]
fn dataloader_rs(m: &Bound<'_, PyModule>) -> PyResult<()> {
    // PyO3 asserts, on the first attach from a thread that is not attached, that the
    // interpreter is initialized. Do that first attach here: a pool thread attaching while
    // the interpreter finishes shutting down would panic on it instead of being refused.
    m.py().detach(|| Python::try_attach(|_| ()));
    m.add_class::<PyDataloader>()?;
    m.add_class::<PyDataloaderIter>()?;
    m.add_class::<PyDatasetBase>()?;
    m.add_class::<PyAsyncDatasetBase>()?;
    m.add_class::<PyRustDataset>()?;
    m.add_class::<PyHttpDataset>()?;
    m.add_class::<PyFileDataset>()?;
    m.add_class::<PyObjectStoreDataset>()?;
    m.add_function(wrap_pyfunction!(pool::keep_pools_alive, m)?)?;
    m.add_class::<PyImage>()?;
    m.add_class::<PyImageBatch>()?;
    m.add_class::<PyImageCollator>()?;
    m.add_class::<PyImageTransform>()?;
    // `dataloader_rs.transforms` re-exports these from the extension module.
    for transform in [
        wrap_pyfunction!(image::decode, m)?,
        wrap_pyfunction!(image::to_rgb, m)?,
        wrap_pyfunction!(image::to_rgba, m)?,
        wrap_pyfunction!(image::to_gray, m)?,
        wrap_pyfunction!(image::crop, m)?,
        wrap_pyfunction!(image::center_crop, m)?,
        wrap_pyfunction!(image::pad, m)?,
        wrap_pyfunction!(image::pad_to, m)?,
        wrap_pyfunction!(image::resize, m)?,
        wrap_pyfunction!(image::hflip, m)?,
        wrap_pyfunction!(image::vflip, m)?,
        wrap_pyfunction!(image::rotate90, m)?,
        wrap_pyfunction!(image::to_float, m)?,
        wrap_pyfunction!(image::normalize, m)?,
        wrap_pyfunction!(image::to_chw, m)?,
        wrap_pyfunction!(image::to_hwc, m)?,
        wrap_pyfunction!(image::compose, m)?,
    ] {
        m.add_function(transform)?;
    }
    Ok(())
}