Repository: microsoft/torchgeo

Issue:
Chesapeake datasets don't extract properly
The ChesapeakeMD dataset fails when attempting to extract the downloaded zip file as `zipfile.ZipFile` doesn't support the deflate64 compression type that `_MD_STATEWIDE.zip` uses.


Discussion:
I tested the rest of them, the following don't work for the same reason:
- ChesapeakeNY
- ChesapeakeVA

I also directly committed a fix for the Chesapeake7 dataset to main.
I think `zipfile.ZipFile` does support deflate64, you just have to explicitly pass `compression=zipfile.ZIP_DEFLATED`. See https://docs.python.org/3/library/zipfile.html#zipfile.ZIP_DEFLATED.

I was hoping this would be done automatically. Torchvision had some logic to inject this but it wasn't clear to me how it worked. Will investigate more.
So it turns out deflate and deflate64 are very different. The latter is a proprietary compression scheme that is not supported by zipfile. We could use https://pypi.org/project/zipfile-deflate64/, or call `unzip` or `7z` from the command line.

Resolving patch:
diff --git a/torchgeo/datasets/chesapeake.py b/torchgeo/datasets/chesapeake.py
--- a/torchgeo/datasets/chesapeake.py
+++ b/torchgeo/datasets/chesapeake.py
@@ -233,7 +233,15 @@ class ChesapeakeDE(Chesapeake):
 
 
 class ChesapeakeMD(Chesapeake):
-    """This subset of the dataset contains data only for Maryland."""
+    """This subset of the dataset contains data only for Maryland.
+
+    .. note::
+
+       This dataset requires the following additional library to be installed:
+
+       * `zipfile-deflate64 <https://pypi.org/project/zipfile-deflate64/>`_ to extract
+         the proprietary deflate64 compressed zip file.
+    """
 
     base_folder = "MD"
     filename = "MD_STATEWIDE.tif"
@@ -242,7 +250,15 @@ class ChesapeakeMD(Chesapeake):
 
 
 class ChesapeakeNY(Chesapeake):
-    """This subset of the dataset contains data only for New York."""
+    """This subset of the dataset contains data only for New York.
+
+    .. note::
+
+       This dataset requires the following additional library to be installed:
+
+       * `zipfile-deflate64 <https://pypi.org/project/zipfile-deflate64/>`_ to extract
+         the proprietary deflate64 compressed zip file.
+    """
 
     base_folder = "NY"
     filename = "NY_STATEWIDE.tif"
@@ -260,7 +276,15 @@ class ChesapeakePA(Chesapeake):
 
 
 class ChesapeakeVA(Chesapeake):
-    """This subset of the dataset contains data only for Virginia."""
+    """This subset of the dataset contains data only for Virginia.
+
+    .. note::
+
+       This dataset requires the following additional library to be installed:
+
+       * `zipfile-deflate64 <https://pypi.org/project/zipfile-deflate64/>`_ to extract
+         the proprietary deflate64 compressed zip file.
+    """
 
     base_folder = "VA"
     filename = "CIC2014_VA_STATEWIDE.tif"
diff --git a/torchgeo/datasets/utils.py b/torchgeo/datasets/utils.py
--- a/torchgeo/datasets/utils.py
+++ b/torchgeo/datasets/utils.py
@@ -11,7 +11,6 @@
 import os
 import sys
 import tarfile
