Fix caching, decode images ahead of time (#229)

* revert `Only request the same information once`

* improve log macro formatting

* random improvements and fix cache

* switch to unbounded channel

* Revert "switch to unbounded channel"

This reverts commit b0d2592.

* switch to unbounded channel

* predecode images

* clippy

* add missed files to commit

* fix timeout bug in overtime

* fix timeout bug in game

* fix alphagen

* Fix assets to prevent alpha in color display

* Display team names in overtime/sudden death page

* Add PMK mode to alphagen

* Re-generate images with alphagen PMK mode

* Fix timeout flag animations

---------

Co-authored-by: Tristan Debrunner <tdebrunner@atlantissports.org>

alphagen/src/lib.rs

@@ -1,25 +1,70 @@
-use image::io::Reader;
-use image::GenericImageView;
-use image::ImageFormat;
-use rayon::iter::ParallelBridge;
-use rayon::iter::ParallelIterator;
-use std::fs;
-use std::io::BufWriter;
-use std::io::Cursor;
-use std::path::PathBuf;
+use image::{io::Reader, GenericImageView, ImageBuffer, ImageFormat, LumaA, Rgba};
+use rayon::iter::{ParallelBridge, ParallelIterator};
+use std::{
+    fs,
+    io::{BufWriter, Cursor},
+    path::PathBuf,
+};
 
 /// Processes all files in `paths` and writes output images to `dir_output`
 pub fn on_paths(paths: Vec<&PathBuf>, dir_output: PathBuf) {
+    let process = |in_pix: &Rgba<u8>, out_pix: &mut LumaA<u8>| {
+        out_pix.0[0] = in_pix[3];
+        out_pix.0[1] = in_pix[3];
+    };
+
+    process_images_on_paths(paths, dir_output, process);
+}
+
+/// Processes all files in `paths` and writes output images to `dir_output`.
+/// Removes alpha channel from images.
+pub fn remove_alpha_on_paths(paths: Vec<&PathBuf>, dir_output: PathBuf) {
+    let process = |in_pix: &Rgba<u8>, out_pix: &mut Rgba<u8>| {
+        out_pix.0[0] = in_pix[0];
+        out_pix.0[1] = in_pix[1];
+        out_pix.0[2] = in_pix[2];
+        if in_pix[3] == 0 {
+            out_pix.0[3] = 0;
+        } else {
+            out_pix.0[3] = 255;
+        }
+    };
+
+    process_images_on_paths(paths, dir_output, process);
+}
+
+/// Processes all files in `paths` and writes output images to `dir_output`.
+/// Pre-multiplies the colors with the alpha channel for use with the ATEM.
+pub fn pre_multiply_on_paths(paths: Vec<&PathBuf>, dir_output: PathBuf) {
+    let process = |in_pix: &Rgba<u8>, out_pix: &mut Rgba<u8>| {
+        out_pix.0[0] = ((in_pix[0] as u16 * in_pix[3] as u16) / 255) as u8;
+        out_pix.0[1] = ((in_pix[1] as u16 * in_pix[3] as u16) / 255) as u8;
+        out_pix.0[2] = ((in_pix[2] as u16 * in_pix[3] as u16) / 255) as u8;
+        out_pix.0[3] = in_pix[3];
+    };
+
+    process_images_on_paths(paths, dir_output, process);
+}
+
+fn process_images_on_paths<O>(
+    paths: Vec<&PathBuf>,
+    dir_output: PathBuf,
+    process: fn(&Rgba<u8>, &mut O),
+) where
+    O: image::Pixel + image::PixelWithColorType,
+    [O::Subpixel]: image::EncodableLayout,
+{
     paths.iter().par_bridge().for_each(|path| {
-        let file = image::open(path).unwrap();
-        let mut output_image_buff = image::GrayAlphaImage::new(file.width(), file.height());
+        let file =
+            image::open(path).unwrap_or_else(|_| panic!("Couldn't open image at {:?}", path));
+        let mut output_image_buff =
+            image::ImageBuffer::<O, Vec<O::Subpixel>>::new(file.width(), file.height());
         let mut pixs = output_image_buff.pixels_mut();
         let mut file_out = fs::File::create(dir_output.join(path.file_name().unwrap()))
             .expect("Couldn't create output file");
-        for (_, _, alpha_channel) in file.pixels() {
+        for (_, _, pixel) in file.pixels() {
             let p = pixs.next().unwrap();
-            p.0[0] = alpha_channel[3];
-            p.0[1] = alpha_channel[3];
+            process(&pixel, p);
         }
         output_image_buff
             .write_to(&mut file_out, image::ImageFormat::Png)
@@ -40,10 +85,30 @@ pub fn on_raw(input: &[u8]) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
         .decode()?;
     let mut img_out = image::GrayAlphaImage::new(img_in.width(), img_in.height());
     let mut pixs = img_out.pixels_mut();
-    for (_, _, alpha_channel) in img_in.pixels() {
+    for (_, _, pixel) in img_in.pixels() {
+        let p = pixs.next().unwrap();
+        p.0[0] = pixel[3];
+        p.0[1] = pixel[3];
+    }
+    let mut writer = BufWriter::new(Cursor::new(Vec::new()));
+    img_out.write_to(&mut writer, ImageFormat::Png)?;
+    Ok(writer.into_inner()?.into_inner())
+}
+
+/// Process raw rgba8 image data
+pub fn on_raw_rgba8(
+    width: u32,
+    height: u32,
+    buff: Vec<u8>,
+) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
+    let img_in: ImageBuffer<Rgba<u8>, Vec<u8>> =
+        ImageBuffer::from_raw(width, height, buff).unwrap();
+    let mut img_out = image::GrayAlphaImage::new(img_in.width(), img_in.height());
+    let mut pixs = img_out.pixels_mut();
+    for Rgba([_, _, _, alpha_channel]) in img_in.pixels() {
         let p = pixs.next().unwrap();
-        p.0[0] = alpha_channel[3];
-        p.0[1] = alpha_channel[3];
+        p.0[0] = *alpha_channel;
+        p.0[1] = *alpha_channel;
     }
     let mut writer = BufWriter::new(Cursor::new(Vec::new()));
     img_out.write_to(&mut writer, ImageFormat::Png)?;

alphagen/src/main.rs

@@ -1,4 +1,4 @@
-use alphagen::on_paths;
+use alphagen::{on_paths, pre_multiply_on_paths, remove_alpha_on_paths};
 use clap::Parser;
 use log::warn;
 use rayon::iter::ParallelBridge;
@@ -8,6 +8,22 @@ use std::path::PathBuf;
 #[derive(Parser, Debug)]
 #[command(author, version, about)]
 struct Args {
+    #[clap(
+        short,
+        help = "Remove alpha channel from images",
+        conflicts_with = "pre_multiply",
+        default_value_t = false
+    )]
+    remove_alpha_mode: bool,
+
+    #[clap(
+        short,
+        help = "Pre-multiply the color image",
+        conflicts_with = "remove_alpha_mode",
+        default_value_t = false
+    )]
+    pre_multiply: bool,
+
     #[clap(help = "Input files", required = true)]
     input_location: Vec<PathBuf>,
 
@@ -38,6 +54,12 @@ fn main() {
         })
         .collect::<Vec<_>>();
 
-    assert!(paths.is_empty(), "No valid input file paths!");
-    on_paths(paths, dir_output);
+    assert!(!paths.is_empty(), "No valid input file paths!");
+    if args.remove_alpha_mode {
+        remove_alpha_on_paths(paths, dir_output)
+    } else if args.pre_multiply {
+        pre_multiply_on_paths(paths, dir_output)
+    } else {
+        on_paths(paths, dir_output);
+    }
 }

