rustytorch/demos/segmentation-shared/src/class_sets.rs

// Predefined class sets for common segmentation datasets

use crate::ClassInfo;

/// PASCAL VOC 2012 class set (21 classes)
pub const PASCAL_VOC_CLASSES: &[ClassInfo] = &[
    ClassInfo {
        id: 0,
        name: "background",
        color: "#000000",
    },
    ClassInfo {
        id: 1,
        name: "aeroplane",
        color: "#800000",
    },
    ClassInfo {
        id: 2,
        name: "bicycle",
        color: "#008000",
    },
    ClassInfo {
        id: 3,
        name: "bird",
        color: "#808000",
    },
    ClassInfo {
        id: 4,
        name: "boat",
        color: "#000080",
    },
    ClassInfo {
        id: 5,
        name: "bottle",
        color: "#800080",
    },
    ClassInfo {
        id: 6,
        name: "bus",
        color: "#008080",
    },
    ClassInfo {
        id: 7,
        name: "car",
        color: "#808080",
    },
    ClassInfo {
        id: 8,
        name: "cat",
        color: "#C00000",
    },
    ClassInfo {
        id: 9,
        name: "chair",
        color: "#00C000",
    },
    ClassInfo {
        id: 10,
        name: "cow",
        color: "#C0C000",
    },
    ClassInfo {
        id: 11,
        name: "diningtable",
        color: "#0000C0",
    },
    ClassInfo {
        id: 12,
        name: "dog",
        color: "#C000C0",
    },
    ClassInfo {
        id: 13,
        name: "horse",
        color: "#00C0C0",
    },
    ClassInfo {
        id: 14,
        name: "motorbike",
        color: "#C0C0C0",
    },
    ClassInfo {
        id: 15,
        name: "person",
        color: "#FF0000",
    },
    ClassInfo {
        id: 16,
        name: "pottedplant",
        color: "#00FF00",
    },
    ClassInfo {
        id: 17,
        name: "sheep",
        color: "#FFFF00",
    },
    ClassInfo {
        id: 18,
        name: "sofa",
        color: "#0000FF",
    },
    ClassInfo {
        id: 19,
        name: "train",
        color: "#FF00FF",
    },
    ClassInfo {
        id: 20,
        name: "tvmonitor",
        color: "#00FFFF",
    },
];

/// ADE20K dataset class set (150 classes - subset shown)
/// Full dataset contains 150 classes, showing first 30 for brevity
pub const ADE20K_CLASSES: &[ClassInfo] = &[
    ClassInfo {
        id: 0,
        name: "background",
        color: "#000000",
    },
    ClassInfo {
        id: 1,
        name: "wall",
        color: "#7F0000",
    },
    ClassInfo {
        id: 2,
        name: "building",
        color: "#007F00",
    },
    ClassInfo {
        id: 3,
        name: "sky",
        color: "#7F7F00",
    },
    ClassInfo {
        id: 4,
        name: "floor",
        color: "#00007F",
    },
    ClassInfo {
        id: 5,
        name: "tree",
        color: "#7F007F",
    },
    ClassInfo {
        id: 6,
        name: "ceiling",
        color: "#007F7F",
    },
    ClassInfo {
        id: 7,
        name: "road",
        color: "#7F7F7F",
    },
    ClassInfo {
        id: 8,
        name: "bed",
        color: "#BF0000",
    },
    ClassInfo {
        id: 9,
        name: "windowpane",
        color: "#00BF00",
    },
    ClassInfo {
        id: 10,
        name: "grass",
        color: "#BFBF00",
    },
    ClassInfo {
        id: 11,
        name: "cabinet",
        color: "#0000BF",
    },
    ClassInfo {
        id: 12,
        name: "sidewalk",
        color: "#BF00BF",
    },
    ClassInfo {
        id: 13,
        name: "person",
        color: "#00BFBF",
    },
    ClassInfo {
        id: 14,
        name: "earth",
        color: "#BFBFBF",
    },
    ClassInfo {
        id: 15,
        name: "door",
        color: "#FF4000",
    },
    ClassInfo {
        id: 16,
        name: "table",
        color: "#40FF00",
    },
    ClassInfo {
        id: 17,
        name: "mountain",
        color: "#FFFF40",
    },
    ClassInfo {
        id: 18,
        name: "plant",
        color: "#4000FF",
    },
    ClassInfo {
        id: 19,
        name: "curtain",
        color: "#FF40FF",
    },
    ClassInfo {
        id: 20,
        name: "chair",
        color: "#40FFFF",
    },
    ClassInfo {
        id: 21,
        name: "car",
        color: "#FF8000",
    },
    ClassInfo {
        id: 22,
        name: "water",
        color: "#80FF00",
    },
    ClassInfo {
        id: 23,
        name: "painting",
        color: "#FFFF80",
    },
    ClassInfo {
        id: 24,
        name: "sofa",
        color: "#8000FF",
    },
    ClassInfo {
        id: 25,
        name: "shelf",
        color: "#FF80FF",
    },
    ClassInfo {
        id: 26,
        name: "house",
        color: "#80FFFF",
    },
    ClassInfo {
        id: 27,
        name: "sea",
        color: "#FFC000",
    },
    ClassInfo {
        id: 28,
        name: "mirror",
        color: "#C0FF00",
    },
    ClassInfo {
        id: 29,
        name: "rug",
        color: "#FFFFC0",
    },
];

