feat: generate thumbnails

This commit is contained in:
kieran 2025-01-25 23:22:39 +00:00
parent b45018d0de
commit 6763e53d41
No known key found for this signature in database
GPG Key ID: DE71CEB3925BE941
3 changed files with 145 additions and 39 deletions

View File

@ -89,6 +89,10 @@ async fn main() -> Result<(), Error> {
{
rocket = rocket.mount("/", routes::nip96_routes());
}
#[cfg(feature = "media-compression")]
{
rocket = rocket.mount("/", routes![routes::get_blob_thumb]);
}
if let Err(e) = rocket.launch().await {
error!("Rocker error {}", e);
Err(Error::from(e))

View File

@ -1,8 +1,7 @@
use std::path::{Path, PathBuf};
use anyhow::{bail, Error, Result};
use ffmpeg_rs_raw::ffmpeg_sys_the_third::AVPixelFormat::AV_PIX_FMT_YUV420P;
use ffmpeg_rs_raw::{Demuxer, DemuxerInfo, Encoder, StreamType, Transcoder};
use ffmpeg_rs_raw::{Demuxer, DemuxerInfo, Encoder, Muxer, StreamType, Transcoder};
use std::path::{Path, PathBuf};
use uuid::Uuid;
#[cfg(feature = "labels")]
@ -21,7 +20,7 @@ impl WebpProcessor {
Self
}
pub fn process_file(
pub fn compress(
&mut self,
input: &Path,
mime_type: &str,
@ -68,6 +67,45 @@ impl WebpProcessor {
})
}
}
pub fn thumbnail(&mut self, input: &Path, out_path: &Path) -> Result<()> {
use ffmpeg_rs_raw::ffmpeg_sys_the_third::AVCodecID::AV_CODEC_ID_WEBP;
unsafe {
let mut input = Demuxer::new(input.to_str().unwrap())?;
let probe = input.probe_input()?;
let image_stream = probe
.streams
.iter()
.find(|c| c.stream_type == StreamType::Video)
.ok_or(Error::msg("No image found, cant compress"))?;
let w = 512u16;
let scale = w as f32 / image_stream.width as f32;
let h = (image_stream.height as f32 * scale) as u16;
let enc = Encoder::new(AV_CODEC_ID_WEBP)?
.with_height(h as i32)
.with_width(w as i32)
.with_pix_fmt(AV_PIX_FMT_YUV420P)
.with_framerate(1.0)?
.open(None)?;
let mut trans = Transcoder::new_custom_io(
input,
Muxer::builder()
.with_output_path(out_path.to_str().unwrap(), Some("webp"))?
.build()?,
);
trans.transcode_stream(image_stream, enc)?;
trans.run(None)?;
Ok(())
}
}
}
pub struct NewFileProcessorResult {
@ -92,7 +130,7 @@ pub fn compress_file(
if mime_type.starts_with("image/") {
let mut proc = WebpProcessor::new();
return proc.process_file(stream, mime_type, out_dir);
return proc.compress(stream, mime_type, out_dir);
}
bail!("No media processor")
}

View File

@ -1,5 +1,7 @@
use crate::db::{Database, FileUpload};
use crate::filesystem::FileStore;
#[cfg(feature = "media-compression")]
use crate::processing::WebpProcessor;
pub use crate::routes::admin::admin_routes;
#[cfg(feature = "blossom")]
pub use crate::routes::blossom::blossom_routes;
@ -16,6 +18,7 @@ use rocket::http::{ContentType, Header, Status};
use rocket::response::Responder;
use rocket::serde::Serialize;
use rocket::{Request, Response, State};
use std::env::temp_dir;
use std::io::SeekFrom;
use std::ops::Range;
use std::pin::{pin, Pin};
@ -170,49 +173,61 @@ impl AsyncRead for RangeBody {
impl<'r> Responder<'r, 'static> for FilePayload {
fn respond_to(self, request: &'r Request<'_>) -> rocket::response::Result<'static> {
let mut response = Response::new();
response.set_header(Header::new("cache-control", "max-age=31536000, immutable"));
// handle ranges
#[cfg(feature = "ranges")]
{
response.set_header(Header::new("accept-ranges", "bytes"));
if let Some(r) = request.headers().get("range").next() {
if let Ok(ranges) = parse_range_header(r) {
if ranges.ranges.len() > 1 {
warn!("Multipart ranges are not supported, fallback to non-range request");
response.set_streamed_body(self.file);
} else {
const MAX_UNBOUNDED_RANGE: u64 = 1024 * 1024;
let single_range = ranges.ranges.first().unwrap();
let range_start = match single_range.start {
StartPosition::Index(i) => i,
StartPosition::FromLast(i) => self.info.size - i,
};
let range_end = match single_range.end {
EndPosition::Index(i) => i,
EndPosition::LastByte => {
(range_start + MAX_UNBOUNDED_RANGE).min(self.info.size)
}
};
let r_len = range_end - range_start;
let r_body = RangeBody::new(self.file, range_start..range_end);
response.set_status(Status::PartialContent);
response.set_header(Header::new("content-length", r_len.to_string()));
response.set_header(Header::new(
"content-range",
format!("bytes {}-{}/{}", range_start, range_end - 1, self.info.size),
));
response.set_streamed_body(Box::pin(r_body));
}
}
const MAX_UNBOUNDED_RANGE: u64 = 1024 * 1024;
// only use range response for files > 1MiB
if self.info.size < MAX_UNBOUNDED_RANGE {
response.set_sized_body(None, self.file);
} else {
response.set_streamed_body(self.file);
response.set_header(Header::new("accept-ranges", "bytes"));
if let Some(r) = request.headers().get("range").next() {
if let Ok(ranges) = parse_range_header(r) {
if ranges.ranges.len() > 1 {
warn!(
"Multipart ranges are not supported, fallback to non-range request"
);
response.set_streamed_body(self.file);
} else {
let single_range = ranges.ranges.first().unwrap();
let range_start = match single_range.start {
StartPosition::Index(i) => i,
StartPosition::FromLast(i) => self.info.size - i,
};
let range_end = match single_range.end {
EndPosition::Index(i) => i,
EndPosition::LastByte => {
(range_start + MAX_UNBOUNDED_RANGE).min(self.info.size)
}
};
let r_len = range_end - range_start;
let r_body = RangeBody::new(self.file, range_start..range_end);
response.set_status(Status::PartialContent);
response.set_header(Header::new("content-length", r_len.to_string()));
response.set_header(Header::new(
"content-range",
format!(
"bytes {}-{}/{}",
range_start,
range_end - 1,
self.info.size
),
));
response.set_streamed_body(Box::pin(r_body));
}
}
} else {
response.set_sized_body(None, self.file);
}
}
}
#[cfg(not(feature = "ranges"))]
{
response.set_streamed_body(self.file);
response.set_header(Header::new("content-length", self.info.size.to_string()));
response.set_sized_body(None, self.file);
}
if let Ok(ct) = ContentType::from_str(&self.info.mime_type) {
@ -352,6 +367,55 @@ pub async fn head_blob(sha256: &str, fs: &State<FileStore>) -> Status {
}
}
/// Generate thumbnail for image / video
#[cfg(feature = "media-compression")]
#[rocket::get("/thumb/<sha256>")]
pub async fn get_blob_thumb(
sha256: &str,
fs: &State<FileStore>,
db: &State<Database>,
) -> Result<FilePayload, Status> {
let sha256 = if sha256.contains(".") {
sha256.split('.').next().unwrap()
} else {
sha256
};
let id = if let Ok(i) = hex::decode(sha256) {
i
} else {
return Err(Status::NotFound);
};
if id.len() != 32 {
return Err(Status::NotFound);
}
if let Ok(Some(info)) = db.get_file(&id).await {
let file_path = fs.get(&id);
let mut thumb_file = temp_dir().join(format!("thumb_{}", sha256));
thumb_file.set_extension("webp");
if !thumb_file.exists() {
let mut p = WebpProcessor::new();
if p.thumbnail(&file_path, &thumb_file).is_err() {
return Err(Status::InternalServerError);
}
};
if let Ok(f) = File::open(&thumb_file).await {
return Ok(FilePayload {
file: f,
info: FileUpload {
size: thumb_file.metadata().unwrap().len(),
mime_type: "image/webp".to_string(),
..info
},
});
}
}
Err(Status::NotFound)
}
/// Legacy URL redirect for void.cat uploads
#[rocket::get("/d/<id>")]
pub async fn void_cat_redirect(id: &str, settings: &State<Settings>) -> Option<NamedFile> {