Обработка video audio на Rust

t.me/oneRustnoqRust

Для чего нужна данная статья? :

Написать собственный анализатор кадров с использованием модели машинного обучения для распознавания объектов.

Для чего Вам это уметь? :

Отказаться от конвертации и использовать ML - tch-rs для интеграции PyTorch, без сохранения в промежуточные форматы.

Архитектура проекта обработки видео и аудио:

1. Обработка видео: Извлечение кадров из видеофайла.
   Анализ кадров с использованием модели машинного обучения (например, распознавание объектов).
   
2. Обработка аудио: Извлечение аудиопотока из видео.
   Распознавание речи с помощью ML-модели.
   
3. Машинное обучение: Интеграция предобученных моделей для анализа видео (например, YOLO) и аудио (например, Whisper).
   Обработка результатов и их синхронизация.
   
4. Синхронизация и вывод: Сопоставление временных меток кадров и аудио для создания единого результата.
   

Для реализации используются следующие библиотеки Rust:

• ffmpeg-next — для работы с видео и аудио.
• tch-rs — для интеграции моделей PyTorch (например, для обработки видео).
• image — для работы с изображениями (кадрами).
• Псевдокод для аудио ML, так как полноценная поддержка моделей вроде Whisper в Rust ограничена и может потребовать внешних интеграций.

Зависимости

Добавьте следующие зависимости в Cargo.toml:

[dependencies]

ffmpeg-next = "6.0"

tch = "0.13"

image = "0.24"

anyhow = "1.0"

Код проекта

1. Извлечение кадров из видео

Этот модуль извлекает кадры из видеофайла и сохраняет их для дальнейшей обработки.

use ffmpeg_next::{

format::{self, Pixel},

frame, media,

software::scaling::{self, Flags},

Error,

};

use std::path::Path;

fn extract_frames(video_path: &Path, output_dir: &Path) -> Result<Vec<f64>, Error> {

ffmpeg_next::init()?;

let mut ictx = format::input(&video_path)?;

let video_stream = ictx

.streams()

.best(media::Type::Video)

.ok_or(Error::StreamNotFound)?;

let video_stream_index = video_stream.index();

let context_decoder = ffmpeg_next::codec::context::Context::from_parameters(video_stream.parameters())?;

let mut decoder = context_decoder.decoder().video()?;

let mut scaler = scaling::Context::get(

decoder.format(),

decoder.width(),

decoder.height(),

Pixel::RGB24,

decoder.width(),

decoder.height(),

Flags::BILINEAR,

)?;

let mut timestamps = Vec::new();

let mut frame_index = 0;

for (stream, packet) in ictx.packets() {

if stream.index() == video_stream_index {

decoder.send_packet(&packet)?;

let mut decoded = frame::Video::empty();

while decoder.receive_frame(&mut decoded).is_ok() {

let timestamp = packet.pts().unwrap_or(0) as f64 * stream.time_base().numerator() as f64

/ stream.time_base().denominator() as f64;

timestamps.push(timestamp);

let mut rgb_frame = frame::Video::empty();

scaler.run(&decoded, &mut rgb_frame)?;

let frame_path = output_dir.join(format!("frame_{:04}.png", frame_index));

image::save_buffer(

&frame_path,

rgb_frame.data(0),

rgb_frame.width(),

rgb_frame.height(),

image::ColorType::Rgb8,

).map_err(|e| Error::Other { error: Box::new(e) })?;

frame_index += 1;

}

}

}

decoder.send_eof()?;

Ok(timestamps)

}

2. Обработка кадров с помощью ML

Для анализа кадров используется предобученная модель (например, YOLO) через tch-rs.

use tch::{nn, Device, Tensor, vision::image, TchError};

use std::path::Path;

struct FrameResult {

timestamp: f64,

objects: Vec<(String, f32, [f32; 4])>, // (метка, вероятность, [x_min, y_min, x_max, y_max])

}

