working?

Cargo.lock

@@ -298,16 +298,6 @@ dependencies = [
  "rayon",
 ]
 
-[[package]]
-name = "ffimage_yuv"
-version = "0.9.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "e74255e35266c3ede4909fd078b939f869cf87b2fdba8f3bf878fcd2d1a10eda"
-dependencies = [
- "ffimage",
- "num-traits",
-]
-
 [[package]]
 name = "gif"
 version = "0.11.2"
@@ -629,7 +619,6 @@ version = "0.1.0"
 dependencies = [
  "color-eyre",
  "ffimage",
- "ffimage_yuv",
  "image",
  "rsmpeg",
 ]

Cargo.toml

@@ -6,7 +6,6 @@ edition = "2018"
 [dependencies]
 image = "0.23"
 ffimage = "0.9"
-ffimage_yuv = "0.9"
 rsmpeg = "0.6"
 color-eyre = "0.5"
 

src/framestream.rs

@@ -3,8 +3,11 @@
 use color_eyre::{eyre::eyre, Report, Result};
 use rsmpeg::avcodec::{AVCodec, AVCodecContext, AVCodecID, AVCodecParametersRef, AVPacket};
 use rsmpeg::avformat::{AVFormatContextInput, AVStreamRef};
-use rsmpeg::avutil::AVFrame;
+use rsmpeg::avutil::{AVFrame, AVImage, AVPixelFormat};
+use rsmpeg::error::RsmpegError;
 use rsmpeg::ffi;
+use std::borrow::BorrowMut;
+use std::slice;
 
 pub trait FormatContextInputExt {
     fn into_packets(self) -> PacketIterator;
@@ -19,6 +22,7 @@ pub struct VideoStreamInfo {
     pub height: i32,
     pub index: i32,
     pub codec: AVCodecID,
+    pub format: AVPixelFormat,
 }
 
 impl VideoStreamInfo {
@@ -31,6 +35,7 @@ impl VideoStreamInfo {
             height,
             index: index as i32,
             codec: codec_params.codec_id,
+            format: codec_params.format,
         }
     }
 }
@@ -66,6 +71,7 @@ impl FormatContextInputExt for AVFormatContextInput {
             codec_context.open(None)?;
             (codec_context, info)
         };
