improve: cut sequences by max 6 gigapixel of images (mapillary#616)

ptpt · web-flow · commit 58e66e938a36 · 2023-05-16T17:14:54.000-07:00
* improve: cut sequences by max 6 gigapixel of images

* fix types

* add tests

* fix tests
diff --git a/mapillary_tools/constants.py b/mapillary_tools/constants.py
@@ -44,3 +44,5 @@
 MAX_SEQUENCE_LENGTH = int(os.getenv(_ENV_PREFIX + "MAX_SEQUENCE_LENGTH", 500))
 # Max file size per sequence (sum of image filesizes in the sequence)
 MAX_SEQUENCE_FILESIZE: str = os.getenv(_ENV_PREFIX + "MAX_SEQUENCE_FILESIZE", "2G")
+# Max number of pixels per sequence (sum of image pixels in the sequence)
+MAX_SEQUENCE_PIXELS: str = os.getenv(_ENV_PREFIX + "MAX_SEQUENCE_PIXELS", "6G")
diff --git a/mapillary_tools/exif_read.py b/mapillary_tools/exif_read.py
@@ -110,8 +110,6 @@ class ExifRead:
     EXIF class for reading exif from an image
     """
 
-    path_or_stream: T.Union[Path, T.BinaryIO]
-
     def __init__(
         self, path_or_stream: T.Union[Path, T.BinaryIO], details: bool = False
     ) -> None:
@@ -125,7 +123,6 @@ def __init__(
             self.tags = exifread.process_file(
                 path_or_stream, details=details, debug=True
             )
-        self.path_or_stream = path_or_stream
 
     def extract_altitude(self) -> T.Optional[float]:
         """
@@ -285,6 +282,22 @@ def extract_model(self) -> T.Optional[str]:
         """
         return self._extract_alternative_fields(["EXIF LensModel", "Image Model"], str)
 
+    def extract_width(self) -> T.Optional[int]:
+        """
+        Extract image width in pixels
+        """
+        return self._extract_alternative_fields(
+            ["Image ImageWidth", "EXIF ExifImageWidth"], int
+        )
+
+    def extract_height(self) -> T.Optional[int]:
+        """
+        Extract image height in pixels
+        """
+        return self._extract_alternative_fields(
+            ["Image ImageLength", "EXIF ExifImageLength"], int
+        )
+
     def extract_orientation(self) -> int:
         """
         Extract image orientation
diff --git a/mapillary_tools/geotag/geotag_from_exif.py b/mapillary_tools/geotag/geotag_from_exif.py
@@ -75,6 +75,8 @@ def to_description(self) -> T.List[types.ImageMetadataOrError]:
                 alt=exif.extract_altitude(),
                 angle=exif.extract_direction(),
                 time=geo.as_unix_time(timestamp),
+                width=exif.extract_width(),
+                height=exif.extract_height(),
                 MAPOrientation=exif.extract_orientation(),
                 MAPDeviceMake=exif.extract_make(),
                 MAPDeviceModel=exif.extract_model(),
diff --git a/mapillary_tools/geotag/geotag_from_gpx.py b/mapillary_tools/geotag/geotag_from_gpx.py
@@ -31,13 +31,29 @@ def __init__(
         self.use_gpx_start_time = use_gpx_start_time
         self.use_image_start_time = use_image_start_time
         self.offset_time = offset_time
+        self._exifs: T.Dict[Path, ExifRead] = {}
 
     def read_image_time(self, image_path: Path) -> T.Optional[float]:
-        image_time = ExifRead(image_path).extract_capture_time()
+        exif = self._exifs.get(image_path)
+        if exif is None:
+            exif = ExifRead(image_path)
+            self._exifs[image_path] = exif
+        image_time = exif.extract_capture_time()
         if image_time is None:
             return None
         return geo.as_unix_time(image_time)
 
+    def read_image_size(
+        self, image_path: Path
+    ) -> T.Tuple[T.Optional[int], T.Optional[int]]:
+        exif = self._exifs.get(image_path)
+        if exif is None:
+            exif = ExifRead(image_path)
+            self._exifs[image_path] = exif
+        width = exif.extract_width()
+        height = exif.extract_height()
+        return width, height
+
     def to_description(self) -> T.List[types.ImageMetadataOrError]:
         metadatas: T.List[types.ImageMetadataOrError] = []
 
@@ -165,13 +181,16 @@ def to_description(self) -> T.List[types.ImageMetadataOrError]:
                     continue
 
             interpolated = geo.interpolate(sorted_points, image_time)
+            width, height = self.read_image_size(image_path)
             image_metadata = types.ImageMetadata(
                 filename=image_path,
                 lat=interpolated.lat,
                 lon=interpolated.lon,
                 alt=interpolated.alt,
                 angle=interpolated.angle,
                 time=interpolated.time,
+                width=width,
+                height=height,
             )
             metadatas.append(image_metadata)
 
diff --git a/mapillary_tools/process_sequence_properties.py b/mapillary_tools/process_sequence_properties.py
@@ -2,10 +2,13 @@
 import itertools
 import math
 import typing as T
+import logging
 
 from . import constants, geo, types
 from .exceptions import MapillaryDuplicationError, MapillaryBadParameterError
 
+LOG = logging.getLogger(__name__)
+
 
 Point = T.TypeVar("Point", bound=geo.Point)
 PointSequence = T.List[Point]
@@ -95,26 +98,63 @@ def cut_sequence(
     sequence: T.List[types.ImageMetadata],
     max_images: int,
     max_sequence_filesize: int,
+    max_sequence_pixels: int,
 ) -> T.List[T.List[types.ImageMetadata]]:
     """
     Cut a sequence into multiple sequences by max_images or max filesize
     """
     sequences: T.List[T.List[types.ImageMetadata]] = []
     last_sequence_file_size = 0
+    last_sequence_pixels = 0
+
     for image in sequence:
+        # decent default values if width/height not available
+        width = 1024 if image.width is None else image.width
+        height = 1024 if image.height is None else image.height
+
         filesize = os.path.getsize(image.filename)
-        if len(sequences) == 0 or (
-            sequences[-1]
-            and (
-                max_images < len(sequences[-1])
-                or max_sequence_filesize < last_sequence_file_size + filesize
-            )
-        ):
+
+        if len(sequences) == 0:
+            start_new_sequence = True
+        else:
+            if sequences[-1]:
+                if max_images < len(sequences[-1]):
+                    LOG.debug(
+                        "Cut the sequence because the current sequence (%s) reaches the max number of images (%s)",
+                        len(sequences[-1]),
+                        max_images,
+                    )
+                    start_new_sequence = True
+                elif max_sequence_filesize < last_sequence_file_size + filesize:
+                    LOG.debug(
+                        "Cut the sequence because the current sequence (%s) reaches the max filesize (%s)",
+                        last_sequence_file_size + filesize,
+                        max_sequence_filesize,
+                    )
+                    start_new_sequence = True
+                elif max_sequence_pixels < last_sequence_pixels + width * height:
+                    LOG.debug(
+                        "Cut the sequence because the current sequence (%s) reaches the max pixels (%s)",
+                        last_sequence_pixels + width * height,
+                        max_sequence_pixels,
+                    )
+                    start_new_sequence = True
+                else:
+                    start_new_sequence = False
+            else:
+                start_new_sequence = False
+
+        if start_new_sequence:
             sequences.append([])
             last_sequence_file_size = 0
+            last_sequence_pixels = 0
+
         sequences[-1].append(image)
         last_sequence_file_size += filesize
+        last_sequence_pixels += width * height
+
     assert sum(len(s) for s in sequences) == len(sequence)
+
     return sequences
 
 
@@ -192,8 +232,21 @@ def _parse_filesize_in_bytes(filesize_str: str) -> int:
         return int(filesize_str)
 
 
+def _parse_pixels(pixels_str: str) -> int:
+    pixels_str = pixels_str.strip().upper()
+
+    if pixels_str.endswith("K"):
+        return int(pixels_str[:-1]) * 1000
+    elif pixels_str.endswith("M"):
+        return int(pixels_str[:-1]) * 1000 * 1000
+    elif pixels_str.endswith("G"):
+        return int(pixels_str[:-1]) * 1000 * 1000 * 1000
+    else:
+        return int(pixels_str)
+
+
 def process_sequence_properties(
-    metadatas: T.List[types.MetadataOrError],
+    metadatas: T.Sequence[types.MetadataOrError],
     cutoff_distance=constants.CUTOFF_DISTANCE,
     cutoff_time=constants.CUTOFF_TIME,
     interpolate_directions=False,
@@ -209,6 +262,13 @@ def process_sequence_properties(
             f"Expect the envvar {constants._ENV_PREFIX}MAX_SEQUENCE_FILESIZE to be a valid filesize that ends with B, K, M, or G, but got {constants.MAX_SEQUENCE_FILESIZE}"
         )
 
+    try:
+        max_sequence_pixels = _parse_pixels(constants.MAX_SEQUENCE_PIXELS)
+    except ValueError:
+        raise MapillaryBadParameterError(
+            f"Expect the envvar {constants._ENV_PREFIX}MAX_SEQUENCE_PIXELS to be a valid number of pixels that ends with K, M, or G, but got {constants.MAX_SEQUENCE_PIXELS}"
+        )
+
     error_metadatas: T.List[types.ErrorMetadata] = []
     image_metadatas: T.List[types.ImageMetadata] = []
     video_metadatas: T.List[types.VideoMetadata] = []
@@ -262,6 +322,7 @@ def process_sequence_properties(
             dedups,
             constants.MAX_SEQUENCE_LENGTH,
             max_sequence_filesize_in_bytes,
+            max_sequence_pixels,
         )
 
         for c in cut:
diff --git a/mapillary_tools/types.py b/mapillary_tools/types.py
@@ -44,6 +44,8 @@ class FileType(enum.Enum):
 class ImageMetadata(geo.Point):
     filename: Path
     # filetype is always FileType.IMAGE
+    width: T.Optional[int]
+    height: T.Optional[int]
     MAPSequenceUUID: T.Optional[str] = None
     MAPDeviceMake: T.Optional[str] = None
     MAPDeviceModel: T.Optional[str] = None
@@ -518,6 +520,8 @@ def _from_image_desc(desc) -> ImageMetadata:
         alt=desc.get("MAPAltitude"),
         time=geo.as_unix_time(map_capture_time_to_datetime(desc["MAPCaptureTime"])),
         angle=desc.get("MAPCompassHeading", {}).get("TrueHeading"),
+        width=None,
+        height=None,
         **kwargs,
     )
 
diff --git a/tests/cli/exif_read.py b/tests/cli/exif_read.py
@@ -13,6 +13,7 @@ def main():
     parsed_args = parser.parse_args()
     for image_path in utils.find_images([Path(p) for p in parsed_args.path]):
         exif = ExifRead(image_path, details=True)
+        pprint.pprint(exif.tags)
         pprint.pprint(
             {
                 "filename": image_path,
@@ -24,6 +25,8 @@ def main():
                 "lon_lat": exif.extract_lon_lat(),
                 "make": exif.extract_make(),
                 "model": exif.extract_model(),
+                "width": exif.extract_width(),
+                "height": exif.extract_height(),
             }
         )
 
diff --git a/tests/integration/test_process.py b/tests/integration/test_process.py
@@ -163,12 +163,12 @@ def test_process_images_with_overwrite_all_EXIF_tags(setup_data: py.path.local):
     assert x.returncode == 0, x.stderr
     expected_descs = [
         {
-            **_DEFAULT_EXPECTED_DESCS["DSC00001.JPG"],
+            **_DEFAULT_EXPECTED_DESCS["DSC00001.JPG"],  # type: ignore
             "filename": str(Path(setup_data, "images", "DSC00001.JPG")),
             "MAPCaptureTime": "2018_06_08_20_24_13_500",
         },
         {
-            **_DEFAULT_EXPECTED_DESCS["DSC00497.JPG"],
+            **_DEFAULT_EXPECTED_DESCS["DSC00497.JPG"],  # type: ignore
             "filename": str(Path(setup_data, "images", "DSC00497.JPG")),
             "MAPCaptureTime": "2018_06_08_20_32_30_500",
         },
diff --git a/tests/unit/test_sequence_processing.py b/tests/unit/test_sequence_processing.py
diff --git a/tests/unit/test_types.py b/tests/unit/test_types.py

Original file line number	Diff line number	Diff line change
`@@ -13,6 +13,7 @@ def main():`
`13`	`13`	`parsed_args = parser.parse_args()`
`14`	`14`	`for image_path in utils.find_images([Path(p) for p in parsed_args.path]):`
`15`	`15`	`exif = ExifRead(image_path, details=True)`
	`16`	`+ pprint.pprint(exif.tags)`
`16`	`17`	`pprint.pprint(`
`17`	`18`	`{`
`18`	`19`	`"filename": image_path,`
`@@ -24,6 +25,8 @@ def main():`
`24`	`25`	`"lon_lat": exif.extract_lon_lat(),`
`25`	`26`	`"make": exif.extract_make(),`
`26`	`27`	`"model": exif.extract_model(),`
	`28`	`+ "width": exif.extract_width(),`
	`29`	`+ "height": exif.extract_height(),`
`27`	`30`	`}`
`28`	`31`	`)`
`29`	`32`