overlay/src/main.rs

@@ -64,10 +64,10 @@ pub struct State {
 
 // TODO: Change this to return Result. We're not rn cause from_file_with_format
 // panics anyways if image bytes is invalid
-pub fn texture_from_bytes(bytes: Vec<u8>) -> Texture {
+pub fn texture_from_image(image: network::Image) -> Texture {
     Texture {
-        color: Texture2D::from_file_with_format(&bytes, None),
-        alpha: alphagen::on_raw(&bytes)
+        color: Texture2D::from_rgba8(image.0, image.1, &image.2),
+        alpha: alphagen::on_raw_rgba8(image.0 as u32, image.1 as u32, image.2)
             .map(|bytes| Texture2D::from_file_with_format(&bytes, None))
             .expect("Failed to decode image"),
     }
@@ -88,7 +88,7 @@ impl State {
             self.black = TeamInfo::from(black);
             self.white = TeamInfo::from(white);
             self.start_time = start_time;
-            self.sponsor_logo = sponsor_logo.map(texture_from_bytes);
+            self.sponsor_logo = sponsor_logo.map(texture_from_image);
             self.referees = referees.into_iter().map(Member::from).collect();
             self.pool = pool;
         }
@@ -115,8 +115,8 @@ impl From<network::MemberRaw> for Member {
             name: member_raw.name,
             role: member_raw.role,
             number: member_raw.number,
-            picture: member_raw.picture.map(texture_from_bytes),
-            geared_picture: member_raw.geared_picture.map(texture_from_bytes),
+            picture: member_raw.picture.map(texture_from_image),
+            geared_picture: member_raw.geared_picture.map(texture_from_image),
         }
     }
 }
@@ -146,7 +146,7 @@ impl From<TeamInfoRaw> for TeamInfo {
                 .into_iter()
                 .map(Member::from)
                 .collect(),
-            flag: team_info_raw.flag.map(texture_from_bytes),
+            flag: team_info_raw.flag.map(texture_from_image),
         }
     }
 }