/// Cityscapes dataset class set (19 classes)
pub const CITYSCAPES_CLASSES: &[ClassInfo] = &[
    ClassInfo {
        id: 0,
        name: "road",
        color: "#804080",
    },
    ClassInfo {
        id: 1,
        name: "sidewalk",
        color: "#F423E8",
    },
    ClassInfo {
        id: 2,
        name: "building",
        color: "#464646",
    },
    ClassInfo {
        id: 3,
        name: "wall",
        color: "#6496C8",
    },
    ClassInfo {
        id: 4,
        name: "fence",
        color: "#BE9664",
    },
    ClassInfo {
        id: 5,
        name: "pole",
        color: "#999999",
    },
    ClassInfo {
        id: 6,
        name: "traffic_light",
        color: "#FAE619",
    },
    ClassInfo {
        id: 7,
        name: "traffic_sign",
        color: "#DCE620",
    },
    ClassInfo {
        id: 8,
        name: "vegetation",
        color: "#6B8E23",
    },
    ClassInfo {
        id: 9,
        name: "terrain",
        color: "#98FB98",
    },
    ClassInfo {
        id: 10,
        name: "sky",
        color: "#4682B4",
    },
    ClassInfo {
        id: 11,
        name: "person",
        color: "#DC143C",
    },
    ClassInfo {
        id: 12,
        name: "rider",
        color: "#FF0000",
    },
    ClassInfo {
        id: 13,
        name: "car",
        color: "#00008B",
    },
    ClassInfo {
        id: 14,
        name: "truck",
        color: "#1E1E46",
    },
    ClassInfo {
        id: 15,
        name: "bus",
        color: "#3C3C64",
    },
    ClassInfo {
        id: 16,
        name: "train",
        color: "#505064",
    },
    ClassInfo {
        id: 17,
        name: "motorcycle",
        color: "#0000E6",
    },
    ClassInfo {
        id: 18,
        name: "bicycle",
        color: "#770BB3",
    },
];

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_pascal_voc_has_21_classes() {
        assert_eq!(PASCAL_VOC_CLASSES.len(), 21);
    }

    #[test]
    fn test_pascal_voc_class_ids_sequential() {
        for (i, class) in PASCAL_VOC_CLASSES.iter().enumerate() {
            assert_eq!(class.id, i);
        }
    }

    #[test]
    fn test_pascal_voc_all_colors_valid() {
        for class in PASCAL_VOC_CLASSES {
            assert!(
                class.validate_color().is_ok(),
                "Invalid color for class {}: {}",
                class.name,
                class.color
            );
        }
    }

    #[test]
    fn test_pascal_voc_background_is_first() {
        assert_eq!(PASCAL_VOC_CLASSES[0].name, "background");
        assert_eq!(PASCAL_VOC_CLASSES[0].id, 0);
    }

    #[test]
    fn test_pascal_voc_person_class_exists() {
        let person_class = PASCAL_VOC_CLASSES.iter().find(|c| c.name == "person");
        assert!(person_class.is_some());
        assert_eq!(person_class.unwrap().id, 15);
    }

    #[test]
    fn test_pascal_voc_all_names_non_empty() {
        for class in PASCAL_VOC_CLASSES {
            assert!(!class.name.is_empty());
        }
    }

    #[test]
    fn test_pascal_voc_colors_are_hex() {
        for class in PASCAL_VOC_CLASSES {
            assert!(class.color.starts_with('#'));
            assert_eq!(class.color.len(), 7);
        }
    }

    #[test]
    fn test_ade20k_has_30_classes() {
        assert_eq!(ADE20K_CLASSES.len(), 30);
    }

    #[test]
    fn test_ade20k_class_ids_sequential() {
        for (i, class) in ADE20K_CLASSES.iter().enumerate() {
            assert_eq!(class.id, i);
        }
    }

    #[test]
    fn test_ade20k_all_colors_valid() {
        for class in ADE20K_CLASSES {
            assert!(
                class.validate_color().is_ok(),
                "Invalid color for class {}: {}",
                class.name,
                class.color
            );
        }
    }

    #[test]
    fn test_ade20k_background_is_first() {
        assert_eq!(ADE20K_CLASSES[0].name, "background");
        assert_eq!(ADE20K_CLASSES[0].id, 0);
    }

    #[test]
    fn test_ade20k_person_class_exists() {
        let person_class = ADE20K_CLASSES.iter().find(|c| c.name == "person");
        assert!(person_class.is_some());
        assert_eq!(person_class.unwrap().id, 13);
    }

    #[test]
    fn test_ade20k_all_names_non_empty() {
        for class in ADE20K_CLASSES {
            assert!(!class.name.is_empty());
        }
    }

    #[test]
    fn test_cityscapes_has_19_classes() {
        assert_eq!(CITYSCAPES_CLASSES.len(), 19);
    }

    #[test]
    fn test_cityscapes_class_ids_sequential() {
        for (i, class) in CITYSCAPES_CLASSES.iter().enumerate() {
            assert_eq!(class.id, i);
        }
    }

    #[test]
    fn test_cityscapes_all_colors_valid() {
        for class in CITYSCAPES_CLASSES {
            assert!(
                class.validate_color().is_ok(),
                "Invalid color for class {}: {}",
                class.name,
                class.color
            );
        }
    }

    #[test]
    fn test_cityscapes_road_is_first() {
        assert_eq!(CITYSCAPES_CLASSES[0].name, "road");
        assert_eq!(CITYSCAPES_CLASSES[0].id, 0);
    }

    #[test]
    fn test_cityscapes_person_class_exists() {
        let person_class = CITYSCAPES_CLASSES.iter().find(|c| c.name == "person");
        assert!(person_class.is_some());
        assert_eq!(person_class.unwrap().id, 11);
    }

    #[test]
    fn test_cityscapes_all_names_non_empty() {
        for class in CITYSCAPES_CLASSES {
            assert!(!class.name.is_empty());
        }
    }

    #[test]
    fn test_cityscapes_vehicle_classes_exist() {
        let vehicle_classes = ["car", "truck", "bus", "train", "motorcycle", "bicycle"];
        for vehicle in vehicle_classes {
            let class = CITYSCAPES_CLASSES.iter().find(|c| c.name == vehicle);
            assert!(class.is_some(), "Vehicle class '{}' not found", vehicle);
        }
    }

    #[test]
    fn test_all_class_sets_have_unique_ids() {
        let check_unique = |classes: &[ClassInfo], name: &str| {
            let mut ids = classes.iter().map(|c| c.id).collect::<Vec<_>>();
            ids.sort_unstable();
            ids.dedup();
            assert_eq!(ids.len(), classes.len(), "{} has duplicate IDs", name);
        };

        check_unique(PASCAL_VOC_CLASSES, "PASCAL_VOC_CLASSES");
        check_unique(ADE20K_CLASSES, "ADE20K_CLASSES");
        check_unique(CITYSCAPES_CLASSES, "CITYSCAPES_CLASSES");
    }

    #[test]
    fn test_all_class_sets_can_convert_to_rgb() {
        let test_set = |classes: &[ClassInfo], name: &str| {
            for class in classes {
                assert!(
                    class.to_rgb().is_ok(),
                    "{} class {} has invalid color",
                    name,
                    class.name
                );
            }
        };

        test_set(PASCAL_VOC_CLASSES, "PASCAL_VOC");
        test_set(ADE20K_CLASSES, "ADE20K");
        test_set(CITYSCAPES_CLASSES, "CITYSCAPES");
    }
}