CDL ordinal label map

torchgeo/datasets/cdl.py

@@ -12,7 +12,7 @@
 from rasterio.crs import CRS
 
 from .geo import RasterDataset
-from .utils import download_url, extract_archive
+from .utils import BoundingBox, download_url, extract_archive
 
 
 class CDL(RasterDataset):
@@ -28,6 +28,9 @@ class CDL(RasterDataset):
     annually using moderate resolution satellite imagery and extensive agricultural
     ground truth.
 
+    The dataset contains 133 classes, for a description of the classes see the
+    xls file at the top of `this page <https://www.nass.usda.gov/Research_and_Science/Cropland/sarsfaqs2.php>`_.
+
     If you use this dataset in your research, please cite it using the following format:
 
     * https://www.nass.usda.gov/Research_and_Science/Cropland/sarsfaqs2.php#Section1_14.0
@@ -320,6 +323,143 @@ class CDL(RasterDataset):
         255: (0, 0, 0, 255),
     }
 
+    ordinal_label_map = {
+        0: 0,
+        1: 1,
+        2: 2,
+        3: 3,
+        4: 4,
+        5: 5,
+        6: 6,
+        10: 7,
+        11: 8,
+        12: 9,
+        13: 10,
+        14: 11,
+        21: 12,
+        22: 13,
+        23: 14,
+        24: 15,
+        25: 16,
+        26: 17,
+        27: 18,
+        28: 19,
+        29: 20,
+        30: 21,
+        31: 22,
+        32: 23,
+        33: 24,
+        34: 25,
+        35: 26,
+        36: 27,
+        37: 28,
+        38: 29,
+        39: 30,
+        41: 31,
+        42: 32,
+        43: 33,
+        44: 34,
+        45: 35,
+        46: 36,
+        47: 37,
+        48: 38,
+        49: 39,
+        50: 40,
+        51: 41,
+        52: 42,
+        53: 43,
+        54: 44,
+        55: 45,
+        56: 46,
+        57: 47,
+        58: 48,
+        59: 49,
+        60: 50,
+        61: 51,
+        63: 52,
+        64: 53,
+        65: 54,
+        66: 55,
+        67: 56,
+        68: 57,
+        69: 58,
+        70: 59,
+        71: 60,
+        72: 61,
+        74: 62,
+        75: 63,
+        76: 64,
+        77: 65,
+        81: 66,
+        82: 67,
+        83: 68,
+        87: 69,
+        88: 70,
+        92: 71,
+        111: 72,
+        112: 73,
+        121: 74,
+        122: 75,
+        123: 76,
+        124: 77,
+        131: 78,
+        141: 79,
+        142: 80,
+        143: 81,
+        152: 82,
+        176: 83,
+        190: 84,
+        195: 85,
+        204: 86,
+        205: 87,
+        206: 88,
+        207: 89,
+        208: 90,
+        209: 91,
+        210: 92,
+        211: 93,
+        212: 94,
+        213: 95,
+        214: 96,
+        215: 97,
+        216: 98,
+        217: 99,
+        218: 100,
+        219: 101,
+        220: 102,
+        221: 103,
+        222: 104,
+        223: 105,
+        224: 106,
+        225: 107,
+        226: 108,
+        227: 109,
+        228: 110,
+        229: 111,
+        230: 112,
+        231: 113,
+        232: 114,
+        233: 115,
+        234: 116,
+        235: 117,
+        236: 118,
+        237: 119,
+        238: 120,
+        239: 121,
+        240: 122,
+        241: 123,
+        242: 124,
+        243: 125,
+        244: 126,
+        245: 127,
+        246: 128,
+        247: 129,
+        248: 130,
+        249: 131,
+        250: 132,
+        254: 133,
+    }
+
     def __init__(
         self,
         root: str = "data",
@@ -356,6 +496,28 @@ def __init__(
 
         super().__init__(root, crs, res, transforms=transforms, cache=cache)
 
+    def __getitem__(self, query: BoundingBox) -> dict[str, Any]:
+        """Retrieve mask and metadata indexed by query.
+
+        Args:
+            query: (minx, maxx, miny, maxy, mint, maxt) coordinates to index
+
+        Returns:
+            sample of mask and metadata at that index
+
+        Raises:
+            IndexError: if query is not found in the index
+        """
+        sample = super().__getitem__(query)
+
+        mask = sample["mask"]
+        for k, v in self.ordinal_label_map.items():
+            mask[mask == k] = v
+
+        sample["mask"] = mask
+
+        return sample
+
     def _verify(self) -> None:
         """Verify the integrity of the dataset.