feat(dataset push): object_detection + keypoint_detection

internal/cli/dataset.go

@@ -72,14 +72,15 @@ func newDatasetPushCmd() *cobra.Command {
  // Ingest-spec flags. image_classification + the tabular /
  // time-series family are supported today; text + detection +
  // segmentation land in later increments.
- table string
- category string
- intent string
- labelColumn string
- targetSize string
- schemaFlag string
- labelPolicy string
- timeColumn string
+ table string
+ category string
+ intent string
+ labelColumn string
+ targetSize string
+ schemaFlag string
+ labelPolicy string
+ timeColumn string
+ numberOfKeypoints int
 
  // Operations flags.
  dryRun bool
@@ -153,6 +154,7 @@ Exit codes:
  Spec: push.SpecArgs{
  Table: table, Category: category, Intent: intent,
  LabelColumn: labelColumn, LabelPolicy: labelPolicy, TimeColumn: timeColumn,
+ NumberOfKeypoints: numberOfKeypoints,
  },
  TargetSizeFlag: targetSize,
  SchemaFlag: schemaFlag,
@@ -198,6 +200,8 @@ Exit codes:
  "passthrough|bucket (default bucket — bins the target so the raw value never leaves the cluster)")
 	cmd.Flags().StringVar(&timeColumn, "time-column", "",
  "time_to_event_prediction only: name of the time/duration column (default: a column named \"time\")")
+	cmd.Flags().IntVar(&numberOfKeypoints, "number-of-keypoints", 0,
+ "keypoint_detection only: number of keypoints per sample (required; e.g. 17 for COCO pose)")
 
 	cmd.Flags().BoolVar(&dryRun, "dry-run", false,
  "validate + discover + walk, but don't create any cluster resources")
@@ -280,21 +284,24 @@ func runDatasetPush(ctx context.Context, out, errOut io.Writer, a runDatasetPush
 	case a.Spec.Category == "":
  // Left empty by a caller; let the schema produce the canonical
  // "category is required" error downstream.
-	case push.IsTabular(a.Spec.Category) || push.IsText(a.Spec.Category) || a.Spec.Category == "image_classification":
+	case push.IsTabular(a.Spec.Category) || push.IsText(a.Spec.Category) ||
+ a.Spec.Category == "image_classification" ||
+ a.Spec.Category == "object_detection" ||
+ a.Spec.Category == "keypoint_detection":
  // supported
 	case push.IsImage(a.Spec.Category):
+ // semantic_segmentation / instance_segmentation
  return &exitError{code: 2, err: fmt.Errorf(
- "category %q isn't supported by the CLI yet — it needs annotation/mask "+
- "sidecar staging that's coming in a later release. Supported image "+
- "category: image_classification.", a.Spec.Category)}
+ "category %q isn't supported by the CLI yet. semantic_segmentation is "+
+ "blocked on the ingestor's mask-sidecar support (data-ingestors#136), and "+
+ "instance_segmentation isn't implemented. Supported image categories: "+
+ "image_classification, object_detection, keypoint_detection.", a.Spec.Category)}
 	default:
  return &exitError{code: 2, err: fmt.Errorf(
- "category %q isn't supported by the CLI yet. Supported: image_classification, "+
- "text_classification, masked_language_modeling, and the tabular / "+
- "time-series family (tabular_classification, tabular_regression, "+
- "time_series_forecasting, time_to_event_prediction). (Object detection / "+
- "keypoint / segmentation are coming; use the helm flow for those meanwhile.)",
- a.Spec.Category)}
+ "category %q isn't a recognized task category. Supported: image_classification, "+
+ "object_detection, keypoint_detection, text_classification, "+
+ "masked_language_modeling, tabular_classification, tabular_regression, "+
+ "time_series_forecasting, time_to_event_prediction.", a.Spec.Category)}
 	}
 
 	// 3. Walk the local directory FIRST (local "fail fast"), dispatched
@@ -312,7 +319,10 @@ func runDatasetPush(ctx context.Context, out, errOut io.Writer, a runDatasetPush
  layout, err = push.DiscoverTabular(a.LocalPath)
 	case push.IsText(a.Spec.Category):
  layout, err = push.DiscoverText(a.Spec.Category, a.LocalPath)
+	case a.Spec.Category == "object_detection":
+ layout, err = push.DiscoverObjectDetection(a.LocalPath)
 	default:
+ // image_classification + keypoint_detection: labels.csv + images/.
  layout, err = push.Discover(a.LocalPath)
 	}
 	if err != nil {
@@ -347,6 +357,14 @@ func runDatasetPush(ctx context.Context, out, errOut io.Writer, a runDatasetPush
  }
  }
 	case push.IsImage(a.Spec.Category):
