Initial commit
This commit is contained in:
redclawsystems
@@ -0,0 +1,723 @@
|
||||
use crate::error::{MultimodalError, Result};
|
||||
use crate::vision::vit::{LayerNorm, MultiHeadAttention, TransformerBlock, TransformerConfig};
|
||||
use rtx_tensor::Device;
|
||||
use rtx_tensor::Tensor;
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct TimeSformerConfig {
|
||||
pub image_size: usize,
|
||||
pub patch_size: usize,
|
||||
pub in_channels: usize,
|
||||
pub embed_dim: usize,
|
||||
pub depth: usize,
|
||||
pub num_heads: usize,
|
||||
pub mlp_ratio: f32,
|
||||
pub num_frames: usize,
|
||||
pub num_classes: usize,
|
||||
pub dropout: f32,
|
||||
pub attention_dropout: f32,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct DividedSpaceTimeAttentionConfig {
|
||||
pub embed_dim: usize,
|
||||
pub num_heads: usize,
|
||||
pub dropout: f32,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct TimeSformerBlockConfig {
|
||||
pub embed_dim: usize,
|
||||
pub num_heads: usize,
|
||||
pub mlp_ratio: f32,
|
||||
pub dropout: f32,
|
||||
pub attention_dropout: f32,
|
||||
}
|
||||
|
||||
pub struct VideoPatchEmbedding {
|
||||
conv: Tensor, // 3D convolution weights
|
||||
bias: Tensor,
|
||||
num_patches: usize,
|
||||
embed_dim: usize,
|
||||
num_frames: usize,
|
||||
device: Device,
|
||||
}
|
||||
|
||||
impl VideoPatchEmbedding {
|
||||
pub fn new(
|
||||
image_size: usize,
|
||||
patch_size: usize,
|
||||
in_channels: usize,
|
||||
embed_dim: usize,
|
||||
num_frames: usize,
|
||||
device: &Device,
|
||||
) -> Result<Self> {
|
||||
if !image_size.is_multiple_of(patch_size) {
|
||||
return Err(MultimodalError::InvalidPatchSize {
|
||||
patch_size,
|
||||
image_size,
|
||||
});
|
||||
}
|
||||
|
||||
let num_patches = (image_size / patch_size).pow(2);
|
||||
|
||||
// 2D convolution for spatial patches (we handle temporal dimension separately)
|
||||
let fan_out = embed_dim * patch_size * patch_size;
|
||||
let bound = (6.0 / (in_channels * patch_size * patch_size + fan_out) as f32).sqrt();
|
||||
|
||||
// Use randn and scale to [-bound, bound] range
|
||||
let conv = Tensor::randn(&[embed_dim, in_channels, patch_size, patch_size], device)?
|
||||
.mul_scalar(bound / 3.0)?; // Scale by bound/3 for similar distribution
|
||||
|
||||
let bias = Tensor::zeros([embed_dim], device)?;
|
||||
|
||||
Ok(Self {
|
||||
conv,
|
||||
bias,
|
||||
num_patches,
|
||||
embed_dim,
|
||||
num_frames,
|
||||
device: device.clone(),
|
||||
})
|
||||
}
|
||||
|
||||
pub fn num_patches(&self) -> usize {
|
||||
self.num_patches
|
||||
}
|
||||
|
||||
pub fn embed_dim(&self) -> usize {
|
||||
self.embed_dim
|
||||
}
|
||||
|
||||
pub fn forward(&self, x: &Tensor) -> Result<Tensor> {
|
||||
// x shape: [batch_size, channels, frames, height, width]
|
||||
let _batch_size = x.shape()[0];
|
||||
let _channels = x.shape()[1];
|
||||
let frames = x.shape()[2];
|
||||
let _height = x.shape()[3];
|
||||
let _width = x.shape()[4];
|
||||
|
||||
// Process each frame separately
|
||||
let mut frame_patches = Vec::new();
|
||||
for t in 0..frames {
|
||||
// Extract frame: [batch, channels, height, width]
|
||||
let frame = x.slice(2, t, t + 1)?.squeeze(Some(2))?;
|
||||
|
||||
// Apply 2D convolution to create patches
|
||||
let patches = frame.conv2d(&self.conv, Some(&self.bias), 1, 0, 1, 1)?;
|
||||
|
||||
// Flatten spatial dimensions
|
||||
let batch_size = patches.shape()[0];
|
||||
let embed_dim = patches.shape()[1];
|
||||
let h_patches = patches.shape()[2];
|
||||
let w_patches = patches.shape()[3];
|
||||
|
||||
let flattened = patches.reshape([batch_size, embed_dim, h_patches * w_patches])?;
|
||||
let transposed = flattened.transpose(1, 2)?; // [batch, num_patches, embed_dim]
|
||||
|
||||
frame_patches.push(transposed);
|
||||
}
|
||||
|
||||
// Stack frame patches: [batch, frames, num_patches, embed_dim]
|
||||
let video_patches = Tensor::stack(&frame_patches, 1)?;
|
||||
|
||||
Ok(video_patches)
|
||||
}
|
||||
}
|
||||
|
||||
pub struct TemporalPositionalEncoding {
|
||||
embeddings: Tensor,
|
||||
max_frames: usize,
|
||||
device: Device,
|
||||
}
|
||||
|
||||
impl TemporalPositionalEncoding {
|
||||
pub fn new(max_frames: usize, embed_dim: usize, device: &Device) -> Result<Self> {
|
||||
let embeddings = Tensor::randn(&[1, max_frames, embed_dim], device)?;
|
||||
|
||||
Ok(Self {
|
||||
embeddings,
|
||||
max_frames,
|
||||
device: device.clone(),
|
||||
})
|
||||
}
|
||||
|
||||
pub fn forward(&self, num_frames: usize) -> Result<Tensor> {
|
||||
if num_frames > self.max_frames {
|
||||
return Err(MultimodalError::InvalidSequenceLength {
|
||||
seq_len: num_frames,
|
||||
max_len: self.max_frames,
|
||||
});
|
||||
}
|
||||
|
||||
let indices = Tensor::arange(0, num_frames as i64, &self.device)?
|
||||
.unsqueeze(0)?
|
||||
.unsqueeze(2)?;
|
||||
|
||||
Ok(self.embeddings.gather(1, &indices)?)
|
||||
}
|
||||
}
|
||||
|
||||
pub struct SpatialPositionalEncoding {
|
||||
embeddings: Tensor,
|
||||
max_patches: usize,
|
||||
device: Device,
|
||||
}
|
||||
|
||||
impl SpatialPositionalEncoding {
|
||||
pub fn new(max_patches: usize, embed_dim: usize, device: &Device) -> Result<Self> {
|
||||
let embeddings = Tensor::randn(&[1, max_patches, embed_dim], device)?;
|
||||
|
||||
Ok(Self {
|
||||
embeddings,
|
||||
max_patches,
|
||||
device: device.clone(),
|
||||
})
|
||||
}
|
||||
|
||||
pub fn forward(&self, num_patches: usize) -> Result<Tensor> {
|
||||
if num_patches > self.max_patches {
|
||||
return Err(MultimodalError::InvalidSequenceLength {
|
||||
seq_len: num_patches,
|
||||
max_len: self.max_patches,
|
||||
});
|
||||
}
|
||||
|
||||
let indices = Tensor::arange(0, num_patches as i64, &self.device)?
|
||||
.unsqueeze(0)?
|
||||
.unsqueeze(2)?;
|
||||
|
||||
Ok(self.embeddings.gather(1, &indices)?)
|
||||
}
|
||||
}
|
||||
|
||||
pub struct SpaceOnlyAttention {
|
||||
attention: MultiHeadAttention,
|
||||
device: Device,
|
||||
}
|
||||
|
||||
impl SpaceOnlyAttention {
|
||||
pub fn new(embed_dim: usize, num_heads: usize, _dropout: f32, device: &Device) -> Result<Self> {
|
||||
let attention = MultiHeadAttention::new(embed_dim, num_heads, device)?;
|
||||
|
||||
Ok(Self {
|
||||
attention,
|
||||
device: device.clone(),
|
||||
})
|
||||
}
|
||||
|
||||
pub fn forward(&self, x: &Tensor) -> Result<Tensor> {
|
||||
// x shape: [batch, frames, patches, embed_dim]
|
||||
let _batch_size = x.shape()[0];
|
||||
let frames = x.shape()[1];
|
||||
let _patches = x.shape()[2];
|
||||
let _embed_dim = x.shape()[3];
|
||||
|
||||
// Process each frame independently
|
||||
let mut frame_outputs = Vec::new();
|
||||
for t in 0..frames {
|
||||
let frame = x.slice(1, t, t + 1)?.squeeze(Some(1))?; // [batch, patches, embed_dim]
|
||||
let frame_out = self.attention.forward(&frame)?;
|
||||
frame_outputs.push(frame_out);
|
||||
}
|
||||
|
||||
// Stack outputs: [batch, frames, patches, embed_dim]
|
||||
let output = Tensor::stack(&frame_outputs, 1)?;
|
||||
|
||||
Ok(output)
|
||||
}
|
||||
}
|
||||
|
||||
pub struct TimeOnlyAttention {
|
||||
attention: MultiHeadAttention,
|
||||
device: Device,
|
||||
}
|
||||
|
||||
impl TimeOnlyAttention {
|
||||
pub fn new(embed_dim: usize, num_heads: usize, _dropout: f32, device: &Device) -> Result<Self> {
|
||||
let attention = MultiHeadAttention::new(embed_dim, num_heads, device)?;
|
||||
|
||||
Ok(Self {
|
||||
attention,
|
||||
device: device.clone(),
|
||||
})
|
||||
}
|
||||
|
||||
pub fn forward(&self, x: &Tensor) -> Result<Tensor> {
|
||||
// x shape: [batch, frames, patches, embed_dim]
|
||||
let _batch_size = x.shape()[0];
|
||||
let _frames = x.shape()[1];
|
||||
let patches = x.shape()[2];
|
||||
let _embed_dim = x.shape()[3];
|
||||
|
||||
// Process each spatial location independently across time
|
||||
let mut patch_outputs = Vec::new();
|
||||
for p in 0..patches {
|
||||
let patch_temporal = x.slice(2, p, p + 1)?.squeeze(Some(2))?; // [batch, frames, embed_dim]
|
||||
let patch_out = self.attention.forward(&patch_temporal)?;
|
||||
patch_outputs.push(patch_out);
|
||||
}
|
||||
|
||||
// Stack outputs and permute back: [batch, patches, frames, embed_dim] -> [batch, frames, patches, embed_dim]
|
||||
let stacked = Tensor::stack(&patch_outputs, 1)?; // [batch, patches, frames, embed_dim]
|
||||
let output = stacked.transpose(1, 2)?; // [batch, frames, patches, embed_dim]
|
||||
|
||||
Ok(output)
|
||||
}
|
||||
}
|
||||
|
||||
pub struct JointSpaceTimeAttention {
|
||||
attention: MultiHeadAttention,
|
||||
device: Device,
|
||||
}
|
||||
|
||||
impl JointSpaceTimeAttention {
|
||||
pub fn new(embed_dim: usize, num_heads: usize, _dropout: f32, device: &Device) -> Result<Self> {
|
||||
let attention = MultiHeadAttention::new(embed_dim, num_heads, device)?;
|
||||
|
||||
Ok(Self {
|
||||
attention,
|
||||
device: device.clone(),
|
||||
})
|
||||
}
|
||||
|
||||
pub fn forward(&self, x: &Tensor) -> Result<Tensor> {
|
||||
// x shape: [batch, frames, patches, embed_dim]
|
||||
let batch_size = x.shape()[0];
|
||||
let frames = x.shape()[1];
|
||||
let patches = x.shape()[2];
|
||||
let embed_dim = x.shape()[3];
|
||||
|
||||
// Flatten spatial-temporal dimensions for joint attention
|
||||
let x_flat = x.reshape([batch_size, frames * patches, embed_dim])?;
|
||||
let out_flat = self.attention.forward(&x_flat)?;
|
||||
|
||||
// Reshape back to original dimensions
|
||||
let output = out_flat.reshape([batch_size, frames, patches, embed_dim])?;
|
||||
|
||||
Ok(output)
|
||||
}
|
||||
}
|
||||
|
||||
pub struct DividedSpaceTimeAttention {
|
||||
space_attention: SpaceOnlyAttention,
|
||||
time_attention: TimeOnlyAttention,
|
||||
device: Device,
|
||||
}
|
||||
|
||||
impl DividedSpaceTimeAttention {
|
||||
pub fn new(config: &DividedSpaceTimeAttentionConfig, device: &Device) -> Result<Self> {
|
||||
let space_attention =
|
||||
SpaceOnlyAttention::new(config.embed_dim, config.num_heads, config.dropout, device)?;
|
||||
|
||||
let time_attention =
|
||||
TimeOnlyAttention::new(config.embed_dim, config.num_heads, config.dropout, device)?;
|
||||
|
||||
Ok(Self {
|
||||
space_attention,
|
||||
time_attention,
|
||||
device: device.clone(),
|
||||
})
|
||||
}
|
||||
|
||||
pub fn forward(&self, x: &Tensor) -> Result<Tensor> {
|
||||
// Apply spatial attention first
|
||||
let x = self.space_attention.forward(x)?;
|
||||
|
||||
// Then apply temporal attention
|
||||
let x = self.time_attention.forward(&x)?;
|
||||
|
||||
Ok(x)
|
||||
}
|
||||
}
|
||||
|
||||
pub struct TimeSformerBlock {
|
||||
temporal_attn: TimeOnlyAttention,
|
||||
spatial_attn: SpaceOnlyAttention,
|
||||
mlp: TransformerBlock, // Reuse the MLP from TransformerBlock
|
||||
ln1: LayerNorm,
|
||||
ln2: LayerNorm,
|
||||
ln3: LayerNorm,
|
||||
device: Device,
|
||||
}
|
||||
|
||||
impl TimeSformerBlock {
|
||||
pub fn new(config: &TimeSformerBlockConfig, device: &Device) -> Result<Self> {
|
||||
let temporal_attn = TimeOnlyAttention::new(
|
||||
config.embed_dim,
|
||||
config.num_heads,
|
||||
config.attention_dropout,
|
||||
device,
|
||||
)?;
|
||||
|
||||
let spatial_attn = SpaceOnlyAttention::new(
|
||||
config.embed_dim,
|
||||
config.num_heads,
|
||||
config.attention_dropout,
|
||||
device,
|
||||
)?;
|
||||
|
||||
let mlp_config = TransformerConfig {
|
||||
embed_dim: config.embed_dim,
|
||||
num_heads: config.num_heads,
|
||||
mlp_ratio: config.mlp_ratio,
|
||||
dropout: config.dropout,
|
||||
attention_dropout: config.attention_dropout,
|
||||
};
|
||||
let mlp = TransformerBlock::new(&mlp_config, device)?;
|
||||
|
||||
let ln1 = LayerNorm::new(config.embed_dim, device)?;
|
||||
let ln2 = LayerNorm::new(config.embed_dim, device)?;
|
||||
let ln3 = LayerNorm::new(config.embed_dim, device)?;
|
||||
|
||||
Ok(Self {
|
||||
temporal_attn,
|
||||
spatial_attn,
|
||||
mlp,
|
||||
ln1,
|
||||
ln2,
|
||||
ln3,
|
||||
device: device.clone(),
|
||||
})
|
||||
}
|
||||
|
||||
pub fn forward(&self, x: &Tensor) -> Result<Tensor> {
|
||||
// Temporal attention
|
||||
let temp_out = self.temporal_attn.forward(&self.ln1.forward(x)?)?;
|
||||
let x = (x + &temp_out)?;
|
||||
|
||||
// Spatial attention
|
||||
let spat_out = self.spatial_attn.forward(&self.ln2.forward(&x)?)?;
|
||||
let x = (&x + &spat_out)?;
|
||||
|
||||
// We need to reshape for the MLP which expects [batch, seq, embed]
|
||||
let batch_size = x.shape()[0];
|
||||
let frames = x.shape()[1];
|
||||
let patches = x.shape()[2];
|
||||
let embed_dim = x.shape()[3];
|
||||
|
||||
let x_flat = x.reshape([batch_size, frames * patches, embed_dim])?;
|
||||
let mlp_out = self.mlp.forward(&x_flat)?;
|
||||
let mlp_out = mlp_out.reshape([batch_size, frames, patches, embed_dim])?;
|
||||
|
||||
Ok(mlp_out)
|
||||
}
|
||||
}
|
||||
|
||||
pub struct TimeSformerModel {
|
||||
patch_embed: VideoPatchEmbedding,
|
||||
cls_token: Tensor,
|
||||
temporal_pos_embed: TemporalPositionalEncoding,
|
||||
spatial_pos_embed: SpatialPositionalEncoding,
|
||||
blocks: Vec<TimeSformerBlock>,
|
||||
ln_f: LayerNorm,
|
||||
head: Tensor,
|
||||
bias: Tensor,
|
||||
config: TimeSformerConfig,
|
||||
device: Device,
|
||||
}
|
||||
|
||||
impl TimeSformerModel {
|
||||
pub fn new(config: &TimeSformerConfig, device: &Device) -> Result<Self> {
|
||||
let patch_embed = VideoPatchEmbedding::new(
|
||||
config.image_size,
|
||||
config.patch_size,
|
||||
config.in_channels,
|
||||
config.embed_dim,
|
||||
config.num_frames,
|
||||
device,
|
||||
)?;
|
||||
|
||||
// Class token for each frame
|
||||
let cls_token = Tensor::randn(&[1, config.num_frames, 1, config.embed_dim], device)?;
|
||||
|
||||
let temporal_pos_embed =
|
||||
TemporalPositionalEncoding::new(config.num_frames, config.embed_dim, device)?;
|
||||
|
||||
let spatial_pos_embed = SpatialPositionalEncoding::new(
|
||||
patch_embed.num_patches() + 1, // +1 for cls token
|
||||
config.embed_dim,
|
||||
device,
|
||||
)?;
|
||||
|
||||
let block_config = TimeSformerBlockConfig {
|
||||
embed_dim: config.embed_dim,
|
||||
num_heads: config.num_heads,
|
||||
mlp_ratio: config.mlp_ratio,
|
||||
dropout: config.dropout,
|
||||
attention_dropout: config.attention_dropout,
|
||||
};
|
||||
|
||||
let mut blocks = Vec::new();
|
||||
for _ in 0..config.depth {
|
||||
blocks.push(TimeSformerBlock::new(&block_config, device)?);
|
||||
}
|
||||
|
||||
let ln_f = LayerNorm::new(config.embed_dim, device)?;
|
||||
let head = Tensor::randn(&[config.num_classes, config.embed_dim], device)?;
|
||||
let bias = Tensor::zeros([config.num_classes], device)?;
|
||||
|
||||
Ok(Self {
|
||||
patch_embed,
|
||||
cls_token,
|
||||
temporal_pos_embed,
|
||||
spatial_pos_embed,
|
||||
blocks,
|
||||
ln_f,
|
||||
head,
|
||||
bias,
|
||||
config: config.clone(),
|
||||
device: device.clone(),
|
||||
})
|
||||
}
|
||||
|
||||
pub fn forward(&self, x: &Tensor) -> Result<Tensor> {
|
||||
let batch_size = x.shape()[0];
|
||||
|
||||
// Video patch embedding: [batch, frames, patches, embed_dim]
|
||||
let patches = self.patch_embed.forward(x)?;
|
||||
|
||||
// Add cls tokens for each frame
|
||||
let cls_tokens = self.cls_token.clone().expand(&[
|
||||
batch_size,
|
||||
self.config.num_frames,
|
||||
1,
|
||||
self.config.embed_dim,
|
||||
])?;
|
||||
let x = Tensor::cat(&[cls_tokens, patches], 2)?; // Concat along patch dimension
|
||||
|
||||
// Add positional encodings
|
||||
let spatial_pos = self.spatial_pos_embed.forward(x.shape()[2])?; // [1, patches+1, embed_dim]
|
||||
let temporal_pos = self.temporal_pos_embed.forward(self.config.num_frames)?; // [1, frames, embed_dim]
|
||||
|
||||
// Broadcast and add positional encodings
|
||||
let spatial_pos_expanded = spatial_pos.unsqueeze(1)?.expand(&[
|
||||
batch_size,
|
||||
self.config.num_frames,
|
||||
x.shape()[2],
|
||||
self.config.embed_dim,
|
||||
])?;
|
||||
let temporal_pos_expanded = temporal_pos.unsqueeze(2)?.expand(&[
|
||||
batch_size,
|
||||
self.config.num_frames,
|
||||
x.shape()[2],
|
||||
self.config.embed_dim,
|
||||
])?;
|
||||
|
||||
let x = ((&x + &spatial_pos_expanded)? + temporal_pos_expanded)?;
|
||||
|
||||
// Pass through TimeSformer blocks
|
||||
let mut x = x;
|
||||
for block in &self.blocks {
|
||||
x = block.forward(&x)?;
|
||||
}
|
||||
|
||||
// Final layer norm
|
||||
let x = self.ln_f.forward(&x)?;
|
||||
|
||||
// Classification head (use cls tokens)
|
||||
let cls_outputs = x.slice(2, 0, 1)?.squeeze(Some(2))?; // [batch, frames, embed_dim]
|
||||
let cls_pooled = cls_outputs.mean(&[1], false)?; // Pool across time: [batch, embed_dim]
|
||||
let logits = cls_pooled
|
||||
.matmul(&self.head.transpose(0, 1)?)?
|
||||
.add(&self.bias)?;
|
||||
|
||||
Ok(logits)
|
||||
}
|
||||
|
||||
pub fn extract_features(&self, x: &Tensor) -> Result<Tensor> {
|
||||
let batch_size = x.shape()[0];
|
||||
|
||||
// Video patch embedding
|
||||
let patches = self.patch_embed.forward(x)?;
|
||||
|
||||
// Add cls tokens for each frame
|
||||
let cls_tokens = self.cls_token.clone().expand(&[
|
||||
batch_size,
|
||||
self.config.num_frames,
|
||||
1,
|
||||
self.config.embed_dim,
|
||||
])?;
|
||||
let x = Tensor::cat(&[cls_tokens, patches], 2)?;
|
||||
|
||||
// Add positional encodings
|
||||
let spatial_pos = self.spatial_pos_embed.forward(x.shape()[2])?;
|
||||
let temporal_pos = self.temporal_pos_embed.forward(self.config.num_frames)?;
|
||||
|
||||
let spatial_pos_expanded = spatial_pos.unsqueeze(1)?.expand(&[
|
||||
batch_size,
|
||||
self.config.num_frames,
|
||||
x.shape()[2],
|
||||
self.config.embed_dim,
|
||||
])?;
|
||||
let temporal_pos_expanded = temporal_pos.unsqueeze(2)?.expand(&[
|
||||
batch_size,
|
||||
self.config.num_frames,
|
||||
x.shape()[2],
|
||||
self.config.embed_dim,
|
||||
])?;
|
||||
|
||||
let x = ((&x + &spatial_pos_expanded)? + temporal_pos_expanded)?;
|
||||
|
||||
// Pass through blocks
|
||||
let mut x = x;
|
||||
for block in &self.blocks {
|
||||
x = block.forward(&x)?;
|
||||
}
|
||||
|
||||
// Return cls token features across time
|
||||
let cls_features = x.slice(2, 0, 1)?.squeeze(Some(2))?; // [batch, frames, embed_dim]
|
||||
|
||||
Ok(cls_features)
|
||||
}
|
||||
}
|
||||
|
||||
pub struct VideoFrameSampler {
|
||||
num_frames: usize,
|
||||
uniform_sampling: bool,
|
||||
device: Device,
|
||||
}
|
||||
|
||||
impl VideoFrameSampler {
|
||||
pub fn new(num_frames: usize, uniform_sampling: bool, device: &Device) -> Result<Self> {
|
||||
Ok(Self {
|
||||
num_frames,
|
||||
uniform_sampling,
|
||||
device: device.clone(),
|
||||
})
|
||||
}
|
||||
|
||||
pub fn sample_frames(&self, video: &Tensor) -> Result<Tensor> {
|
||||
let total_frames = video.shape()[2];
|
||||
|
||||
if total_frames <= self.num_frames {
|
||||
// If video has fewer frames than requested, repeat last frame
|
||||
return Ok(video.clone());
|
||||
}
|
||||
|
||||
if self.uniform_sampling {
|
||||
// Uniform sampling
|
||||
let step = total_frames as f32 / self.num_frames as f32;
|
||||
let mut indices = Vec::new();
|
||||
for i in 0..self.num_frames {
|
||||
let idx = ((i as f32 + 0.5) * step) as usize;
|
||||
indices.push(idx.min(total_frames - 1));
|
||||
}
|
||||
|
||||
let _index_tensor = Tensor::from_vec(
|
||||
indices.into_iter().map(|i| i as f32).collect(),
|
||||
&[self.num_frames],
|
||||
&self.device,
|
||||
)?;
|
||||
|
||||
// For now, return the video as-is (would need proper frame selection)
|
||||
Ok(video.clone())
|
||||
} else {
|
||||
// Random sampling (simplified)
|
||||
let indices: Vec<f32> = (0..self.num_frames)
|
||||
.map(|_| (rand::random::<f32>() * total_frames as f32) as f32)
|
||||
.collect();
|
||||
|
||||
let _index_tensor = Tensor::from_vec(indices, &[self.num_frames], &self.device)?;
|
||||
|
||||
// For now, return the video as-is (would need proper frame selection)
|
||||
Ok(video.clone())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub struct VideoAugmentor {
|
||||
temporal_crop_ratio: f32,
|
||||
spatial_crop_ratio: f32,
|
||||
random_horizontal_flip: bool,
|
||||
device: Device,
|
||||
}
|
||||
|
||||
impl VideoAugmentor {
|
||||
pub fn new(
|
||||
temporal_crop_ratio: f32,
|
||||
spatial_crop_ratio: f32,
|
||||
random_horizontal_flip: bool,
|
||||
device: &Device,
|
||||
) -> Result<Self> {
|
||||
Ok(Self {
|
||||
temporal_crop_ratio,
|
||||
spatial_crop_ratio,
|
||||
random_horizontal_flip,
|
||||
device: device.clone(),
|
||||
})
|
||||
}
|
||||
|
||||
pub fn forward(&self, x: &Tensor) -> Result<Tensor> {
|
||||
// Simplified augmentation - in practice you'd implement proper video augmentation
|
||||
// This is just a placeholder that returns the input
|
||||
Ok(x.clone())
|
||||
}
|
||||
}
|
||||
|
||||
pub struct VideoPreprocessor {
|
||||
target_size: usize,
|
||||
num_frames: usize,
|
||||
fps: f32,
|
||||
device: Device,
|
||||
}
|
||||
|
||||
impl VideoPreprocessor {
|
||||
pub fn new(target_size: usize, num_frames: usize, fps: f32, device: &Device) -> Result<Self> {
|
||||
Ok(Self {
|
||||
target_size,
|
||||
num_frames,
|
||||
fps,
|
||||
device: device.clone(),
|
||||
})
|
||||
}
|
||||
|
||||
pub fn preprocess(&self, raw_frames: &[Vec<Vec<Vec<u8>>>]) -> Result<Tensor> {
|
||||
// Convert raw frames to tensor format
|
||||
let total_frames = raw_frames.len();
|
||||
let height = raw_frames[0].len();
|
||||
let width = raw_frames[0][0].len();
|
||||
let channels = 3;
|
||||
|
||||
// Sample frames if needed
|
||||
let frame_indices = if total_frames > self.num_frames {
|
||||
let step = total_frames as f32 / self.num_frames as f32;
|
||||
(0..self.num_frames)
|
||||
.map(|i| ((i as f32 + 0.5) * step) as usize)
|
||||
.collect::<Vec<_>>()
|
||||
} else {
|
||||
(0..total_frames).collect()
|
||||
};
|
||||
|
||||
// Convert to f32 and normalize
|
||||
let mut video_data = Vec::new();
|
||||
for &frame_idx in &frame_indices {
|
||||
let frame_idx = frame_idx.min(total_frames - 1);
|
||||
let frame = &raw_frames[frame_idx];
|
||||
|
||||
for c in 0..channels {
|
||||
for h in 0..height {
|
||||
for w in 0..width {
|
||||
let pixel_value =
|
||||
if h < frame.len() && w < frame[h].len() && c < frame[h][w].len() {
|
||||
frame[h][w][c] as f32 / 255.0
|
||||
} else {
|
||||
0.0
|
||||
};
|
||||
video_data.push(pixel_value);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let actual_frames = frame_indices.len();
|
||||
let tensor = Tensor::from_vec(
|
||||
video_data,
|
||||
&[1, channels, actual_frames, height, width],
|
||||
&self.device,
|
||||
)?;
|
||||
|
||||
Ok(tensor)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user