fn load_model() -> Result<nn::Sequential, TchError> {

// Загрузка предобученной модели (псевдокод, требуется путь к модели)

let mut vs = nn::VarStore::new(Device::Cuda(0));

let model = nn::seq()

.add(nn::conv2d(&vs.root(), 3, 64, 7, 2)?); // Пример слоя, замените реальной моделью

Ok(model)

}

fn process_frame(frame_path: &Path, model: &nn::Sequential, timestamp: f64) -> Result<FrameResult, TchError> {

let img = image::load(frame_path)?;

let input = img.unsqueeze(0).to_device(Device::Cuda(0));

tch::no_grad(|| {

let output = model.forward(&input);

// Псевдокод для обработки выхода YOLO

let objects = vec![("object".to_string(), 0.95, [100.0, 100.0, 200.0, 200.0])];

Ok(FrameResult { timestamp, objects })

})

}

3. Извлечение аудио из видео

Этот модуль извлекает аудиопоток и сохраняет его в файл.

use ffmpeg_next::{format, media, Error};

use std::path::Path;

fn extract_audio(video_path: &Path, audio_path: &Path) -> Result<(), Error> {

ffmpeg_next::init()?;

let mut ictx = format::input(&video_path)?;

let audio_stream = ictx

.streams()

.best(media::Type::Audio)

.ok_or(Error::StreamNotFound)?;

let mut octx = format::output(&audio_path)?;

let mut ost = octx.add_stream(ffmpeg_next::encoder::find_by_name("aac"))?;

ost.set_parameters(audio_stream.parameters());

let mut encoder = ost.codec().encoder().audio()?;

encoder.set_bit_rate(audio_stream.bit_rate());

encoder.set_sample_rate(audio_stream.sample_rate());

encoder.set_channels(audio_stream.channels());

encoder.open_as(ffmpeg_next::encoder::find_by_name("aac"))?;

for (stream, packet) in ictx.packets() {

if stream.index() == audio_stream.index() {

packet.write_interleaved(&mut octx)?;

}

}

octx.write_trailer()?;

Ok(())

}

4. Обработка аудио с помощью ML

Для распознавания речи используется псевдокод, так как интеграция моделей вроде Whisper в Rust сложна и может потребовать внешних вызовов.

use std::path::Path;

struct AudioResult {

start_time: f64,

end_time: f64,

text: String,

}

fn transcribe_audio(audio_path: &Path) -> Result<Vec<AudioResult>, Box<dyn std::error::Error>> {

// Псевдокод для вызова модели распознавания речи (например, Whisper)

let transcription = vec![AudioResult {

start_time: 0.0,

end_time: 5.0,

text: "Пример текста".to_string(),

}];

Ok(transcription)

}

5. Основная программа и синхронизация

Собираем все компоненты вместе:

use anyhow::Result;

use std::path::Path;

fn main() -> Result<()> {

let video_path = Path::new("input.mp4");

let output_dir = Path::new("frames");

std::fs::create_dir_all(output_dir)?;

// Извлечение кадров

let timestamps = extract_frames(video_path, output_dir)?;

let model = load_model()?;

let mut frame_results = Vec::new();

for (i, ts) in timestamps.into_iter().enumerate() {

let frame_path = output_dir.join(format!("frame_{:04}.png", i));

let result = process_frame(&frame_path, &model, ts)?;

frame_results.push(result);

}

// Извлечение и обработка аудио

let audio_path = Path::new("output.aac");

extract_audio(video_path, audio_path)?;

let audio_results = transcribe_audio(audio_path)?;

// Вывод результатов

println!("Результаты обработки видео:");

for fr in frame_results {

println!("Время: {:.2}s, Объекты: {:?}", fr.timestamp, fr.objects);

}

println!("Результаты обработки аудио:");

for ar in audio_results {

println!("Время: {:.2}s - {:.2}s, Текст: {}", ar.start_time, ar.end_time, ar.text);

}

Ok(())

}