+ // keypoint_detection needs --number-of-keypoints (dataset-
+ // specific, no default). Catch it here with an actionable
+ // message rather than letting the ingestor fail mid-run.
+ if a.Spec.Category == "keypoint_detection" && a.Spec.NumberOfKeypoints <= 0 {
+ return &exitError{code: 2, err: errors.New(
+ "keypoint_detection requires --number-of-keypoints (e.g. " +
+ "--number-of-keypoints 17); it's dataset-specific and has no default")}
+ }
  // Image target resolution: the ingestor's image_classification
  // default is 512x512 and it VALIDATES (it does not resize), so
  // a mismatch hard-fails. Honour an explicit --target-size;
@@ -633,6 +651,9 @@ func printPushPreflight(
 	default:
  _, _ = fmt.Fprintf(out, " labels.csv: %s\n", layout.LabelsCSV)
  _, _ = fmt.Fprintf(out, " images: %d files\n", len(layout.Images))
+ if anns := layout.Sidecars["annotations"]; len(anns) > 0 {
+ _, _ = fmt.Fprintf(out, " annotations: %d files\n", len(anns))
+ }
 	}
 	_, _ = fmt.Fprintf(out, " total size: %s\n", push.HumanBytes(layout.TotalBytes))
 	_, _ = fmt.Fprintln(out)

internal/cli/dataset_test.go

