tenflowers-dataset 0.1.1

Data pipeline and dataset utilities for TenfloweRS
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
59
60
61
62

//! GPU-accelerated image rotation transform
//!
//! This module provides GPU-accelerated image rotation using WGPU compute shaders
//! for data augmentation and geometric transformations.

use tenflowers_core::{Result, Tensor, TensorError};

#[cfg(feature = "gpu")]
use std::sync::Arc;

#[cfg(feature = "gpu")]
use crate::Transform;

#[cfg(feature = "gpu")]
use super::context::GpuContext;

/// GPU-accelerated rotation transform
#[cfg(feature = "gpu")]
pub struct GpuRotation {
    angle_range: (f32, f32),
    context: Arc<GpuContext>,
}

#[cfg(feature = "gpu")]
impl GpuRotation {
    /// Create a new GPU rotation transform
    pub fn new(angle_range: (f32, f32), context: Arc<GpuContext>) -> Result<Self> {
        Ok(Self {
            angle_range,
            context,
        })
    }

    /// Apply rotation to image tensor using GPU
    pub async fn rotate_tensor(&self, input: &Tensor<f32>) -> Result<Tensor<f32>> {
        // Placeholder implementation - would contain full GPU rotation logic
        Ok(input.clone())
    }
}

#[cfg(feature = "gpu")]
impl Transform<f32> for GpuRotation {
    fn apply(&self, sample: (Tensor<f32>, Tensor<f32>)) -> Result<(Tensor<f32>, Tensor<f32>)> {
        let (image_tensor, label_tensor) = sample;
        let rotated_tensor = pollster::block_on(self.rotate_tensor(&image_tensor))?;
        Ok((rotated_tensor, label_tensor))
    }
}

/// CPU fallback rotation transform when GPU is not available
#[cfg(not(feature = "gpu"))]
pub struct GpuRotation;

#[cfg(not(feature = "gpu"))]
impl GpuRotation {
    /// Create a new rotation transform (fallback to CPU)
    pub fn new(_angle_range: (f32, f32), _context: ()) -> Result<Self> {
        Err(TensorError::unsupported_operation_simple(
            "GPU transforms require 'gpu' feature to be enabled".to_string(),
        ))
    }
}