fix(yolo): emit accurate dataset YAML (#68)

Author: amanharshx

src-tauri/src/converter/yolo.rs

@@ -6,6 +6,25 @@ pub struct YoloConverter {
     darknet: bool,
 }
 
+struct SplitPresence {
+    train: bool,
+    val: bool,
+    test: bool,
+}
+
+impl SplitPresence {
+    fn from_data(data: &NDJSONData) -> Self {
+        Self {
+            train: data.images.iter().any(|img| img.split == "train"),
+            val: data
+                .images
+                .iter()
+                .any(|img| img.split == "valid" || img.split == "val"),
+            test: data.images.iter().any(|img| img.split == "test"),
+        }
+    }
+}
+
 fn quote_yaml_scalar(value: &str) -> String {
     serde_json::to_string(value).unwrap_or_else(|_| "\"\"".to_string())
 }
@@ -40,15 +59,26 @@ impl YoloConverter {
     fn create_data_yaml(&self, data: &NDJSONData) -> String {
         let class_names = get_class_list(data);
         let task = &data.metadata.task;
+        // Split presence reflects data.images after download filtering when images are requested.
+        let splits = SplitPresence::from_data(data);
 
         let mut yaml = String::new();
         yaml.push_str("# Generated by YOLO NDJSON Converter\n");
         yaml.push_str(&format!("# Version: {}\n", env!("CARGO_PKG_VERSION")));
         yaml.push_str("# https://github.com/amanharshx/YOLO-Ndjson-Zip\n\n");
         yaml.push_str("path: .\n");
-        yaml.push_str("train: train/images\n");
-        yaml.push_str("val: valid/images\n");
-        yaml.push_str("test: test/images\n");
+        if splits.train {
+            yaml.push_str("train: train/images\n");
+        }
+        if splits.val {
+            yaml.push_str("val: valid/images\n");
+        }
+        if splits.test {
+            yaml.push_str("test: test/images\n");
+        }
+        if !splits.train || !splits.val {
+            yaml.push_str("# WARNING: Ultralytics training requires both train and val splits.\n");
+        }
         yaml.push_str(&format!("nc: {}\n", class_names.len()));
         yaml.push_str("names:\n");
 
@@ -155,10 +185,12 @@ impl Converter for YoloConverter {
             );
         } else {
             // Standard YOLO mode
-            files.insert(
-                "data.yaml".to_string(),
-                self.create_data_yaml(data).into_bytes(),
-            );
+            if task != "classify" {
+                files.insert(
+                    "data.yaml".to_string(),
+                    self.create_data_yaml(data).into_bytes(),
+                );
+            }
             let class_list = get_class_list(data);
             files.insert(
                 "classes.txt".to_string(),
@@ -283,6 +315,92 @@ mod tests {
         }
     }
 
+    fn image_with_split(split: &str) -> ImageEntry {
+        ImageEntry {
+            r#type: "image".to_string(),
+            file: format!("{}.jpg", split),
+            output_file: None,
+            url: String::new(),
+            width: 640,
+            height: 480,
+            split: split.to_string(),
+            annotations: Some(json!({
+                "bboxes": [[0, 0.5, 0.5, 0.2, 0.2]]
+            })),
+        }
+    }
+
+    #[test]
+    fn data_yaml_emits_only_present_splits() {
+        let cases = [
+            (&["train"][..], [true, false, false], true),
+            (&["train", "val"][..], [true, true, false], false),
+            (&["train", "valid", "test"][..], [true, true, true], false),
+            (&["train", "test"][..], [true, false, true], true),
+            (&["val", "test"][..], [false, true, true], true),
+        ];
+
+        for (source_splits, expected, warns) in cases {
+            let images = source_splits
+                .iter()
+                .map(|split| image_with_split(split))
+                .collect();
+            let data = make_data("detect", HashMap::new(), None, images);
+            let files = YoloConverter::new().convert(&data, &HashMap::new());
+            let yaml = std::str::from_utf8(files.get("data.yaml").unwrap()).unwrap();
+
+            let has_line = |line| yaml.lines().any(|candidate| candidate == line);
+            assert_eq!(
+                [
+                    has_line("train: train/images"),
+                    has_line("val: valid/images"),
+                    has_line("test: test/images"),
+                ],
+                expected,
+                "unexpected splits in:\n{yaml}"
+            );
+            assert_eq!(yaml.contains("# WARNING:"), warns);
+        }
+    }
+
+    #[test]
+    fn data_yaml_omits_split_whose_images_all_failed() {
+        let mut data = make_data(
+            "detect",
+            HashMap::new(),
+            None,
+            vec![image_with_split("train"), image_with_split("val")],
+        );
+        data.images.retain(|img| img.split == "train");
+
+        let files = YoloConverter::new().convert(&data, &HashMap::new());
+        let yaml = std::str::from_utf8(files.get("data.yaml").unwrap()).unwrap();
+
+        assert!(yaml.contains("train: train/images"));
+        assert!(!yaml.contains("val: valid/images"));
+        assert!(!files.contains_key("valid/labels/val.txt"));
+    }
+
+    #[test]
+    fn classification_conversion_omits_yaml_metadata() {
+        let data = make_data(
+            "classify",
+            HashMap::from([("0".to_string(), "animal".to_string())]),
+            None,
+            vec![ImageEntry {
+                annotations: Some(json!({ "classification": [0] })),
+                ..image_with_split("train")
+            }],
+        );
+        let downloaded_images =
+            HashMap::from([(image_download_key("train", "train.jpg"), vec![1, 2, 3])]);
+
+        let files = YoloConverter::new().convert(&data, &downloaded_images);
+
+        assert!(files.contains_key("train/animal/train.jpg"));
+        assert!(files.keys().all(|path| !path.ends_with(".yaml")));
+    }
+
     #[test]
     fn create_data_yaml_quotes_class_names() {
         let mut class_names = HashMap::new();