137 lines
5.3 KiB
Python
137 lines
5.3 KiB
Python
"""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)
|