mirror of
https://github.com/icewind1991/purpledot.git
synced 2026-06-03 17:44:11 +02:00
working?
This commit is contained in:
parent
be1dd254e5
commit
9724bdd2b0
5 changed files with 140 additions and 103 deletions
11
Cargo.lock
generated
11
Cargo.lock
generated
|
|
@ -298,16 +298,6 @@ dependencies = [
|
||||||
"rayon",
|
"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]]
|
[[package]]
|
||||||
name = "gif"
|
name = "gif"
|
||||||
version = "0.11.2"
|
version = "0.11.2"
|
||||||
|
|
@ -629,7 +619,6 @@ version = "0.1.0"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"color-eyre",
|
"color-eyre",
|
||||||
"ffimage",
|
"ffimage",
|
||||||
"ffimage_yuv",
|
|
||||||
"image",
|
"image",
|
||||||
"rsmpeg",
|
"rsmpeg",
|
||||||
]
|
]
|
||||||
|
|
|
||||||
|
|
@ -6,7 +6,6 @@ edition = "2018"
|
||||||
[dependencies]
|
[dependencies]
|
||||||
image = "0.23"
|
image = "0.23"
|
||||||
ffimage = "0.9"
|
ffimage = "0.9"
|
||||||
ffimage_yuv = "0.9"
|
|
||||||
rsmpeg = "0.6"
|
rsmpeg = "0.6"
|
||||||
color-eyre = "0.5"
|
color-eyre = "0.5"
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -3,8 +3,11 @@
|
||||||
use color_eyre::{eyre::eyre, Report, Result};
|
use color_eyre::{eyre::eyre, Report, Result};
|
||||||
use rsmpeg::avcodec::{AVCodec, AVCodecContext, AVCodecID, AVCodecParametersRef, AVPacket};
|
use rsmpeg::avcodec::{AVCodec, AVCodecContext, AVCodecID, AVCodecParametersRef, AVPacket};
|
||||||
use rsmpeg::avformat::{AVFormatContextInput, AVStreamRef};
|
use rsmpeg::avformat::{AVFormatContextInput, AVStreamRef};
|
||||||
use rsmpeg::avutil::AVFrame;
|
use rsmpeg::avutil::{AVFrame, AVImage, AVPixelFormat};
|
||||||
|
use rsmpeg::error::RsmpegError;
|
||||||
use rsmpeg::ffi;
|
use rsmpeg::ffi;
|
||||||
|
use std::borrow::BorrowMut;
|
||||||
|
use std::slice;
|
||||||
|
|
||||||
pub trait FormatContextInputExt {
|
pub trait FormatContextInputExt {
|
||||||
fn into_packets(self) -> PacketIterator;
|
fn into_packets(self) -> PacketIterator;
|
||||||
|
|
@ -19,6 +22,7 @@ pub struct VideoStreamInfo {
|
||||||
pub height: i32,
|
pub height: i32,
|
||||||
pub index: i32,
|
pub index: i32,
|
||||||
pub codec: AVCodecID,
|
pub codec: AVCodecID,
|
||||||
|
pub format: AVPixelFormat,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl VideoStreamInfo {
|
impl VideoStreamInfo {
|
||||||
|
|
@ -31,6 +35,7 @@ impl VideoStreamInfo {
|
||||||
height,
|
height,
|
||||||
index: index as i32,
|
index: index as i32,
|
||||||
codec: codec_params.codec_id,
|
codec: codec_params.codec_id,
|
||||||
|
format: codec_params.format,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
@ -66,6 +71,7 @@ impl FormatContextInputExt for AVFormatContextInput {
|
||||||
codec_context.open(None)?;
|
codec_context.open(None)?;
|
||||||
(codec_context, info)
|
(codec_context, info)
|
||||||
};
|
};
|
||||||
|
|
||||||
Ok(VideoFrames {
|
Ok(VideoFrames {
|
||||||
packets: self.into_packets(),
|
packets: self.into_packets(),
|
||||||
codec_context,
|
codec_context,
|
||||||
|
|
@ -97,21 +103,58 @@ pub struct VideoFrames {
|
||||||
pub info: VideoStreamInfo,
|
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 {
|
impl Iterator for VideoFrames {
|
||||||
type Item = Result<AVFrame>;
|
type Item = Result<AVFrame>;
|
||||||
|
|
||||||
fn next(&mut self) -> Option<Result<AVFrame>> {
|
fn next(&mut self) -> Option<Result<AVFrame>> {
|
||||||
self.packets
|
loop {
|
||||||
.next()
|
match self.codec_context.receive_frame() {
|
||||||
.filter(|res| match res {
|
Err(RsmpegError::DecoderDrainError | RsmpegError::DecoderFlushedError) => {}
|
||||||
Ok(packet) => packet.stream_index == self.info.index,
|
result => return Some(result.map_err(Into::into)),
|
||||||
Err(_) => true,
|
}
|
||||||
})
|
if let Err(e) = self.process_packet()? {
|
||||||
.map(|res| {
|
return Some(Err(e));
|
||||||
res.and_then(|packet| {
|
}
|
||||||
self.codec_context.send_packet(Some(&packet))?;
|
}
|
||||||
Ok(self.codec_context.receive_frame()?)
|
}
|
||||||
})
|
}
|
||||||
})
|
|
||||||
|
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]
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
|
||||||
136
src/main.rs
136
src/main.rs
|
|
@ -1,10 +1,13 @@
|
||||||
|
use crate::framestream::AVFrameExt;
|
||||||
use color_eyre::{eyre::eyre, Result};
|
use color_eyre::{eyre::eyre, Result};
|
||||||
use framestream::FormatContextInputExt;
|
use framestream::FormatContextInputExt;
|
||||||
|
use image::buffer::Pixels;
|
||||||
|
use image::{ImageBuffer, Rgba};
|
||||||
use rsmpeg::avformat::AVFormatContextInput;
|
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::ffi::CString;
|
||||||
use std::fs::File;
|
|
||||||
use std::io::Write;
|
|
||||||
use std::slice;
|
|
||||||
|
|
||||||
mod framestream;
|
mod framestream;
|
||||||
|
|
||||||
|
|
@ -13,88 +16,65 @@ fn main() -> Result<()> {
|
||||||
let input = AVFormatContextInput::open(&path)?;
|
let input = AVFormatContextInput::open(&path)?;
|
||||||
let frames = input.into_frames()?;
|
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 = frame?;
|
||||||
let frame_filename = format!("./data/output/frame-{}.pgm", i);
|
|
||||||
|
|
||||||
save_gray_frame(
|
encoder.scale_frame(&frame, 0, frame.height, &mut target_frame)?;
|
||||||
unsafe { slice::from_raw_parts(frame.data[0], (frame.width * frame.height) as usize) },
|
|
||||||
frame.linesize[0] as usize,
|
let image = ImageBuffer::<Rgba<u8>, _>::from_raw(
|
||||||
frame.width as usize,
|
frame.width as u32,
|
||||||
frame.height as usize,
|
frame.height as u32,
|
||||||
frame_filename,
|
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(())
|
Ok(())
|
||||||
}
|
}
|
||||||
|
|
||||||
// fn decode_packet(
|
fn find_purple_dot(pixel: Pixels<Rgba<u8>>, width: usize) -> Option<(usize, usize)> {
|
||||||
// packet: &ffi::AVPacket,
|
let mut center_x = 0;
|
||||||
// codec_context: &mut ffi::AVCodecContext,
|
let mut center_y = 0;
|
||||||
// frame: &mut ffi::AVFrame,
|
let mut count = 0;
|
||||||
// ) -> 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 save_gray_frame(
|
for (i, pixel) in pixel.enumerate() {
|
||||||
buf: &[u8],
|
let y = i / width;
|
||||||
wrap: usize,
|
let x = i % width;
|
||||||
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 in 0..ysize {
|
if pixel[0] > 215 && pixel[1] < 10 && pixel[2] > 215 {
|
||||||
file.write_all(&buf[i * wrap..(i * wrap + xsize)])?;
|
center_x += x;
|
||||||
|
center_y += y;
|
||||||
|
count += 1;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if count > 0 {
|
||||||
|
Some((center_x / count, center_y / count))
|
||||||
|
} else {
|
||||||
|
None
|
||||||
}
|
}
|
||||||
Ok(())
|
|
||||||
}
|
}
|
||||||
|
|
|
||||||
26
src/transcode.rs
Normal file
26
src/transcode.rs
Normal file
|
|
@ -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()?)
|
||||||
|
}
|
||||||
|
}
|
||||||
Loading…
Add table
Add a link
Reference in a new issue