Run Datacomp at scale

components/filter_image_resolution/Dockerfile

@@ -3,15 +3,16 @@ FROM --platform=linux/amd64 python:3.8-slim
 # System dependencies
 RUN apt-get update && \
     apt-get upgrade -y && \
-    apt-get install git -y
+    apt-get install git -y && \
+    apt-get install graphviz -y
 
 # Install requirements
 COPY requirements.txt /
 RUN pip3 install --no-cache-dir -r requirements.txt
 
 # Install Fondant
 # This is split from other requirements to leverage caching
-ARG FONDANT_VERSION=main
+ARG FONDANT_VERSION=f3f3925b8e8f634e2978e5c7fcefa72c53baba7c
 RUN pip3 install fondant[aws,azure,gcp]@git+https://github.com/ml6team/fondant@${FONDANT_VERSION}
 
 # Set the working directory to the component folder

components/filter_image_resolution/fondant_component.yaml

@@ -1,14 +1,14 @@
 name: Filter image resolution
 description: Component that filters images based on minimum size and max aspect ratio
-image: ghcr.io/ml6team/filter_image_resolution:latest
+image: ghcr.io/ml6team/filter_image_resolution:f3f3925b8e8f634e2978e5c7fcefa72c53baba7c
 
 consumes:
   image:
     fields:
       width:
-        type: int16
+        type: int64
       height:
-        type: int16
+        type: int64
 
 args:
   min_image_dim:

components/load_from_hf_hub/Dockerfile

@@ -3,15 +3,16 @@ FROM --platform=linux/amd64 python:3.8-slim
 # System dependencies
 RUN apt-get update && \
     apt-get upgrade -y && \
-    apt-get install git -y
+    apt-get install git -y && \
+    apt-get install graphviz -y
 
 # Install requirements
 COPY requirements.txt /
 RUN pip3 install --no-cache-dir -r requirements.txt
 
 # Install Fondant
 # This is split from other requirements to leverage caching
-ARG FONDANT_VERSION=main
+ARG FONDANT_VERSION=f3f3925b8e8f634e2978e5c7fcefa72c53baba7c
 RUN pip3 install fondant[aws,azure,gcp]@git+https://github.com/ml6team/fondant@${FONDANT_VERSION}
 
 # Set the working directory to the component folder

components/load_from_hf_hub/fondant_component.yaml

@@ -1,6 +1,6 @@
 name: Load from hub
 description: Component that loads a dataset from the hub
-image: ghcr.io/ml6team/load_from_hf_hub:dev
+image: ghcr.io/ml6team/load_from_hf_hub:f3f3925b8e8f634e2978e5c7fcefa72c53baba7c
 
 produces:
   dummy_variable:  #TODO: fill in here
@@ -23,4 +23,9 @@ args:
   n_rows_to_load:
     description: Optional argument that defines the number of rows to load. Useful for testing pipeline runs on a small scale
     type: int
-    default: None
+    default: None
+  dataset_length:
+    description: Optional argument that defines the length of the dataset. Required in case `n_rows_to_load` is specified.
+    type: int
+    default: None
+

components/load_from_hf_hub/src/main.py

@@ -16,6 +16,7 @@ def __init__(self, *_,
              column_name_mapping: dict,
              image_column_names: t.Optional[list],
              n_rows_to_load: t.Optional[int],
+             dataset_length: int,
     ) -> None:
         """
         Args:
@@ -25,11 +26,14 @@ def __init__(self, *_,
                 format the image from HF hub format to a byte string
             n_rows_to_load: optional argument that defines the number of rows to load. Useful for
               testing pipeline runs on a small scale.
+            dataset_length: optional argument that specifies the length of the entire dataset. Only
+            required in case n_rows_to_load is specified.
         """
         self.dataset_name = dataset_name
         self.column_name_mapping = column_name_mapping
         self.image_column_names = image_column_names
         self.n_rows_to_load = n_rows_to_load
+        self.dataset_length = dataset_length
 
     def load(self) -> dd.DataFrame:
         # 1) Load data, read as Dask dataframe
@@ -44,12 +48,27 @@ def load(self) -> dd.DataFrame:
                 )
 
         # 3) Rename columns
+        logger.info("Renaming columns...")
         dask_df = dask_df.rename(columns=self.column_name_mapping)
 
         # 4) Optional: only return specific amount of rows