+
         Ok(VideoFrames {
             packets: self.into_packets(),
             codec_context,
@@ -97,21 +103,58 @@ pub struct VideoFrames {
     pub info: VideoStreamInfo,
 }
 
+impl VideoFrames {
+    fn process_packet(&mut self) -> Option<Result<()>> {
+        let stream_index = self.info.index;
+        let packet = match self
+            .packets
+            .borrow_mut()
+            .filter(|res| match res {
+                Ok(packet) => packet.stream_index == stream_index,
+                Err(_) => true,
+            })
+            .next()
+        {
+            Some(Ok(packet)) => packet,
+            None => return None,
+            Some(Err(e)) => return Some(Err(e)),
+        };
+        Some(
+            self.codec_context
+                .send_packet(Some(&packet))
+                .map_err(Into::into),
+        )
+    }
+}
+
 impl Iterator for VideoFrames {
     type Item = Result<AVFrame>;
 
     fn next(&mut self) -> Option<Result<AVFrame>> {
-        self.packets
-            .next()
-            .filter(|res| match res {
-                Ok(packet) => packet.stream_index == self.info.index,
-                Err(_) => true,
-            })
-            .map(|res| {
-                res.and_then(|packet| {
-                    self.codec_context.send_packet(Some(&packet))?;
-                    Ok(self.codec_context.receive_frame()?)
-                })
-            })
+        loop {
+            match self.codec_context.receive_frame() {
+                Err(RsmpegError::DecoderDrainError | RsmpegError::DecoderFlushedError) => {}
+                result => return Some(result.map_err(Into::into)),
+            }
+            if let Err(e) = self.process_packet()? {
+                return Some(Err(e));
+            }
+        }
+    }
+}
+
+pub trait AVFrameExt {
+    fn data(&self) -> &[u8];
+    fn wrap(&self) -> i32;
+}
+
+impl AVFrameExt for AVFrame {
+    fn data(&self) -> &[u8] {
+        let size = AVImage::get_buffer_size(self.format, self.width, self.height, 1).unwrap();
+        unsafe { slice::from_raw_parts(self.data[0], size as usize) }
+    }
+
+    fn wrap(&self) -> i32 {
+        self.linesize[0]
     }
 }

src/main.rs

@@ -1,10 +1,13 @@
+use crate::framestream::AVFrameExt;
 use color_eyre::{eyre::eyre, Result};
 use framestream::FormatContextInputExt;
+use image::buffer::Pixels;
+use image::{ImageBuffer, Rgba};
 use rsmpeg::avformat::AVFormatContextInput;
+use rsmpeg::avutil::{AVFrameWithImage, AVImage};
+use rsmpeg::ffi::AVPixelFormat_AV_PIX_FMT_RGB32;
+use rsmpeg::swscale::SwsContext;
 use std::ffi::CString;
-use std::fs::File;
-use std::io::Write;
-use std::slice;
 
 mod framestream;
 
@@ -13,88 +16,65 @@ fn main() -> Result<()> {
     let input = AVFormatContextInput::open(&path)?;
     let frames = input.into_frames()?;
 
-    for (i, frame) in frames.take(2).enumerate() {
+    let mut encoder = SwsContext::get_context(
+        frames.info.width,
+        frames.info.height,
+        frames.info.format,
+        frames.info.width,
+        frames.info.height,
+        AVPixelFormat_AV_PIX_FMT_RGB32,
+        0,
+    )
+    .ok_or_else(|| eyre!("Failed to create encoder"))?;
+
+    let image_buffer = AVImage::new(
+        AVPixelFormat_AV_PIX_FMT_RGB32,
+        frames.info.width,
+        frames.info.height,
+        1,
+    )
+    .ok_or_else(|| eyre!("Failed to allocate image buffer"))?;
+    let mut target_frame = AVFrameWithImage::new(image_buffer);
+    let mut last_center = None;
+
+    for (i, frame) in frames.enumerate() {
         let frame = frame?;
-        let frame_filename = format!("./data/output/frame-{}.pgm", i);
 
-        save_gray_frame(
-            unsafe { slice::from_raw_parts(frame.data[0], (frame.width * frame.height) as usize) },
-            frame.linesize[0] as usize,
-            frame.width as usize,
-            frame.height as usize,
-            frame_filename,
-        )?;
+        encoder.scale_frame(&frame, 0, frame.height, &mut target_frame)?;
+
+        let image = ImageBuffer::<Rgba<u8>, _>::from_raw(
+            frame.width as u32,
+            frame.height as u32,
+            target_frame.data(),
+        )
+        .ok_or_else(|| eyre!("Failed to get image buffer"))?;
+
+        last_center = find_purple_dot(image.pixels(), frame.width as usize).or(last_center);
+        let center = last_center.ok_or_else(|| eyre!("No purple dot found"))?;
+        println!("{}, {}, {}", i, center.0, center.1);
     }
     Ok(())
 }
 
-// fn decode_packet(
-//     packet: &ffi::AVPacket,
-//     codec_context: &mut ffi::AVCodecContext,
-//     frame: &mut ffi::AVFrame,
-// ) -> Result<(), String> {
-//     let mut response = unsafe { ffi::avcodec_send_packet(codec_context, packet) };
-//
-//     if response < 0 {
-//         return Err(String::from("Error while sending a packet to the decoder."));
-//     }
-//
-//     while response >= 0 {
-//         response = unsafe { ffi::avcodec_receive_frame(codec_context, frame) };
-//         if response == ffi::AVERROR(ffi::EAGAIN) || response == ffi::AVERROR_EOF {
-//             break;
-//         } else if response < 0 {
-//             return Err(String::from(
-//                 "Error while receiving a frame from the decoder.",
-//             ));
-//         } else {
-//             println!(
-//                 "Frame {} (type={}, size={} bytes) pts {} key_frame {} [DTS {}]",
-//                 codec_context.frame_number,
-//                 unsafe { ffi::av_get_picture_type_char(frame.pict_type) },
-//                 frame.pkt_size,
-//                 frame.pts,
-//                 frame.key_frame,
-//                 frame.coded_picture_number
-//             );
-//
-//             let frame_filename = format!("./data/output/frame-{}.pgm", codec_context.frame_number);
-//             let width = frame.width as usize;
-//             let height = frame.height as usize;
-//             let wrap = frame.linesize[0] as usize;
-//             let data = unsafe { slice::from_raw_parts(frame.data[0], wrap * height) };
-//
-//             if frame.format != AVPixelFormat_AV_PIX_FMT_YUV420P {
-//                 panic!("Input has to be yuv420p, got :{}", frame.format);
-//             }
-//
-//             unsafe {
-//                 // ffi::sws_scale(codec_context, frame.data,
-//                 //                wrap, 0, height, pFrameRGB->data,
-//                 //                pFrameRGB->linesize);
-//             }
-//
-//             dbg!(wrap, width, height);
-//
-//             save_gray_frame(data, wrap, width, height, frame_filename).unwrap();
-//         }
-//     }
-//     Ok(())
-// }
+fn find_purple_dot(pixel: Pixels<Rgba<u8>>, width: usize) -> Option<(usize, usize)> {
+    let mut center_x = 0;
+    let mut center_y = 0;
+    let mut count = 0;
 
-fn save_gray_frame(
-    buf: &[u8],
-    wrap: usize,
-    xsize: usize,
-    ysize: usize,
-    filename: String,
-) -> Result<()> {
-    let mut file = File::create(filename)?;
-    let data = format!("P5\n{} {}\n{}\n", xsize, ysize, 255);
-    file.write_all(data.as_bytes())?;
+    for (i, pixel) in pixel.enumerate() {
+        let y = i / width;
+        let x = i % width;
 
-    for i in 0..ysize {
-        file.write_all(&buf[i * wrap..(i * wrap + xsize)])?;
+        if pixel[0] > 215 && pixel[1] < 10 && pixel[2] > 215 {
+            center_x += x;
+            center_y += y;
+            count += 1;
+        }
+    }
+
+    if count > 0 {
+        Some((center_x / count, center_y / count))
+    } else {
+        None
     }
-    Ok(())
 }

src/transcode.rs

@@ -0,0 +1,26 @@
+use color_eyre::{eyre::eyre, Result};
+use rsmpeg::avcodec::{AVCodec, AVCodecContext, AVCodecID};
+use rsmpeg::avutil::AVFrame;
+
+pub struct Encoder {
+    context: AVCodecContext,
+}
+
+impl Encoder {
+    pub fn new(codec_id: AVCodecID, width: i32, height: i32) -> Result<Encoder> {
+        let decoder = AVCodec::find_decoder(codec_id)
+            .ok_or_else(|| eyre!("No decoder found for codec {}", codec_id))?;
+
+        let mut context = AVCodecContext::new(&decoder);
+        context.set_width(width);
+        context.set_height(height);
+        context.open(None)?;
+
+        Ok(Encoder { context })
+    }
+
+    pub fn encode_frame(&mut self, frame: &AVFrame) -> Result<AVFrame> {
+        self.context.send_frame(Some(frame))?;
+        Ok(self.context.receive_frame()?)
+    }
+}