From ceffc9d60c35aa0b6193f75ed7800fe5389fa9bc Mon Sep 17 00:00:00 2001
From: Benjamin Lefaudeux <ben@photoroom.com>
Date: Tue, 17 Dec 2024 13:03:27 +0000
Subject: [PATCH] supporting 16bits images

---
 pkg/core.go        |  1 +
 pkg/serdes.go      | 15 ++++++++++++++-
 python/go_types.py |  4 +++-
 3 files changed, 18 insertions(+), 2 deletions(-)

diff --git a/pkg/core.go b/pkg/core.go
index 1296b82..a0571db 100644
--- a/pkg/core.go
+++ b/pkg/core.go
@@ -15,6 +15,7 @@ type ImagePayload struct {
 	Height         int // Useful to decode the current payload
 	Width          int
 	Channels       int
+	BitDepth       int
 }
 
 type Sample struct {
diff --git a/pkg/serdes.go b/pkg/serdes.go
index d7e0f60..73d39f7 100644
--- a/pkg/serdes.go
+++ b/pkg/serdes.go
@@ -82,6 +82,8 @@ func imageFromBuffer(buffer []byte, transform *ARAwareTransform, aspect_ratio fl
 	// If requested, re-encode the image to a jpg or png
 	var img_bytes []byte
 	var channels int
+	var bit_depth int
+
 	if pre_encode_image {
 		if err != nil {
 			return nil, -1., err
@@ -93,12 +95,15 @@ func imageFromBuffer(buffer []byte, transform *ARAwareTransform, aspect_ratio fl
 			if err != nil {
 				return nil, -1., err
 			}
+			bit_depth = 8
 		} else {
 			// Re-encode the image to a png
-			img_bytes, _, err = img.ExportPng(vips.NewPngExportParams())
+			png_export_params := vips.NewPngExportParams()
+			img_bytes, _, err = img.ExportPng(png_export_params)
 			if err != nil {
 				return nil, -1., err
 			}
+			bit_depth = png_export_params.Bitdepth
 		}
 		channels = -1 // Signal that we have encoded the image
 	} else {
@@ -107,6 +112,13 @@ func imageFromBuffer(buffer []byte, transform *ARAwareTransform, aspect_ratio fl
 			return nil, -1., err
 		}
 		channels = img.Bands()
+
+		// Define bit depth de facto, not exposed in the vips interface
+		bit_depth = len(img_bytes) / (width * height * channels)
+	}
+
+	if bit_depth == 0 && !pre_encode_image {
+		panic("Bit depth not set")
 	}
 
 	img_payload := ImagePayload{
@@ -116,6 +128,7 @@ func imageFromBuffer(buffer []byte, transform *ARAwareTransform, aspect_ratio fl
 		Height:         height,
 		Width:          width,
 		Channels:       channels,
+		BitDepth:       bit_depth,
 	}
 
 	return &img_payload, aspect_ratio, nil
diff --git a/python/go_types.py b/python/go_types.py
index 6ed28c2..41a2d4b 100644
--- a/python/go_types.py
+++ b/python/go_types.py
@@ -12,8 +12,10 @@ def uint8_array_to_numpy(go_array):
     # We export them from Go with a Channels dimension of -1 to mark them as dimensionless.
     # Anything else is a valid number of channels and will thus lead to a reshape
     num_final_channels = max(go_array.Channels, 1)
+    bit_depth = getattr(go_array, "BitDepth", 8)
+
     length = (
-        go_array.Width * go_array.Height * num_final_channels
+        go_array.Width * go_array.Height * num_final_channels * bit_depth
         if go_array.Channels > 0
         else len(go_array.Data)
     )