Refactor photo metadata workflow

tools/photo_grouping.py

@@ -0,0 +1,136 @@
+"""Burst/HDR group detection for photo metadata cleanup."""
+
+from __future__ import annotations
+
+from collections import defaultdict
+from datetime import datetime, timedelta
+from typing import TYPE_CHECKING
+
+from tools.filenames import naive_wall_time
+
+if TYPE_CHECKING:
+    from photo_metadata import MediaRecord
+
+
+def photo_sort_key(record: MediaRecord) -> tuple:
+    assert record.adjusted_time is not None
+    return (
+        naive_wall_time(record.adjusted_time),
+        record.subsec if record.subsec is not None else -1,
+        record.sequence if record.sequence is not None else -1,
+        record.path.name.lower(),
+    )
+
+
+def group_sort_key(record: MediaRecord) -> tuple:
+    timestamp = naive_wall_time(record.adjusted_time) if record.adjusted_time else datetime.max
+    return (
+        timestamp,
+        record.subsec if record.subsec is not None else -1,
+        record.sequence if record.sequence is not None else -1,
+        str(record.path).lower(),
+    )
+
+
+def detect_photo_groups(records: list[MediaRecord], min_size: int = 2) -> list[list[MediaRecord]]:
+    by_camera: dict[str, list[MediaRecord]] = defaultdict(list)
+    for record in records:
+        if record.is_image and record.adjusted_time is not None:
+            by_camera[record.camera_key].append(record)
+    return [
+        group
+        for camera_records in by_camera.values()
+        for group in detect_groups_for_camera(sorted(camera_records, key=photo_sort_key), min_size)
+    ]
+
+
+def detect_groups_for_camera(records: list[MediaRecord], min_size: int) -> list[list[MediaRecord]]:
+    by_second: dict[datetime, list[MediaRecord]] = defaultdict(list)
+    for record in records:
+        assert record.adjusted_time is not None
+        by_second[naive_wall_time(record.adjusted_time).replace(microsecond=0)].append(record)
+
+    seconds = sorted(by_second)
+    spans: list[list[datetime]] = []
+    for second in seconds:
+        if not spans or second != spans[-1][-1] + timedelta(seconds=1):
+            spans.append([])
+        spans[-1].append(second)
+
+    groups: list[list[MediaRecord]] = []
+    for span in spans:
+        if max(len(by_second[second]) for second in span) < 2:
+            continue
+        span_records = [record for second in span for record in by_second[second]]
+        if len(span_records) >= 2 and all(record.sequence is not None for record in span_records):
+            groups.extend(sequence_groups(span_records, min_size))
+        else:
+            groups.extend(timestamp_groups(span, by_second, min_size))
+    return [group for group in groups if len(group) >= min_size]
+
+
+def sequence_groups(records: list[MediaRecord], min_size: int) -> list[list[MediaRecord]]:
+    groups: list[list[MediaRecord]] = []
+    for record in sorted(records, key=photo_sort_key):
+        if not groups or record.sequence <= groups[-1][-1].sequence:
+            groups.append([])
+        groups[-1].append(record)
+    return [group for group in groups if len(group) >= min_size]
+
+
+def timestamp_groups(
+    span: list[datetime], by_second: dict[datetime, list[MediaRecord]], min_size: int
+) -> list[list[MediaRecord]]:
+    multi_indices = [index for index, second in enumerate(span) if len(by_second[second]) >= 2]
+    if not multi_indices:
+        return []
+
+    clusters: list[list[int]] = [[multi_indices[0]]]
+    for index in multi_indices[1:]:
+        if index - clusters[-1][-1] > 2:
+            clusters.append([])
+        clusters[-1].append(index)
+
+    groups: list[list[MediaRecord]] = []
+    used_seconds: set[datetime] = set()
+    for cluster in clusters:
+        start, end = cluster[0], cluster[-1]
+        if start and len(by_second[span[start - 1]]) == 1:
+            start -= 1
+        if end + 1 < len(span) and len(by_second[span[end + 1]]) == 1:
+            end += 1
+        seconds = [second for second in span[start : end + 1] if second not in used_seconds]
+        used_seconds.update(seconds)
+        records = [record for second in seconds for record in sorted(by_second[second], key=photo_sort_key)]
+        if len(records) >= min_size:
+            groups.append(records)
+    return groups
+
+
+def infer_group_subseconds(groups: list[list[MediaRecord]]) -> None:
+    for group in groups:
+        if any(record.subsec is not None for record in group):
+            continue
+        by_second: dict[datetime, list[MediaRecord]] = defaultdict(list)
+        for record in group:
+            if record.is_image and record.adjusted_time is not None:
+                by_second[naive_wall_time(record.adjusted_time).replace(microsecond=0)].append(record)
+
+        seconds = sorted(by_second)
+        counts = {second: len(records) for second, records in by_second.items()}
+        for index, second in enumerate(seconds):
+            records = sorted(by_second[second], key=photo_sort_key)
+            count = len(records)
+            if count <= 1:
+                continue
+            if len(seconds) == 2:
+                fps = max(counts.values())
+            elif index == 0:
+                fps = max(count, counts[seconds[1]])
+            elif index == len(seconds) - 1:
+                fps = max(count, counts[seconds[-2]])
+            else:
+                fps = count
+            start_slot = fps - count if index == 0 and fps > count else 0
+            for item_index, record in enumerate(records):
+                record.inferred_subsec = int((start_slot + item_index) * 1000 / fps)