datago 2025.3.10__tar.gz → 2025.3.11__tar.gz

{datago-2025.3.10 → datago-2025.3.11}/Cargo.lock

@@ -387,7 +387,7 @@ dependencies = [
 
 [[package]]
 name = "datago"
-version = "2025.3.10"
+version = "2025.3.11"
 dependencies = [
  "clap",
  "image",

{datago-2025.3.10 → datago-2025.3.11}/Cargo.toml

@@ -1,7 +1,7 @@
 [package]
 name = "datago"
 edition = "2021"
-version = "2025.3.10"
+version = "2025.3.11"
 
 [lib]
 # exposed by pyo3

{datago-2025.3.10 → datago-2025.3.11}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: datago
-Version: 2025.3.10
+Version: 2025.3.11
 Classifier: Programming Language :: Rust
 Classifier: Programming Language :: Python :: Implementation :: CPython
 Classifier: Programming Language :: Python :: Implementation :: PyPy

{datago-2025.3.10 → datago-2025.3.11}/python/test_datago_filesystem.py

@@ -3,28 +3,31 @@ from datago import DatagoClient
 import json
 import tempfile
 import pytest
-import random
+from io import BytesIO
 
 
-def generate_tmp_files(dir, limit):
+def generate_tmp_files(dir: str, limit: int, rgb16: bool = False, rgba: bool = False):
     for i in range(limit):
         # Prepare an ephemeral test set
-        mode = "RGB" if random.random() > 0.5 else "RGBA"
+        mode = "RGBA" if rgba else "RGB"
         img = Image.new(mode, (100, 100))
 
         # Randomly make the image 16 bits
-        if random.random() > 0.5:
+        if rgb16:
             img = img.convert("I;16")
 
         img.save(dir + f"/test_{i}.png")
 
 
-@pytest.mark.parametrize("pre_encode_images", [False, True])
-def test_get_sample_filesystem(pre_encode_images: bool):
+@pytest.mark.parametrize(
+    ["pre_encode_images", "rgb16", "rgba"],
+    [(a, b, c) for a in [True, False] for b in [True, False] for c in [True, False]],
+)
+def test_get_sample_filesystem(pre_encode_images: bool, rgb16: bool, rgba: bool):
     limit = 10
 
     with tempfile.TemporaryDirectory() as tmpdirname:
-        generate_tmp_files(tmpdirname, limit)
+        generate_tmp_files(tmpdirname, limit, rgb16, rgba)
 
         # Check that we can instantiate a client and get a sample, nothing more
         client_config = {
@@ -39,6 +42,7 @@ def test_get_sample_filesystem(pre_encode_images: bool):
                 "min_aspect_ratio": 0.5,
                 "max_aspect_ratio": 2.0,
                 "pre_encode_images": pre_encode_images,
+                "image_to_rgb8": rgb16 or rgba,
             },
             "limit": limit,
             "prefetch_buffer_size": 64,
@@ -58,9 +62,16 @@ def test_get_sample_filesystem(pre_encode_images: bool):
             assert data.image.width == 100
             assert data.image.height == 100
 
-            if pre_encode_images:
+            if rgb16:
                 assert data.image.bit_depth == 8
 
+            # Open the image in python scope and check properties
+            if pre_encode_images:
+                test_image = Image.open(BytesIO(data.image.data))
+                assert test_image.width == 100
+                assert test_image.height == 100
+                assert test_image.mode == "RGB"
+
         assert count == limit
 
 
@@ -112,4 +123,4 @@ def test_random_walk():
 
 
 if __name__ == "__main__":
-    test_get_sample_filesystem(True)
+    test_get_sample_filesystem(True, True, True)

{datago-2025.3.10 → datago-2025.3.11}/src/client.rs

@@ -33,6 +33,7 @@ pub struct DatagoClient {
     // Sample processing
     image_transform: Option<ARAwareTransform>,
     encode_images: bool,
+    image_to_rgb8: bool, // Convert all images to RGB 8 bits format
 
     // Threads
     pinger: Option<thread::JoinHandle<()>>,
@@ -55,12 +56,13 @@ impl DatagoClient {
 
         let mut image_transform: Option<ARAwareTransform> = None;
         let mut encode_images = false;
-
+        let mut image_to_rgb8 = false;
         if let Some(image_config) = config.image_config {
             if image_config.crop_and_resize {
                 image_transform = Some(image_config.get_ar_aware_transform());
             }
             encode_images = image_config.pre_encode_images;
+            image_to_rgb8 = image_config.image_to_rgb8;
         }
 
         DatagoClient {
@@ -79,6 +81,7 @@ impl DatagoClient {
             samples_rx,
             image_transform,
             encode_images,
+            image_to_rgb8,
             pinger: None,
             feeder: None,
             worker: None,
@@ -157,6 +160,7 @@ impl DatagoClient {
         let samples_tx_local = self.samples_tx.clone();
         let local_image_transform = self.image_transform.clone();
         let encode_images = self.encode_images;
+        let image_to_rgb8 = self.image_to_rgb8;
 
         match self.source_type {
             SourceType::Db => {
@@ -168,6 +172,7 @@ impl DatagoClient {
                         samples_tx_local,
                         local_image_transform,
                         encode_images,
+                        image_to_rgb8,
                         limit,
                     );
                 }));
@@ -179,6 +184,7 @@ impl DatagoClient {
                         samples_tx_local,
                         local_image_transform,
                         encode_images,
+                        image_to_rgb8,
                         limit,
                     );
                 }));

{datago-2025.3.10 → datago-2025.3.11}/src/image_processing.rs

@@ -1,6 +1,9 @@
-// --- Sample data structures - these will be exposed to the Python world ---------------------------------------------------------------------------------------------------------------------------------------------------------------
+use crate::structs::ImagePayload;
 use serde::Deserialize;
 use serde::Serialize;
+use std::io::Cursor;
+
+// --- Sample data structures - these will be exposed to the Python world ---------------------------------------------------------------------------------------------------------------------------------------------------------------
 
 #[derive(Debug, Serialize, Deserialize)]
 pub struct ImageTransformConfig {
@@ -9,7 +12,12 @@ pub struct ImageTransformConfig {
     pub downsampling_ratio: i32,
     pub min_aspect_ratio: f64,
     pub max_aspect_ratio: f64,
+
+    #[serde(default)]
     pub pre_encode_images: bool,
+
+    #[serde(default)]
+    pub image_to_rgb8: bool, // Convert all images to RGB 8 bits format
 }
 
 impl ImageTransformConfig {
@@ -130,6 +138,70 @@ impl ARAwareTransform {
     }
 }
 
+// ------------------------------------------------------------------------
+pub async fn image_to_payload(
+    mut image: image::DynamicImage,
+    img_tfm: &Option<ARAwareTransform>,
+    aspect_ratio: &String,
+    encode_images: bool,
+    img_to_rgb8: bool,
+) -> Result<ImagePayload, image::ImageError> {
+    let original_height = image.height() as usize;
+    let original_width = image.width() as usize;
+    let mut channels = image.color().channel_count() as i8;
+    let mut bit_depth =
+        (image.color().bits_per_pixel() / image.color().channel_count() as u16) as usize;
+
+    // Optionally transform the additional image in the same way the main image was
+    if let Some(img_tfm) = img_tfm {
+        let aspect_ratio_input = if aspect_ratio.is_empty() {
+            None
+        } else {
+            Some(aspect_ratio)
+        };
+        image = img_tfm.crop_and_resize(&image, aspect_ratio_input).await;
+    }
+
+    let height = image.height() as usize;
+    let width = image.width() as usize;
+
+    // Image to RGB8 if requested
+    if img_to_rgb8 && image.color() != image::ColorType::Rgb8 {
+        image = image::DynamicImage::ImageRgb8(image.to_rgb8());
+        bit_depth = 8;
+        channels = 3;
+        assert!((image.color().bits_per_pixel() / image.color().channel_count() as u16) == 8);
+    }
+
+    // Encode the image if needed
+    let mut image_bytes: Vec<u8> = Vec::new();
+    if encode_images {
+        if image
+            .write_to(&mut Cursor::new(&mut image_bytes), image::ImageFormat::Png)
+            .is_err()
+        {
+            return Err(image::ImageError::IoError(std::io::Error::new(
+                std::io::ErrorKind::Other,
+                "Failed to encode image",
+            )));
+        }
+
+        channels = -1; // Signal the fact that the image is encoded
+    } else {
+        image_bytes = image.into_bytes();
+    }
+
+    Ok(ImagePayload {
+        data: image_bytes,
+        original_height,
+        original_width,
+        height,
+        width,
+        channels,
+        bit_depth,
+    })
+}
+
 // ------------------------------------------------------------------------------------------------------------------------
 
 #[cfg(test)]
@@ -147,6 +219,7 @@ mod tests {
             min_aspect_ratio: 0.5,
             max_aspect_ratio: 2.0,
             pre_encode_images: false,
+            image_to_rgb8: false,
         };
 
         let transform = config.get_ar_aware_transform();

{datago-2025.3.10 → datago-2025.3.11}/src/main.rs

@@ -94,6 +94,7 @@ fn main() {
             min_aspect_ratio: 0.5,
             max_aspect_ratio: 2.0,
             pre_encode_images: true,
+            image_to_rgb8: false,
         })
     } else {
         None

{datago-2025.3.10 → datago-2025.3.11}/src/worker_files.rs

@@ -2,7 +2,6 @@ use crate::image_processing;
 use crate::structs::{ImagePayload, Sample};
 use std::cmp::min;
 use std::collections::HashMap;
-use std::io::Cursor;
 use std::sync::Arc;
 
 async fn image_from_path(path: &str) -> Result<image::DynamicImage, image::ImageError> {
@@ -19,58 +18,18 @@ async fn image_payload_from_path(
     path: &str,
     img_tfm: &Option<image_processing::ARAwareTransform>,
     encode_images: bool,
+    img_to_rgb8: bool,
 ) -> Result<ImagePayload, image::ImageError> {
     match image_from_path(path).await {
-        Ok(mut new_image) => {
-            let original_height = new_image.height() as usize;
-            let original_width = new_image.width() as usize;
-            let mut channels = new_image.color().channel_count() as i8;
-            let mut bit_depth = (new_image.color().bits_per_pixel()
-                / new_image.color().channel_count() as u16)
-                as usize;
-
-            // Optionally transform the additional image in the same way the main image was
-            if let Some(img_tfm) = img_tfm {
-                new_image = img_tfm.crop_and_resize(&new_image, None).await;
-            }
-
-            let height = new_image.height() as usize;
-            let width = new_image.width() as usize;
-
-            // Encode the image if needed
-            let mut image_bytes: Vec<u8> = Vec::new();
-            if encode_images {
-                if new_image.color() != image::ColorType::Rgb8 {
-                    new_image = image::DynamicImage::ImageRgb8(new_image.to_rgb8());
-                    bit_depth = (new_image.color().bits_per_pixel()
-                        / new_image.color().channel_count() as u16)
-                        as usize;
-                }
-
-                if new_image
-                    .write_to(&mut Cursor::new(&mut image_bytes), image::ImageFormat::Png)
-                    .is_err()
-                {
-                    return Err(image::ImageError::IoError(std::io::Error::new(
-                        std::io::ErrorKind::Other,
-                        "Failed to encode image",
-                    )));
-                }
-
-                channels = -1; // Signal the fact that the image is encoded
-            } else {
-                image_bytes = new_image.into_bytes();
-            }
-
-            Ok(ImagePayload {
-                data: image_bytes,
-                original_height,
-                original_width,
-                height,
-                width,
-                channels,
-                bit_depth,
-            })
+        Ok(new_image) => {
+            image_processing::image_to_payload(
+                new_image,
+                img_tfm,
+                &"".to_string(),
+                encode_images,
+                img_to_rgb8,
+            )
+            .await
         }
         Err(e) => Err(e),
     }
@@ -80,9 +39,17 @@ async fn pull_sample(
     sample_json: serde_json::Value,
     img_tfm: Arc<Option<image_processing::ARAwareTransform>>,
     encode_images: bool,
+    img_to_rgb8: bool,
     samples_tx: kanal::Sender<Option<Sample>>,
 ) -> Result<(), ()> {
-    match image_payload_from_path(sample_json.as_str().unwrap(), &img_tfm, encode_images).await {
+    match image_payload_from_path(
+        sample_json.as_str().unwrap(),
+        &img_tfm,
+        encode_images,
+        img_to_rgb8,
+    )
+    .await
+    {
         Ok(image) => {
             let sample = Sample {
                 id: sample_json.to_string(),
@@ -127,6 +94,7 @@ async fn async_pull_samples(
     samples_tx: kanal::Sender<Option<Sample>>,
     image_transform: Option<image_processing::ARAwareTransform>,
     encode_images: bool,
+    img_to_rgb8: bool,
     limit: usize,
 ) {
     // We use async-await here, to better use IO stalls
@@ -148,6 +116,7 @@ async fn async_pull_samples(
             received,
             shareable_img_tfm.clone(),
             encode_images,
+            img_to_rgb8,
             samples_tx.clone(),
         )));
 
@@ -177,6 +146,7 @@ pub fn pull_samples(
     samples_tx: kanal::Sender<Option<Sample>>,
     image_transform: Option<image_processing::ARAwareTransform>,
     encode_images: bool,
+    img_to_rgb8: bool,
     limit: usize,
 ) {
     tokio::runtime::Builder::new_multi_thread()
@@ -189,6 +159,7 @@ pub fn pull_samples(
                 samples_tx,
                 image_transform,
                 encode_images,
+                img_to_rgb8,
                 limit,
             )
             .await;

{datago-2025.3.10 → datago-2025.3.11}/src/worker_http.rs

@@ -5,7 +5,6 @@ use reqwest_middleware::{ClientBuilder, ClientWithMiddleware};
 use reqwest_retry::{policies::ExponentialBackoff, RetryTransientMiddleware};
 use serde::{Deserialize, Serialize};
 use std::cmp::min;
-use std::io::Cursor;
 use std::sync::Arc;
 
 // We'll share a single connection pool across all worker threads
@@ -83,72 +82,20 @@ async fn image_payload_from_url(
     img_tfm: &Option<image_processing::ARAwareTransform>,
     aspect_ratio: &String,
     encode_images: bool,
+    img_to_rgb8: bool,
 ) -> Result<ImagePayload, image::ImageError> {
     let retries = 5;
 
     match image_from_url(client, url, retries).await {
-        Ok(mut new_image) => {
-            let original_height = new_image.height() as usize;
-            let original_width = new_image.width() as usize;
-            let mut channels = new_image.color().channel_count() as i8;
-            let mut bit_depth = (new_image.color().bits_per_pixel()
-                / new_image.color().channel_count() as u16)
-                as usize;
-
-            // Optionally transform the additional image in the same way the main image was
-            if let Some(img_tfm) = img_tfm {
-                let aspect_ratio_input = if aspect_ratio.is_empty() {
-                    None
-                } else {
-                    Some(aspect_ratio)
-                };
-
-                // TODO: tokio::spawn this
-                new_image = img_tfm
-                    .crop_and_resize(&new_image, aspect_ratio_input)
-                    .await;
-            }
-
-            let height = new_image.height() as usize;
-            let width = new_image.width() as usize;
-
-            // Encode the image if needed
-            let mut image_bytes: Vec<u8> = Vec::new();
-
-            if encode_images {
-                // If the image is 16 bits or 2 channels, standardize to 8 bits, 3 channels
-                if new_image.color() != image::ColorType::Rgb8 {
-                    new_image = image::DynamicImage::ImageRgb8(new_image.to_rgb8());
-                    bit_depth = (new_image.color().bits_per_pixel()
-                        / new_image.color().channel_count() as u16)
-                        as usize;
-                }
-
-                // Pre-encode the payload as requested, we move everything to PNGs
-                if new_image
-                    .write_to(&mut Cursor::new(&mut image_bytes), image::ImageFormat::Png)
-                    .is_err()
-                {
-                    return Err(image::ImageError::IoError(std::io::Error::new(
-                        std::io::ErrorKind::Other,
-                        "Failed to encode image",
-                    )));
-                }
-
-                channels = -1; // Signal the fact that the image is encoded
-            } else {
-                image_bytes = new_image.into_bytes();
-            }
-
-            Ok(ImagePayload {
-                data: image_bytes,
-                original_height,
-                original_width,
-                height,
-                width,
-                channels,
-                bit_depth,
-            })
+        Ok(new_image) => {
+            image_processing::image_to_payload(
+                new_image,
+                img_tfm,
+                aspect_ratio,
+                encode_images,
+                img_to_rgb8,
+            )
+            .await
         }
         Err(e) => Err(e),
     }
@@ -159,6 +106,7 @@ async fn pull_sample(
     sample_json: serde_json::Value,
     img_tfm: Arc<Option<image_processing::ARAwareTransform>>,
     encode_images: bool,
+    img_to_rgb8: bool,
     samples_tx: Arc<kanal::Sender<Option<Sample>>>,
 ) -> Result<(), ()> {
     // Deserialize the sample metadata
@@ -175,6 +123,7 @@ async fn pull_sample(
             &img_tfm,
             &String::new(),
             encode_images,
+            img_to_rgb8,
         )
         .await
         {
@@ -211,6 +160,7 @@ async fn pull_sample(
                     &img_tfm,
                     &aspect_ratio,
                     encode_images,
+                    img_to_rgb8,
                 )
                 .await
                 {
@@ -235,6 +185,7 @@ async fn pull_sample(
                     &img_tfm,
                     &aspect_ratio,
                     encode_images,
+                    false, // Masks are not converted to RGB8
                 )
                 .await
                 {
@@ -306,6 +257,7 @@ async fn async_pull_samples(
     samples_tx: kanal::Sender<Option<Sample>>,
     image_transform: Option<image_processing::ARAwareTransform>,
     encode_images: bool,
+    img_to_rgb8: bool,
     limit: usize,
 ) {
     // TODO: Join with the other workers' implementation, same logic
@@ -332,6 +284,7 @@ async fn async_pull_samples(
             received,
             shareable_img_tfm.clone(),
             encode_images,
+            img_to_rgb8,
             shareable_channel_tx.clone(),
         )));
 
@@ -362,6 +315,7 @@ pub fn pull_samples(
     samples_tx: kanal::Sender<Option<Sample>>,
     image_transform: Option<image_processing::ARAwareTransform>,
     encode_images: bool,
+    img_to_rgb8: bool,
     limit: usize,
 ) {
     tokio::runtime::Builder::new_multi_thread()
@@ -376,6 +330,7 @@ pub fn pull_samples(
                 samples_tx,
                 image_transform,
                 encode_images,
+                img_to_rgb8,
                 limit,
             )
             .await;

{datago-2025.3.10 → datago-2025.3.11}/tests/client_test.rs

@@ -172,6 +172,7 @@ fn test_crop_resize() {
         min_aspect_ratio: 0.5,
         max_aspect_ratio: 2.0,
         pre_encode_images: false,
+        image_to_rgb8: false
     });
 
     let mut client = DatagoClient::new(config.to_string());
@@ -203,6 +204,7 @@ fn test_img_compression() {
         min_aspect_ratio: 0.5,
         max_aspect_ratio: 2.0,
         pre_encode_images: true, // new part being tested
+        image_to_rgb8: false
     });
 
     let mut client = DatagoClient::new(config.to_string());

{datago-2025.3.10 → datago-2025.3.11}/src/generator_http.rs

@@ -10,6 +10,7 @@ use crate::worker_http::SharedClient;
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct SourceDBConfig {
     pub sources: String,
+    pub page_size: usize,
 
     #[serde(default)]
     pub sources_ne: String,
@@ -61,7 +62,6 @@ pub struct SourceDBConfig {
 
     #[serde(default)]
     pub random_sampling: bool,
-    pub page_size: usize,
 }
 
 // TODO: Derive from the above