FrenBot/src/discord/image.rs
2025-06-24 16:12:53 -06:00

365 lines
10 KiB
Rust

use crate::album_manager::{AlbumQuery, ImageQuery, ImageSort};
use crate::discord::Context;
use crate::error::Error;
use crate::models::motivation::{Motivation, MotivationConfig};
use magick_rust::{CompositeOperator, DrawingWand, FilterType, MagickWand, PixelWand};
use poise::serenity_prelude::builder::{CreateAttachment, CreateMessage};
use rand::Rng;
use rand::distr::{Distribution, StandardUniform};
use reqwest::{Client, Url};
use serde::Deserialize;
use std::borrow::Cow;
#[derive(Debug, Deserialize)]
pub enum Alignment {
Left,
Center,
Right,
}
impl Alignment {
pub fn offset(&self) -> f64 {
match self {
Alignment::Left => 0.25,
Alignment::Center => 0.5,
Alignment::Right => 0.75,
}
}
}
impl Distribution<Alignment> for StandardUniform {
fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> Alignment {
match rng.random_range(0..=2) {
0 => Alignment::Left,
1 => Alignment::Center,
_ => Alignment::Right,
}
}
}
pub async fn create_greenscreen_image(
background_image_url: Url,
foreground_image_url: Url,
alignment: Alignment,
) -> Result<Vec<u8>, Error> {
let client = Client::new();
let background_image_blob = client
.get(background_image_url)
.send()
.await?
.bytes()
.await?
.to_vec();
let foreground_image_blob = client
.get(foreground_image_url)
.send()
.await?
.bytes()
.await?
.to_vec();
let background_wand = MagickWand::new();
let foreground_wand = MagickWand::new();
background_wand.read_image_blob(background_image_blob)?;
foreground_wand.read_image_blob(foreground_image_blob)?;
let background_height = background_wand.get_image_height();
let background_width = background_wand.get_image_width();
let foreground_height = foreground_wand.get_image_height();
let foreground_width = foreground_wand.get_image_width();
let ratio = foreground_height as f64 / foreground_width as f64;
let (scale_height, scale_width) =
if foreground_height > background_height && foreground_width > background_width {
(
background_height as f64 / foreground_height as f64,
background_width as f64 / foreground_width as f64,
)
} else if foreground_height > background_height {
let scale = background_height as f64 / foreground_height as f64;
(scale, scale * ratio)
} else if foreground_width > background_width {
let scale = background_width as f64 / foreground_width as f64;
(scale * ratio, scale)
} else {
(1.0, 1.0)
};
foreground_wand.scale_image(scale_width * 0.6, scale_height * 0.6, FilterType::Lanczos2)?;
let pos_x = background_width as f64 * alignment.offset()
- foreground_wand.get_image_width() as f64 / 2.0;
let pos_y = background_height - foreground_wand.get_image_height();
background_wand.compose_images(
&foreground_wand,
CompositeOperator::Atop,
true,
pos_x as isize,
pos_y as isize,
)?;
Ok(background_wand.write_image_blob("png")?)
}
pub async fn create_overlay_image(
background_image_url: Url,
overlay_image_url: Url,
) -> Result<Vec<u8>, Error> {
let client = Client::new();
let background_image_blob = client
.get(background_image_url)
.send()
.await?
.bytes()
.await?
.to_vec();
let overlay_image_blob = client
.get(overlay_image_url)
.send()
.await?
.bytes()
.await?
.to_vec();
let background_wand = MagickWand::new();
let foreground_wand = MagickWand::new();
background_wand.read_image_blob(background_image_blob)?;
foreground_wand.read_image_blob(overlay_image_blob)?;
let overlay_height = background_wand.get_image_height();
let overlay_width = background_wand.get_image_width();
foreground_wand.resize_image(overlay_width, overlay_height, FilterType::Lanczos2)?;
background_wand.compose_images(&foreground_wand, CompositeOperator::Atop, true, 0, 0)?;
Ok(background_wand.write_image_blob("png")?)
}
pub async fn create_motivation_image(motivation: Motivation) -> Result<Vec<u8>, Error> {
let client = Client::new();
let motivation_image_blob = client
.get(motivation.image.link)
.send()
.await?
.bytes()
.await?
.to_vec();
let text = format!("{} {}", motivation.action, motivation.goal);
let mut wand = MagickWand::new();
let mut border_wand = PixelWand::new();
wand.read_image_blob(motivation_image_blob)?;
border_wand.set_color(&motivation.border_color)?;
let width = wand.get_image_width();
let border = width / 100;
wand.border_image(
&border_wand,
border,
border,
magick_rust::bindings::CompositeOperator_OverCompositeOp.into(),
)?;
border_wand.set_color("black")?;
let width = wand.get_image_width();
let border = width * 20 / 100;
wand.border_image(
&border_wand,
border,
border,
magick_rust::bindings::CompositeOperator_OverCompositeOp.into(),
)?;
let text_pos_x = wand.get_image_width() as f64 / 2.0;
let text_pos_y = wand.get_image_height() as f64 - (wand.get_image_height() as f64 * 0.05);
let mut text_wand = DrawingWand::new();
let mut text_color_wand = PixelWand::new();
text_color_wand.set_color("white")?;
text_wand.set_fill_color(&text_color_wand);
text_wand.set_font_size(0.07 * (wand.get_image_width() as f64));
text_wand.set_text_alignment(magick_rust::AlignType::Center);
wand.annotate_image(&text_wand, text_pos_x, text_pos_y, 0.0, &text)?;
Ok(wand.write_image_blob("jpg")?)
}
/// Become motivated
#[poise::command(prefix_command, category = "Image")]
pub async fn motivation(
ctx: Context<'_>,
#[description = "Album to use for the motivation"] album_name: Option<String>,
) -> Result<(), Error> {
let motivation = MotivationConfig::generate_motivation(
&ctx.data().db,
&ctx.data().picox,
"white",
album_name,
)
.await?;
let image = create_motivation_image(motivation).await?;
ctx.channel_id()
.send_message(
&ctx,
CreateMessage::new()
.content("Today's motivation")
.add_file(CreateAttachment::bytes(
Cow::from(image),
"motivate.jpg".to_string(),
)),
)
.await?;
Ok(())
}
/// Greenscreen your friends!
#[poise::command(prefix_command, category = "Image")]
pub async fn green_screen(
ctx: Context<'_>,
#[description = "Album to use as the foreground image"] foreground: Option<String>,
#[description = "Album to use for the background image"] background: Option<String>,
) -> Result<(), Error> {
let foreground = ctx
.data()
.picox
.query_image(ImageQuery {
album: foreground,
tags: vec![],
order: ImageSort::Random,
limit: 1,
})
.await?
.first()
.cloned()
.unwrap();
let background = ctx
.data()
.picox
.query_image(ImageQuery {
album: background,
tags: vec![],
order: ImageSort::Random,
limit: 1,
})
.await?
.first()
.cloned()
.unwrap();
let alignment: Alignment = rand::random();
let image =
create_greenscreen_image(background.link.clone(), foreground.link.clone(), alignment)
.await?;
ctx.channel_id()
.send_message(
&ctx,
CreateMessage::new()
.content("i am photoshop pro:")
.add_file(CreateAttachment::bytes(
Cow::from(image),
"greenscreen.jpg".to_string(),
)),
)
.await?;
Ok(())
}
/// Border
#[poise::command(prefix_command, category = "Image")]
pub async fn overlay(
ctx: Context<'_>,
#[description = "Album to use as the border image"] overlay: Option<String>,
#[description = "Album to use for the background image"] background: Option<String>,
) -> Result<(), Error> {
let foreground = ctx
.data()
.picox
.query_image(ImageQuery {
album: overlay,
tags: vec![],
order: ImageSort::Random,
limit: 1,
})
.await?
.first()
.cloned()
.unwrap();
let background = ctx
.data()
.picox
.query_image(ImageQuery {
album: background,
tags: vec![],
order: ImageSort::Random,
limit: 1,
})
.await?
.first()
.cloned()
.unwrap();
let image = create_overlay_image(background.link.clone(), foreground.link.clone()).await?;
ctx.channel_id()
.send_message(
&ctx,
CreateMessage::new()
.content("I made dis:")
.add_file(CreateAttachment::bytes(
Cow::from(image),
"overlay.jpg".to_string(),
)),
)
.await?;
Ok(())
}
/// Add an album to the pool of images to be used for motivations
#[poise::command(prefix_command, category = "Image")]
pub async fn motivation_add_album(
ctx: Context<'_>,
#[description = "Album to add to the motivation pool"] album: String,
) -> Result<(), Error> {
let albums = ctx
.data()
.picox
.query_album(AlbumQuery {
album_name: Some(album),
})
.await?;
let album = if let Some(album) = albums.first().cloned() {
album
} else {
ctx.reply("Fake album tbh, try again!").await?;
return Ok(());
};
MotivationConfig::add_album(&ctx.data().db, album.id)?;
ctx.reply("Done ;)").await?;
Ok(())
}