-import zipfile
 from dataclasses import dataclass
 from datetime import datetime, timedelta
 from typing import (
@@ -77,6 +76,27 @@ def __exit__(self, exc_type: None, exc_value: None, traceback: None) -> None:
             pass
 
 
+class _zipfile:
+    class ZipFile:
+        def __init__(self, *args: Any, **kwargs: Any) -> None:
+            self.args = args
+            self.kwargs = kwargs
+
+        def __enter__(self) -> Any:
+            try:
+                # Supports normal zip files, proprietary deflate64 compression algorithm
+                import zipfile_deflate64 as zipfile
+            except ImportError:
+                # Only supports normal zip files
+                # https://github.com/python/mypy/issues/1153
+                import zipfile  # type: ignore[no-redef]
+
+            return zipfile.ZipFile(*self.args, **self.kwargs)
+
+        def __exit__(self, exc_type: None, exc_value: None, traceback: None) -> None:
+            pass
+
+
 def extract_archive(src: str, dst: Optional[str] = None) -> None:
     """Extract an archive.
 
@@ -96,7 +116,7 @@ def extract_archive(src: str, dst: Optional[str] = None) -> None:
             (".tar", ".tar.gz", ".tar.bz2", ".tar.xz", ".tgz", ".tbz2", ".tbz", ".txz"),
             tarfile.open,
         ),
-        (".zip", zipfile.ZipFile),
+        (".zip", _zipfile.ZipFile),
     ]
 
     for suffix, extractor in suffix_and_extractor:


Test patch:
diff --git a/tests/conf/naipchesapeake.yaml b/tests/conf/naipchesapeake.yaml
--- a/tests/conf/naipchesapeake.yaml
+++ b/tests/conf/naipchesapeake.yaml
@@ -9,7 +9,7 @@ experiment:
     learning_rate: 1e-3
     learning_rate_schedule_patience: 2
     in_channels: 4
-    num_classes: 13
+    num_classes: 14
     num_filters: 1
     ignore_zeros: False
   datamodule:
diff --git a/tests/data/chesapeake/BAYWIDE/Baywide_13Class_20132014.tif b/tests/data/chesapeake/BAYWIDE/Baywide_13Class_20132014.tif
Binary files a/tests/data/chesapeake/BAYWIDE/Baywide_13Class_20132014.tif and b/tests/data/chesapeake/BAYWIDE/Baywide_13Class_20132014.tif differ
diff --git a/tests/data/chesapeake/BAYWIDE/Baywide_13Class_20132014.zip b/tests/data/chesapeake/BAYWIDE/Baywide_13Class_20132014.zip
Binary files a/tests/data/chesapeake/BAYWIDE/Baywide_13Class_20132014.zip and b/tests/data/chesapeake/BAYWIDE/Baywide_13Class_20132014.zip differ
diff --git a/tests/data/chesapeake/BAYWIDE/data.py b/tests/data/chesapeake/BAYWIDE/data.py
new file mode 100755
--- /dev/null
+++ b/tests/data/chesapeake/BAYWIDE/data.py
@@ -0,0 +1,95 @@
+#!/usr/bin/env python3
+
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+import hashlib
+import os
+import subprocess
+
+import numpy as np
+import rasterio
+from rasterio.crs import CRS
+from rasterio.transform import Affine
+
+SIZE = 128  # image width/height
+NUM_CLASSES = 14
+
+np.random.seed(0)
+
+filename = "Baywide_13Class_20132014"
+wkt = """
+PROJCS["USA_Contiguous_Albers_Equal_Area_Conic_USGS_version",
+    GEOGCS["NAD83",
+        DATUM["North_American_Datum_1983",
+            SPHEROID["GRS 1980",6378137,298.257222101004,
+                AUTHORITY["EPSG","7019"]],
+            AUTHORITY["EPSG","6269"]],
+        PRIMEM["Greenwich",0],
+        UNIT["degree",0.0174532925199433,
+            AUTHORITY["EPSG","9122"]],
+        AUTHORITY["EPSG","4269"]],
+    PROJECTION["Albers_Conic_Equal_Area"],
+    PARAMETER["latitude_of_center",23],
+    PARAMETER["longitude_of_center",-96],
+    PARAMETER["standard_parallel_1",29.5],
+    PARAMETER["standard_parallel_2",45.5],
+    PARAMETER["false_easting",0],
+    PARAMETER["false_northing",0],
+    UNIT["metre",1,
+        AUTHORITY["EPSG","9001"]],
+    AXIS["Easting",EAST],
+    AXIS["Northing",NORTH]]
+"""
+cmap = {
+    0: (0, 0, 0, 255),
+    1: (0, 197, 255, 255),
+    2: (0, 168, 132, 255),
+    3: (38, 115, 0, 255),
+    4: (76, 230, 0, 255),
+    5: (163, 255, 115, 255),
+    6: (255, 170, 0, 255),
+    7: (255, 0, 0, 255),
+    8: (156, 156, 156, 255),
+    9: (0, 0, 0, 255),
+    10: (115, 115, 0, 255),
+    11: (230, 230, 0, 255),
+    12: (255, 255, 115, 255),
+    13: (197, 0, 255, 255),
+}
+
+
+meta = {
+    "driver": "GTiff",
+    "dtype": "uint8",
+    "nodata": None,
+    "width": SIZE,
+    "height": SIZE,
+    "count": 1,
+    "crs": CRS.from_wkt(wkt),
+    "transform": Affine(1.0, 0.0, 1303555.0000000005, 0.0, -1.0, 2535064.999999998),
+}
+
+# Remove old data
+if os.path.exists(f"{filename}.tif"):
+    os.remove(f"{filename}.tif")
+
+# Create raster file
+with rasterio.open(f"{filename}.tif", "w", **meta) as f:
+    data = np.random.randint(NUM_CLASSES, size=(SIZE, SIZE), dtype=np.uint8)
+    f.write(data, 1)
+    f.write_colormap(1, cmap)
+
+# Create zip file
+# 7z required to create a zip file using the proprietary DEFLATE64 compression algorithm
+# https://github.com/brianhelba/zipfile-deflate64/issues/19#issuecomment-1006077294
+subprocess.run(
+    ["7z", "a", f"{filename}.zip", "-mm=DEFLATE64", f"{filename}.tif"],
+    capture_output=True,
+    check=True,
+)
+
+# Compute checksums
+with open(f"{filename}.zip", "rb") as f:
+    md5 = hashlib.md5(f.read()).hexdigest()
+    print(repr(md5))
diff --git a/tests/datasets/test_chesapeake.py b/tests/datasets/test_chesapeake.py
--- a/tests/datasets/test_chesapeake.py
+++ b/tests/datasets/test_chesapeake.py
@@ -33,10 +33,11 @@ class TestChesapeake13:
     def dataset(
         self, monkeypatch: Generator[MonkeyPatch, None, None], tmp_path: Path
     ) -> Chesapeake13:
+        pytest.importorskip("zipfile_deflate64")
         monkeypatch.setattr(  # type: ignore[attr-defined]
             torchgeo.datasets.chesapeake, "download_url", download_url
         )
-        md5 = "9557b609e614a1f79dec6eb1bb3f3a06"
+        md5 = "fe35a615b8e749b21270472aa98bb42c"
         monkeypatch.setattr(Chesapeake13, "md5", md5)  # type: ignore[attr-defined]
         url = os.path.join(
             "tests", "data", "chesapeake", "BAYWIDE", "Baywide_13Class_20132014.zip"
diff --git a/tests/datasets/test_utils.py b/tests/datasets/test_utils.py
--- a/tests/datasets/test_utils.py
+++ b/tests/datasets/test_utils.py
@@ -41,7 +41,7 @@ def mock_missing_module(monkeypatch: Generator[MonkeyPatch, None, None]) -> None
     import_orig = builtins.__import__
 
     def mocked_import(name: str, *args: Any, **kwargs: Any) -> Any:
-        if name in ["rarfile", "radiant_mlhub"]:
+        if name in ["radiant_mlhub", "rarfile", "zipfile_deflate64"]:
             raise ImportError()
         return import_orig(name, *args, **kwargs)
 
@@ -93,11 +93,15 @@ def test_mock_missing_module(mock_missing_module: None) -> None:
         os.path.join("cowc_detection", "COWC_test_list_detection.txt.bz2"),
         os.path.join("vhr10", "NWPU VHR-10 dataset.rar"),
         os.path.join("landcoverai", "landcover.ai.v1.zip"),
+        os.path.join("chesapeake", "BAYWIDE", "Baywide_13Class_20132014.zip"),
         os.path.join("sen12ms", "ROIs1158_spring_lc.tar.gz"),
     ],
 )
 def test_extract_archive(src: str, tmp_path: Path) -> None:
-    pytest.importorskip("rarfile", minversion="3")
+    if src.endswith(".rar"):
+        pytest.importorskip("rarfile", minversion="3")
+    if src.startswith("chesapeake"):
+        pytest.importorskip("zipfile_deflate64")
     extract_archive(os.path.join("tests", "data", src), str(tmp_path))
 
 
@@ -111,6 +115,11 @@ def test_missing_rarfile(mock_missing_module: None) -> None:
         )
 
 
+def test_missing_zipfile_deflate64(mock_missing_module: None) -> None:
+    # Should fallback on Python builtin zipfile
+    extract_archive(os.path.join("tests", "data", "landcoverai", "landcover.ai.v1.zip"))
+
+
 def test_unsupported_scheme() -> None:
     with pytest.raises(
         RuntimeError, match="src file has unknown archival/compression scheme"
diff --git a/tests/trainers/test_segmentation.py b/tests/trainers/test_segmentation.py
--- a/tests/trainers/test_segmentation.py
+++ b/tests/trainers/test_segmentation.py
@@ -50,6 +50,9 @@ def test_trainer(
         name: str,
         classname: Type[LightningDataModule],
     ) -> None:
+        if name == "naipchesapeake":
+            pytest.importorskip("zipfile_deflate64")
+
         conf = OmegaConf.load(os.path.join("tests", "conf", name + ".yaml"))
         conf_dict = OmegaConf.to_object(conf.experiment)
         conf_dict = cast(Dict[Any, Dict[Any, Any]], conf_dict)