From 7517b7fd242d1adbe29ae7b0e4d4c6bffa0a5950 Mon Sep 17 00:00:00 2001
From: Diego Dorado <diegodorado@gmail.com>
Date: Mon, 26 Aug 2024 14:41:40 -0300
Subject: [PATCH] encode images with alpha channel if present for preview

---
 yazi-adapter/src/image.rs  | 27 +++++++++++++++++----------
 yazi-adapter/src/iterm2.rs | 33 +++++++++++++++++++++++----------
 2 files changed, 40 insertions(+), 20 deletions(-)

diff --git a/yazi-adapter/src/image.rs b/yazi-adapter/src/image.rs
index 4535087c..d7b42575 100644
--- a/yazi-adapter/src/image.rs
+++ b/yazi-adapter/src/image.rs
@@ -2,7 +2,7 @@ use std::{fs::File, io::BufReader, path::{Path, PathBuf}};
 
 use anyhow::Result;
 use exif::{In, Tag};
-use image::{codecs::jpeg::JpegEncoder, imageops::{self, FilterType}, DynamicImage, Limits};
+use image::{codecs::{jpeg::JpegEncoder, png::PngEncoder}, imageops::{self, FilterType}, DynamicImage, ImageEncoder, Limits};
 use ratatui::layout::Rect;
 use yazi_config::{PREVIEW, TASKS};
 
@@ -30,18 +30,25 @@ impl Image {
 				img = img.resize(w, h, Self::filter());
 			}
 
+			let mut buf = Vec::new();
 			img = Self::rotate(img, orientation);
-			if !matches!(img, DynamicImage::ImageRgb8(_)) {
-				img = DynamicImage::ImageRgb8(img.into_rgb8());
+
+			match img {
+				DynamicImage::ImageRgb8(_) => {
+					JpegEncoder::new_with_quality(&mut buf, PREVIEW.image_quality).encode_image(&img)?;
+				}
+				_ => {
+					let (width, height) = (img.width(), img.height());
+					let img_rgba8 = img.into_rgba8(); // Consume `img` here
+					PngEncoder::new(&mut buf).write_image(
+						&img_rgba8,
+						width,
+						height,
+						image::ExtendedColorType::Rgba8,
+					)?;
+				}
 			}
 
-			let mut buf = Vec::new();
-			JpegEncoder::new_with_quality(&mut buf, PREVIEW.image_quality).encode(
-				img.as_bytes(),
-				img.width(),
-				img.height(),
-				img.color().into(),
-			)?;
 			Ok::<_, anyhow::Error>(buf)
 		})
 		.await??;
diff --git a/yazi-adapter/src/iterm2.rs b/yazi-adapter/src/iterm2.rs
index d71cbf93..3b7108d2 100644
--- a/yazi-adapter/src/iterm2.rs
+++ b/yazi-adapter/src/iterm2.rs
@@ -1,10 +1,11 @@
 use std::{io::Write, path::Path};
 
 use anyhow::Result;
-use base64::{engine::{general_purpose::STANDARD, Config}, Engine};
+use base64::{engine::general_purpose::STANDARD, Engine};
 use crossterm::{cursor::MoveTo, queue};
-use image::{codecs::jpeg::JpegEncoder, DynamicImage};
+use image::{codecs::{jpeg::JpegEncoder, png::PngEncoder}, DynamicImage, ImageEncoder};
 use ratatui::layout::Rect;
+use yazi_config::PREVIEW;
 
 use super::image::Image;
 use crate::{adapter::Adapter, Emulator, CLOSE, START};
@@ -38,22 +39,34 @@ impl Iterm2 {
 
 	async fn encode(img: DynamicImage) -> Result<Vec<u8>> {
 		tokio::task::spawn_blocking(move || {
-			let mut jpg = vec![];
-			JpegEncoder::new_with_quality(&mut jpg, 75).encode_image(&img)?;
+			let width = img.width();
+			let height = img.height();
 
-			let len = base64::encoded_len(jpg.len(), STANDARD.config().encode_padding());
-			let mut buf = Vec::with_capacity(200 + len.unwrap_or(1 << 16));
+			let mut img_buf = vec![];
 
+			if img.color().has_alpha() {
+				PngEncoder::new(&mut img_buf).write_image(
+					&img.into_rgba8(),
+					width,
+					height,
+					image::ExtendedColorType::Rgba8,
+				)?;
+			} else {
+				JpegEncoder::new_with_quality(&mut img_buf, PREVIEW.image_quality).encode_image(&img)?;
+			}
+
+			let mut buf = vec![];
 			write!(
 				buf,
 				"{}]1337;File=inline=1;size={};width={}px;height={}px;doNotMoveCursor=1:{}\x07{}",
 				START,
-				jpg.len(),
-				img.width(),
-				img.height(),
-				STANDARD.encode(&jpg),
+				img_buf.len(),
+				width,
+				height,
+				STANDARD.encode(&img_buf),
 				CLOSE
 			)?;
+
 			Ok(buf)
 		})
 		.await?