From 712f913478952c22dce9e10e3f84fb3fafa6d39e Mon Sep 17 00:00:00 2001
From: ninjacat272 <84151999+ninjacat272@users.noreply.github.com>
Date: Wed, 9 Sep 2026 11:11:35 -0700
Subject: [PATCH] Fix HuggingFaceDataLoader mask casting for point cloud
 datasets

_cast_single previously cast any "mask" column to Image() unconditionally,
assuming masks are always pixel maps. Point cloud datasets store mask as
Array2D (per-point labels) instead, since there's no image to draw a mask
onto. Only cast to Image() when the column isn't already Image or Array2D.

Tested against Project-AgML/pheno4d_point_cloud_segmentation: mask now
loads correctly as a list instead of being miscast into an image.
---
 agml/data/hf_loader.py | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/agml/data/hf_loader.py b/agml/data/hf_loader.py
index 77fdcdd8..1116e9f7 100644
--- a/agml/data/hf_loader.py
+++ b/agml/data/hf_loader.py
@@ -3,7 +3,7 @@
 from typing import List, Union
 
 try:
-    from datasets import load_dataset, DatasetDict, Image, Sequence, ClassLabel
+    from datasets import load_dataset, DatasetDict, Image, Sequence, ClassLabel, Array2D
 except ImportError:
     raise ImportError(
         "The `datasets` library is required to use the HuggingFaceDataLoader. "
@@ -80,8 +80,10 @@ def _cast_single(self, ds):
         if "image" in features and not isinstance(features["image"], Image):
             ds = ds.cast_column("image", Image())
 
-        # A "mask" column is always a pixel map — cast unconditionally.
-        if "mask" in features and not isinstance(features["mask"], Image):
+        # A "mask" column is a pixel map for image datasets, but a per-point
+        # label array (Array2D) for point cloud datasets. Only cast to
+        # Image() when it isn't already a numeric array type.
+        if "mask" in features and not isinstance(features["mask"], (Image, Array2D)):
             ds = ds.cast_column("mask", Image())
 
         return ds