-        if self.n_rows_to_load:
-            dask_df = dask_df.head(self.n_rows_to_load)
-            dask_df = dd.from_pandas(dask_df, npartitions=1)
+        if self.n_rows_to_load is not None:
+            if self.dataset_length is None:
+                raise ValueError("""Make sure to also specify the length of the entire
+                                 dataset. This is required as otherwise only the first
+                                 partition can be loaded""")
+            logger.info(f"Loading approximately {self.n_rows_to_load} rows...")
+            partition_length = self.dataset_length // dask_df.npartitions
+            npartitions = self.n_rows_to_load // partition_length
+            dask_df = dask_df.head(self.n_rows_to_load, npartitions=npartitions)
+            dask_df = dd.from_pandas(dask_df, npartitions=npartitions)
+            # .reset_index(drop=True) # will reset it from 0 for every partition
+
+        # Set monotonically increasing index
+        logger.info("Setting the index...")
+        dask_df["id"] = 1
+        dask_df["id"] = dask_df.id.cumsum()
+        dask_df = dask_df.set_index("id", sort=True)
 
         return dask_df
 

examples/pipelines/datacomp/components/cluster_image_embeddings/Dockerfile

@@ -11,7 +11,7 @@ RUN pip3 install --no-cache-dir -r requirements.txt
 
 # Install Fondant
 # This is split from other requirements to leverage caching
-ARG FONDANT_VERSION=main
+ARG FONDANT_VERSION=79df895e9d62d2010ccb8d40ee7e4fd4c68f117d
 RUN pip3 install fondant[aws,azure,gcp]@git+https://github.com/ml6team/fondant@${FONDANT_VERSION}
 
 # Set the working directory to the component folder

examples/pipelines/datacomp/components/cluster_image_embeddings/src/main.py

@@ -16,7 +16,7 @@
 class ClusterImageEmbeddingsComponent(DaskTransformComponent):
     """Component that clusters images based on embeddings."""
 
-    def __init__(self, sample_ratio: float, num_clusters: int) -> None:
+    def __init__(self, *_, sample_ratio: float, num_clusters: int) -> None:
         self.sample_ratio = sample_ratio
         self.num_clusters = num_clusters
 

examples/pipelines/datacomp/components/download_images/Dockerfile

@@ -0,0 +1,23 @@
+FROM --platform=linux/amd64 python:3.8-slim
+
+# System dependencies
+RUN apt-get update && \
+    apt-get upgrade -y && \
+    apt-get install git -y
+
+# Install requirements
+COPY requirements.txt /
+RUN pip3 install --no-cache-dir -r requirements.txt
+
+# Install Fondant
+# This is split from other requirements to leverage caching
+ARG FONDANT_VERSION=main
+RUN pip3 install fondant[aws,azure,gcp]@git+https://github.com/ml6team/fondant@${FONDANT_VERSION}
+
+# Set the working directory to the component folder
+WORKDIR /component/src
+
+# Copy over src-files
+COPY src/ .
+
+ENTRYPOINT ["python", "main.py"]

examples/pipelines/datacomp/components/download_images/README.md

@@ -0,0 +1,12 @@
+# download_images
+
+### Description
+This component takes in image URLs as input and downloads the images, along with some metadata (like their height and width).
+The images are stored in a new colum as bytes objects. This component also resizes the images using the [resizer](https://github.com/rom1504/img2dataset/blob/main/img2dataset/resizer.py) function from the img2dataset library.
+
+If the component is unable to retrieve the image at a URL (for any reason), it will return `None` for that particular URL.
+
+### **Inputs/Outputs**
+
+See [`fondant_component.yaml`](fondant_component.yaml) for a more detailed description on all the input/output parameters. 
+

examples/pipelines/datacomp/components/download_images/fondant_component.yaml

@@ -0,0 +1,49 @@
+name: Download images
+description: Component that downloads images based on URLs
+image: ghcr.io/ml6team/download_images:dev
+
+consumes:
+  image:
+    fields:
+      url:
+        type: string
+
+produces:
+  image:
+    fields:
+      data:
+        type: binary
+      width:
+        type: int16
+      height:
+        type: int16
+
+args:
+  timeout:
+    description: Maximum time (in seconds) to wait when trying to download an image
+    type: int
+    default: 10
+  retries:
+    description: Number of times to retry downloading an image if it fails.
+    type: int
+    default: 0
+  image_size:
+    description: Size of the images after resizing.
+    type: int
+    default: 256
+  resize_mode:
+    description: Resize mode to use. One of "no", "keep_ratio", "center_crop", "border".
+    type: str
+    default: 'border'
+  resize_only_if_bigger: 
+    description: If True, resize only if image is bigger than image_size.
+    type: bool
+    default: 'False'
+  min_image_size:
+    description: Minimum size of the images.
+    type: int
+    default: 0
+  max_aspect_ratio:
+    description: Maximum aspect ratio of the images.
+    type: float
+    default: 'inf'

examples/pipelines/datacomp/components/download_images/requirements.txt

@@ -0,0 +1,2 @@
+albumentations==1.3.0
+opencv-python-headless>=4.5.5.62,<5