@@ -74,8 +74,8 @@ func execDatasetPush(t *testing.T, args []string) (exitCode int, stdout, stderr
 func TestDatasetPush_UnsupportedCategory_ExitsTwo(t *testing.T) {
 	root := imgcLayout(t)
 	for _, badCategory := range []string{
- "object_detection", // image category, needs annotation sidecar (later)
- "keypoint_detection", // image category, needs keypoint flags (later)
+ "semantic_segmentation", // blocked on the ingestor (data-ingestors#136)
+ "instance_segmentation", // not implemented
  "definitely-not-a-category", // nonsense; gate catches this too
 	} {
  t.Run(badCategory, func(t *testing.T) {

internal/push/image_extras.go

@@ -0,0 +1,51 @@
+package push
+
+import (
+	"fmt"
+	"path/filepath"
+)
+
+// xmlExtensions are the annotation file types object_detection reads
+// (Pascal VOC XML).
+var xmlExtensions = map[string]struct{}{".xml": {}}
+
+// DiscoverObjectDetection validates a local object_detection dataset:
+//
+// - <root>/labels.csv (required)
+// - <root>/images/* (required)
+// - <root>/annotations/*.xml (required; Pascal VOC)
+//
+// It builds on the image-classification layout (labels.csv + images/)
+// and adds the annotations/ sidecar via the shared sidecar walker, so
+// the existing tar/stream machinery stages annotations under
+// "annotations/".
+func DiscoverObjectDetection(rootDir string) (*LocalLayout, error) {
+	layout, err := Discover(rootDir) // labels.csv + images/ (+ caps + symlink guards)
+	if err != nil {
+ return nil, err
+	}
+
+	annotations, annoBytes, err := discoverSidecarFiles(layout.Root, "annotations", xmlExtensions)
+	if err != nil {
+ return nil, err
+	}
+	if len(annotations) == 0 {
+ return nil, fmt.Errorf(
+ "no .xml annotation files found in %q. object_detection expects "+
+ "<dir>/annotations/*.xml (Pascal VOC).",
+ filepath.Join(layout.Root, "annotations"))
+	}
+	if layout.Sidecars == nil {
+ layout.Sidecars = map[string][]string{}
+	}
+	layout.Sidecars["annotations"] = annotations
+	layout.TotalBytes += annoBytes
+
+	if layout.TotalBytes > MaxTotalBytes {
+ return nil, fmt.Errorf(
+ "dataset is %s, exceeds v0.1 cap of %s. For larger datasets, the "+
+ "cloud-source path is on the v0.2 roadmap (tracebloc/client#147).",
+ HumanBytes(layout.TotalBytes), HumanBytes(MaxTotalBytes))
+	}
+	return layout, nil
+}

internal/push/image_extras_test.go

@@ -0,0 +1,59 @@
+package push
+
+import (
+	"os"
+	"path/filepath"
+	"testing"
+)
+
+// mkODDir builds an object_detection dataset dir: labels.csv +
+// images/001.jpg, plus annotations/001.xml when withAnnotations.
+func mkODDir(t *testing.T, withAnnotations bool) string {
+	t.Helper()
+	dir := t.TempDir()
+	writeFile(t, dir, "labels.csv", "image_label,filename\ncat,001.jpg\n")
+	imgs := filepath.Join(dir, "images")
+	if err := os.MkdirAll(imgs, 0o755); err != nil {
+ t.Fatal(err)
+	}
+	// JPEG magic bytes; Discover checks extension + size, not decode.
+	if err := os.WriteFile(filepath.Join(imgs, "001.jpg"), []byte("\xff\xd8\xff\xe0"), 0o644); err != nil {
+ t.Fatal(err)
+	}
+	if withAnnotations {
+ ann := filepath.Join(dir, "annotations")
+ if err := os.MkdirAll(ann, 0o755); err != nil {
+ t.Fatal(err)
+ }
+ if err := os.WriteFile(filepath.Join(ann, "001.xml"), []byte("<annotation/>"), 0o644); err != nil {
+ t.Fatal(err)
+ }
+	}
+	return dir
+}
+
+// TestDiscoverObjectDetection: a valid OD layout yields images +
+// the annotations sidecar, staged together.
+func TestDiscoverObjectDetection(t *testing.T) {
+	layout, err := DiscoverObjectDetection(mkODDir(t, true))
+	if err != nil {
+ t.Fatalf("DiscoverObjectDetection: %v", err)
+	}
+	if len(layout.Images) != 1 {
+ t.Errorf("images = %d, want 1", len(layout.Images))
+	}
+	if len(layout.Sidecars["annotations"]) != 1 {
+ t.Errorf("annotations = %d, want 1", len(layout.Sidecars["annotations"]))
+	}
+	if got := layout.FileCount(); got != 3 { // labels.csv + image + xml
+ t.Errorf("FileCount = %d, want 3", got)
+	}
+}
+
+// TestDiscoverObjectDetection_MissingAnnotations: OD without an
+// annotations/ directory is a clear error.
+func TestDiscoverObjectDetection_MissingAnnotations(t *testing.T) {
+	if _, err := DiscoverObjectDetection(mkODDir(t, false)); err == nil {
+ t.Error("DiscoverObjectDetection without annotations/ returned nil error")
+	}
+}

internal/push/spec.go

@@ -162,6 +162,14 @@ type SpecArgs struct {
 	// Emitted as the top-level `time_column` field. Empty ⇒ the
 	// ingestor falls back to a column named "time".
 	TimeColumn string
+
+	// NumberOfKeypoints is the keypoints-per-sample count for
+	// keypoint_detection (required there; no convention default).
+	// Emitted under spec.file_options.number_of_keypoints — the schema
+	// allows arbitrary file_options keys, and conventions.resolve reads
+	// it from there, so this needs no top-level schema field. 0 ⇒
+	// unset (ignored for non-keypoint categories).
+	NumberOfKeypoints int
 }
 
 // Build produces the ingest.v1.json-conforming spec map. The
@@ -221,19 +229,38 @@ func (a SpecArgs) buildText(spec map[string]any, prefix string) {
 	}
 }
 
-// buildImage fills in the image-category fields: the images/ sidecar
-// dir, the label column, and the optional target_size override.
+// buildImage fills in the image-family fields for image_classification,
+// object_detection, and keypoint_detection: the images/ dir, the label,
+// object_detection's annotations/ dir, and the resolution overrides.
+//
+// keypoint_detection emits target_size + number_of_keypoints as
+// TOP-LEVEL fields — the schema's keypoint conditional requires them
+// there (both are dataset-specific, no convention defaults), and the
+// ingestor validates against that conditional. image_classification
+// and object_detection emit target_size under spec.file_options (the
+// override key conventions.resolve reads); without it,
+// image_classification defaults to 512x512 and the Image Resolution
+// Validator rejects other sizes.
 func (a SpecArgs) buildImage(spec map[string]any, prefix string) {
-	// Trailing slash on `images` matches the schema example
+	// Trailing slash on the dir fields matches the schema example
 	// (data-ingestors/examples/yaml/image_classification.yaml); the
-	// ingestor treats it as a directory glob.
+	// ingestor treats them as directory globs.
 	spec["images"] = path.Join(prefix, "images") + "/"
 	spec["label"] = a.LabelColumn
-	// Emit the image resolution under spec.file_options.target_size —
-	// the same override key data-ingestors' conventions.resolve
-	// honours. Without it, image_classification defaults to 512x512
-	// and the ingestor's Image Resolution Validator rejects any other
-	// size.
+	if a.Category == "object_detection" {
+ spec["annotations"] = path.Join(prefix, "annotations") + "/"
+	}
+
+	if a.Category == "keypoint_detection" {
+ if len(a.TargetSize) == 2 {
+ spec["target_size"] = []int{a.TargetSize[0], a.TargetSize[1]}
+ }
+ if a.NumberOfKeypoints > 0 {
+ spec["number_of_keypoints"] = a.NumberOfKeypoints
+ }
+ return
+	}
+
 	if len(a.TargetSize) == 2 {
  spec["spec"] = map[string]any{
  "file_options": map[string]any{

internal/push/spec_test.go

@@ -275,6 +275,61 @@ func TestBuild_Tabular_RegressionDefaultsPolicyBucket(t *testing.T) {
 	}
 }
 
+// TestBuild_ImageExtras_PassSchema pins object_detection and
+// keypoint_detection: OD emits an annotations field; keypoint emits
+// number_of_keypoints under spec.file_options (NOT a top-level field,
+// so it validates against the current vendored schema) and no
+// annotations. Both must validate.
+func TestBuild_ImageExtras_PassSchema(t *testing.T) {
+	v, err := schema.NewV1Validator()
+	if err != nil {
+ t.Fatalf("NewV1Validator: %v", err)
+	}
+	validate := func(t *testing.T, spec map[string]any) {
+ b, err := yaml.Marshal(spec)
+ if err != nil {
+ t.Fatalf("marshal: %v", err)
+ }
+ _, errs, parseErr := v.ValidateYAML(b)
+ if parseErr != nil {
+ t.Fatalf("parse: %v\n%s", parseErr, b)
+ }
+ if len(errs) != 0 {
+ t.Fatalf("schema validation failed: %s\n%s", schema.FormatErrors(errs), b)
+ }
+	}
+
+	t.Run("object_detection", func(t *testing.T) {
+ spec := SpecArgs{
+ Table: "od", Category: "object_detection", Intent: "train",
+ LabelColumn: "image_label", TargetSize: []int{448, 448},
+ }.Build()
+ if _, ok := spec["annotations"]; !ok {
+ t.Errorf("OD spec missing annotations: %v", spec)
+ }
+ validate(t, spec)
+	})
+
+	t.Run("keypoint_detection", func(t *testing.T) {
+ spec := SpecArgs{
+ Table: "kp", Category: "keypoint_detection", Intent: "train",
+ LabelColumn: "image_label", TargetSize: []int{448, 448}, NumberOfKeypoints: 9,
+ }.Build()
+ if _, ok := spec["annotations"]; ok {
+ t.Errorf("keypoint spec should not emit annotations: %v", spec)
+ }
+ // keypoint requires target_size + number_of_keypoints TOP-LEVEL
+ // (the schema's keypoint conditional), not under file_options.
+ if spec["number_of_keypoints"] != 9 {
+ t.Errorf("top-level number_of_keypoints = %v, want 9", spec["number_of_keypoints"])
+ }
+ if ts, ok := spec["target_size"].([]int); !ok || len(ts) != 2 || ts[0] != 448 || ts[1] != 448 {
+ t.Errorf("top-level target_size = %#v, want [448 448]", spec["target_size"])
+ }
+ validate(t, spec)
+	})
+}
+
 // TestValidateTableName_Accepts pins the names that MUST pass —
 // the real-world example tables plus a few edge shapes (single
 // char, leading underscore, mixed case, digits). A regression

internal/schema/ingest.v1.json

@@ -138,6 +138,20 @@
  "description": "Name of the time column for time_to_event_prediction. Falls back to a column named `time` if unset."
  },
 
+ "target_size": {
+ "type": "array",
+ "items": { "type": "integer", "minimum": 1 },
+ "minItems": 2,
+ "maxItems": 2,
+ "description": "[height, width] to resize images to. Required for keypoint_detection (no default — depends on the customer's pose model). Other image categories use category defaults if unset."
+ },
+
+ "number_of_keypoints": {
+ "type": "integer",
+ "minimum": 1,
+ "description": "Number of keypoints per sample. Required for keypoint_detection — dataset-specific (e.g. 17 for COCO pose, 9 for the upper-body sample in templates/keypoint_detection)."
+ },
+
  "data_id": {
  "type": "object",
  "additionalProperties": false,
@@ -356,6 +370,14 @@
  "required": ["label"]
  }
  },
+ {
+ "description": "keypoint_detection requires customer-supplied `target_size` and `number_of_keypoints` (no convention defaults — both are dataset-specific).",
+ "if": {
+ "properties": { "category": { "const": "keypoint_detection" } },
+ "required": ["category"]
+ },
+ "then": { "required": ["target_size", "number_of_keypoints"] }
+ },
  {
  "description": "Most categories require `label`.",
  "if": {