thumbhash.h2

thumbhash: namespace = {

PI: _ == 3.14159'26535'89793'23846L;

// down-/upscale to ~100x100, keeping the aprox. ratio
scale_to_size_nn : (src_img, src_w, src_h) -> (out_img : std::vector<u8> = (), out_w : i32, out_h : i32) = {
	out_img = std::vector<u8>();

	// resize to 100x100 (aprox, the longer size will be 100)
	img_size : const = std::max(src_w, src_h);
	out_w = std::lround(100 * src_w / img_size);
	out_h = std::lround(100 * src_h / img_size);
	//out_img = std::vector<u8>(out_w * out_h * 4, 0xff);
	out_img.resize(out_w * out_h * 4, 0xff);

	// perform nn
	for 0 ..< out_h do (dest_y) {
		for 0 ..< out_w do (dest_x) {
			sx : const = std::lround(float(dest_x) / out_w * (src_w - 1));
			sy : const = std::lround(float(dest_y) / out_h * (src_h - 1));

			src_pos : const = (sx * 4) + sy * src_w * 4;
			dst_pos : const = (dest_x * 4) + dest_y * out_w * 4;

			out_img[dst_pos + 0] = src_img[src_pos + 0];
			out_img[dst_pos + 1] = src_img[src_pos + 1];
			out_img[dst_pos + 2] = src_img[src_pos + 2];
			out_img[dst_pos + 3] = src_img[src_pos + 3];
		}
	}
}

encode_channel : (channel, w, h, nx, ny) -> (dc : float, ac : std::vector<float> = (), scale : float) = {
	dc = 0.f;
	ac.reserve(nx * ny / 2);
	scale = 0.f;

	fx := std::vector<float>(w, 0.f);

	for 0 ..< ny do (cy) {
		(copy cx := 0)
		while (cx * ny < nx * (ny - cy)) next cx++ {
			for 0 ..< w do (x) {
				// casting PI down to float here increases the similarity with the target hash
				fx[x] = std::cos(float(PI) / w * cx * (x + 0.5f));
			}

			f : float = 0.f;
			for 0 ..< h do (y) {
				fy : const = std::cos(float(PI) / h * cy * (y + 0.5f));
				for 0 ..< w do (x) {
					f += channel[x + y * w] * fx[x] * fy;
				}
			}
			f /= w * h;

			if (cx > 0 || cy > 0) {
				ac.push_back(f);
				scale = std::max(scale, std::abs(f));
			} else {
				dc = f;
			}
		}
	}

	if (scale > 0) {
		for ac do (inout e) {
			e = 0.5f + 0.5f / scale * e;
		}
	}
}


rgba_to_hash : (rgba : std::span<const u8>, w: i16, h: i16) -> std::vector<u8> = {
	assert(10 <= w <= 100);
	assert(10 <= h <= 100);

	assert(rgba.size() == w * h * 4);

	// avg color
	avg_r := 0.0f;
	avg_g := 0.0f;
	avg_b := 0.0f;
	avg_a := 0.0f; // not avg but sum

	(copy i: i32 = 0)
	while i < rgba.ssize() next i += 4 {
		pixel := rgba.subspan(i, 4);
		alpha := pixel[3] / 255.f;
		avg_r += (pixel[0] / 255.f) * alpha;
		avg_g += (pixel[1] / 255.f) * alpha;
		avg_b += (pixel[2] / 255.f) * alpha;
		avg_a += alpha;
	}

	if avg_a > 0.f {
		avg_r /= avg_a;
		avg_g /= avg_a;
		avg_b /= avg_a;
	}

	has_alpha : const bool = avg_a < w * h;
	l_limit : const i16;
	if has_alpha {
		l_limit = 5;
	} else {
		l_limit = 7;
	}

	longer_side : const i16;
	if w >= h {
		longer_side = w;
	} else {
		longer_side = h;
	}

	l_w : const = std::max<i16>(1, i16(std::lround(float(l_limit * w) / longer_side)));
	l_h : const = std::max<i16>(1, i16(std::lround(float(l_limit * h) / longer_side)));

	l := std::vector<float>(w*h); // luminance
	p := std::vector<float>(w*h); // yellow - blue
	q := std::vector<float>(w*h); // red - green
	a := std::vector<float>(w*h); // alpha

	// Convert the image from RGBA to LPQA (composite atop the average color)
	(copy i: i32 = 0)
	while i*4 < rgba.ssize() next i++ {
		pixel := rgba.subspan(i*4, 4);
		alpha := pixel[3] / 255.f;

		r := avg_r * (1.f - alpha) + pixel[0]/255.f * alpha;
		g := avg_g * (1.f - alpha) + pixel[1]/255.f * alpha;
		b := avg_b * (1.f - alpha) + pixel[2]/255.f * alpha;

		l[i] = (r + g + b) / 3.f;
		p[i] = (r + g) / 2.f - b;
		q[i] = r - g;
		a[i] = alpha;
	}

	// Encode using the DCT into DC (constant) and normalized AC (varying) terms
	dct_l := encode_channel(l, w, h, std::max<i16>(3, l_w), std::max<i16>(3, l_h));
	dct_p := encode_channel(p, w, h, 3, 3);
	dct_q := encode_channel(q, w, h, 3, 3);

	dct_a_dc := 1.f;
	dct_a_ac := std::vector<float>();
	dct_a_scale := 1.f;

	if has_alpha {
		// very meh, either declare a common return type with defaults, or beg for structured bindings
		// or use tuples?
		dct_a := encode_channel(a, w, h, 5, 5);
		dct_a_dc = dct_a.dc;
		dct_a_ac = dct_a.ac;
		dct_a_scale = dct_a.scale;
	}

	is_landscape : const = w > h;

	header24 : const u32 =
		(u32(std::lround(63.f * dct_l.dc)) << 0) |
		(u32(std::lround(31.5f + 31.5f * dct_p.dc)) << 6) |
		(u32(std::lround(31.5f + 31.5f * dct_q.dc)) << 12) |
		(u32(std::lround(31.0f * dct_l.scale)) << 18) |
		(u32(has_alpha) << 23)
	;

	landscape_l : u16;
	if is_landscape {
		landscape_l = l_h;
	} else {
		landscape_l = l_w;
	}

	header16 : const u16 =
		landscape_l |
		(u16(std::lround(63.f * dct_p.scale)) << 3) |
		(u16(std::lround(63.f * dct_q.scale)) << 9) |
		(u16(is_landscape) << 15)
	;

	hash := std::vector<u8>();
	hash.reserve(25);
	hash.push_back(u8((header24 >> 0) & 0xff));
	hash.push_back(u8((header24 >> 8) & 0xff));
	hash.push_back(u8((header24 >> 16) & 0xff));
	hash.push_back(u8((header16 >> 0) & 0xff));
	hash.push_back(u8((header16 >> 8) & 0xff));

	if has_alpha {
		hash.push_back(
			(u8(std::lround(15.f * dct_a_dc)) << 0) |
			(u8(std::lround(15.f * dct_a_scale)) << 4)
		);
	}

	is_odd := false;

	for (dct_l.ac, dct_p.ac, dct_q.ac) do (ac) {
		for ac do (f) {
			v : const = u8(std::lround(15.f * f));
			if is_odd {
				hash.back() |= v << 4;
			} else {
				hash.push_back(v);
			}
			is_odd = !is_odd;
		}
	}
	if has_alpha {
		for dct_a_ac do (f) {
			v : const = u8(std::lround(15.f * f));
			if is_odd {
				hash.back() |= v << 4;
			} else {
				hash.push_back(v);
			}
			is_odd = !is_odd;
		}
	}

	return hash;
}

// read the approx. aspect ratio from the hash
// returns 0.f on error
hash_to_aspect_ratio : (hash : std::span<const u8>) -> float = {
	if hash.ssize() < 5 {
		return 0.f;
	}

	has_alpha : const = (hash[2] & 0x80) != 0;
	l_max : const i32;
	if has_alpha {
		l_max = 5;
	} else {
		l_max = 7;
	}
	l_min : const i32 = hash[3] & 7;

	is_landscape : const = (hash[4] & 0x80) != 0;

	lx : const i32;
	if is_landscape {
		lx = l_max;
	} else {
		lx = l_min;
	}
	ly : const i32;
	if is_landscape {
		ly = l_min;
	} else {
		ly = l_max;
	}

	return float(lx) / float(ly);
}

} // thumbhash