diff --git a/internal/cli/dataset.go b/internal/cli/dataset.go index a5da0e7e..15517673 100644 --- a/internal/cli/dataset.go +++ b/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) diff --git a/internal/cli/dataset_test.go b/internal/cli/dataset_test.go index ed4bef98..5d5cfcfd 100644 --- a/internal/cli/dataset_test.go +++ b/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) { diff --git a/internal/push/image_extras.go b/internal/push/image_extras.go new file mode 100644 index 00000000..5e1e0f15 --- /dev/null +++ b/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: +// +// - /labels.csv (required) +// - /images/* (required) +// - /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 "+ + "/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 +} diff --git a/internal/push/image_extras_test.go b/internal/push/image_extras_test.go new file mode 100644 index 00000000..8ed8d718 --- /dev/null +++ b/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(""), 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") + } +} diff --git a/internal/push/spec.go b/internal/push/spec.go index 4cf901ec..21339763 100644 --- a/internal/push/spec.go +++ b/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{ diff --git a/internal/push/spec_test.go b/internal/push/spec_test.go index 79f3b80e..21670085 100644 --- a/internal/push/spec_test.go +++ b/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 diff --git a/internal/schema/ingest.v1.json b/internal/schema/ingest.v1.json index 21a0c460..d8eb83c3 100644 --- a/internal/schema/ingest.v1.json +++ b/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": {