Add similarity grouping pipeline
This commit is contained in:
@@ -0,0 +1,351 @@
|
||||
import Foundation
|
||||
import Vision
|
||||
|
||||
/// 走査からグループ生成までの一連の処理をまとめる。
|
||||
///
|
||||
/// 安いフィルタから順にかけて候補を絞る:
|
||||
/// 1. ファイルサイズが同じものだけ SHA-256 → 完全同一を確定
|
||||
/// 2. dHash（知覚ハッシュ）で「見た目が近そうな」ペアを抽出
|
||||
/// 3. そのペアに登場する画像だけ Vision の feature print で精査
|
||||
/// 4. Union-Find で連結成分をグループ化
|
||||
enum SimilarityGrouper {
|
||||
private struct Pair: Sendable {
|
||||
let a: Int
|
||||
let b: Int
|
||||
}
|
||||
|
||||
private struct HashEntry: Sendable {
|
||||
let index: Int
|
||||
let hash: UInt64
|
||||
}
|
||||
|
||||
/// feature print を計算できなかったときの保険。dHash がここまで一致していれば同一とみなす。
|
||||
private static let fallbackHammingThreshold = 4
|
||||
|
||||
static func run(
|
||||
root: URL,
|
||||
settings: ScanSettings,
|
||||
cache: FingerprintCache,
|
||||
progress: @escaping @Sendable (ScanProgress) -> Void
|
||||
) async throws -> [SimilarityGroup] {
|
||||
progress(ScanProgress(phase: .scanning))
|
||||
let files = try PhotoScanner.scan(root: root, includesSubfolders: settings.includesSubfolders)
|
||||
try Task.checkCancellation()
|
||||
guard files.count > 1 else { return [] }
|
||||
|
||||
// --- Stage 1: 完全同一をまとめ、以降は各セットの代表1枚だけを扱う ---
|
||||
let leaderOf = try await exactLeaders(files: files, progress: progress)
|
||||
let leaderIndices = files.indices.filter { leaderOf[$0] == $0 }
|
||||
let leaderFiles = leaderIndices.map { files[$0] }
|
||||
|
||||
// --- Stage 2: dHash ---
|
||||
let basicsByURL = try await loadBasics(for: leaderFiles, cache: cache, progress: progress)
|
||||
|
||||
// --- Stage 3: 候補ペア → feature print ---
|
||||
let entries = leaderIndices.compactMap { index -> HashEntry? in
|
||||
guard let basics = basicsByURL[files[index].url] else { return nil }
|
||||
return HashEntry(index: index, hash: basics.dHash)
|
||||
}
|
||||
progress(ScanProgress(phase: .comparing))
|
||||
let pairs = try await candidatePairs(
|
||||
entries,
|
||||
threshold: settings.sensitivity.hammingThreshold
|
||||
)
|
||||
|
||||
var involved = Set<Int>()
|
||||
for pair in pairs {
|
||||
involved.insert(pair.a)
|
||||
involved.insert(pair.b)
|
||||
}
|
||||
let involvedFiles = involved.sorted().map { files[$0] }
|
||||
let printsByURL = try await loadFeaturePrints(
|
||||
for: involvedFiles, cache: cache, progress: progress
|
||||
)
|
||||
|
||||
// --- Stage 4: グループ化 ---
|
||||
progress(ScanProgress(phase: .grouping))
|
||||
var unionFind = UnionFind(count: files.count)
|
||||
for index in files.indices where leaderOf[index] != index {
|
||||
unionFind.union(index, leaderOf[index])
|
||||
}
|
||||
|
||||
let threshold = settings.sensitivity.featurePrintThreshold
|
||||
for pair in pairs {
|
||||
let printA = printsByURL[files[pair.a].url]
|
||||
let printB = printsByURL[files[pair.b].url]
|
||||
|
||||
if let distance = FeaturePrinter.distance(printA, printB) {
|
||||
if distance < threshold { unionFind.union(pair.a, pair.b) }
|
||||
} else if let basicsA = basicsByURL[files[pair.a].url],
|
||||
let basicsB = basicsByURL[files[pair.b].url],
|
||||
DHash.hammingDistance(basicsA.dHash, basicsB.dHash) <= fallbackHammingThreshold {
|
||||
unionFind.union(pair.a, pair.b)
|
||||
}
|
||||
}
|
||||
try Task.checkCancellation()
|
||||
|
||||
let groups = buildGroups(
|
||||
components: unionFind.components(),
|
||||
files: files,
|
||||
leaderOf: leaderOf,
|
||||
basicsByURL: basicsByURL,
|
||||
printsByURL: printsByURL
|
||||
)
|
||||
|
||||
await cache.prune(keeping: files)
|
||||
await cache.save()
|
||||
return groups
|
||||
}
|
||||
|
||||
// MARK: - Stage 1
|
||||
|
||||
/// 各インデックスに「完全同一セットの代表インデックス」を割り当てる。単独なら自分自身。
|
||||
private static func exactLeaders(
|
||||
files: [ScannedFile],
|
||||
progress: @escaping @Sendable (ScanProgress) -> Void
|
||||
) async throws -> [Int] {
|
||||
var leaderOf = Array(files.indices)
|
||||
|
||||
// サイズが違えば中身も違う。同サイズのものだけハッシュすればよい。
|
||||
var bySize: [Int64: [Int]] = [:]
|
||||
for (index, file) in files.enumerated() {
|
||||
bySize[file.fileSize, default: []].append(index)
|
||||
}
|
||||
let candidates = bySize.values.filter { $0.count > 1 }.flatMap { $0 }.sorted()
|
||||
guard !candidates.isEmpty else { return leaderOf }
|
||||
|
||||
progress(ScanProgress(phase: .hashing, completed: 0, total: candidates.count))
|
||||
let hashes = try await mapConcurrently(
|
||||
candidates,
|
||||
onProgress: { done in
|
||||
progress(ScanProgress(phase: .hashing, completed: done, total: candidates.count))
|
||||
}
|
||||
) { index -> String? in
|
||||
try? ContentHasher.sha256(of: files[index].url)
|
||||
}
|
||||
|
||||
var byDigest: [String: [Int]] = [:]
|
||||
for (index, hash) in zip(candidates, hashes) {
|
||||
guard let hash else { continue }
|
||||
byDigest["\(files[index].fileSize):\(hash)", default: []].append(index)
|
||||
}
|
||||
for indices in byDigest.values where indices.count > 1 {
|
||||
guard let leader = indices.min(by: { files[$0].url.path < files[$1].url.path }) else {
|
||||
continue
|
||||
}
|
||||
for index in indices { leaderOf[index] = leader }
|
||||
}
|
||||
return leaderOf
|
||||
}
|
||||
|
||||
// MARK: - Stage 2 / 3
|
||||
|
||||
private static func loadBasics(
|
||||
for files: [ScannedFile],
|
||||
cache: FingerprintCache,
|
||||
progress: @escaping @Sendable (ScanProgress) -> Void
|
||||
) async throws -> [URL: ImageBasics] {
|
||||
var result = await cache.basics(for: files)
|
||||
let missing = files.filter { result[$0.url] == nil }
|
||||
guard !missing.isEmpty else { return result }
|
||||
|
||||
let cached = files.count - missing.count
|
||||
progress(ScanProgress(phase: .fingerprinting, completed: cached, total: files.count))
|
||||
let computed = try await mapConcurrently(
|
||||
missing,
|
||||
onProgress: { done in
|
||||
progress(
|
||||
ScanProgress(phase: .fingerprinting, completed: cached + done, total: files.count)
|
||||
)
|
||||
}
|
||||
) { file -> ImageBasics? in
|
||||
DHash.compute(for: file.url)
|
||||
}
|
||||
|
||||
var fresh: [URL: ImageBasics] = [:]
|
||||
for (file, basics) in zip(missing, computed) {
|
||||
guard let basics else { continue }
|
||||
fresh[file.url] = basics
|
||||
result[file.url] = basics
|
||||
}
|
||||
await cache.store(basics: fresh, for: missing)
|
||||
return result
|
||||
}
|
||||
|
||||
private static func loadFeaturePrints(
|
||||
for files: [ScannedFile],
|
||||
cache: FingerprintCache,
|
||||
progress: @escaping @Sendable (ScanProgress) -> Void
|
||||
) async throws -> [URL: FeaturePrintObservation] {
|
||||
guard !files.isEmpty else { return [:] }
|
||||
var result = await cache.featurePrints(for: files)
|
||||
let missing = files.filter { result[$0.url] == nil }
|
||||
guard !missing.isEmpty else { return result }
|
||||
|
||||
let cached = files.count - missing.count
|
||||
progress(ScanProgress(phase: .comparing, completed: cached, total: files.count))
|
||||
let computed = try await mapConcurrently(
|
||||
missing,
|
||||
onProgress: { done in
|
||||
progress(
|
||||
ScanProgress(phase: .comparing, completed: cached + done, total: files.count)
|
||||
)
|
||||
}
|
||||
) { file -> FeaturePrintObservation? in
|
||||
try? await FeaturePrinter.compute(for: file.url)
|
||||
}
|
||||
|
||||
var fresh: [URL: FeaturePrintObservation] = [:]
|
||||
for (file, print) in zip(missing, computed) {
|
||||
guard let print else { continue }
|
||||
fresh[file.url] = print
|
||||
result[file.url] = print
|
||||
}
|
||||
await cache.store(featurePrints: fresh, for: missing)
|
||||
return result
|
||||
}
|
||||
|
||||
/// dHash のハミング距離がしきい値以下のペアを総当たりで探す。
|
||||
///
|
||||
/// 64bit の XOR + popcount なので1ペアあたりの実費はごく小さい。LSH などで
|
||||
/// 絞ると取りこぼしが出るため、区間をコアに分けて総当たりする方針にしている。
|
||||
private static func candidatePairs(
|
||||
_ entries: [HashEntry],
|
||||
threshold: Int
|
||||
) async throws -> [Pair] {
|
||||
guard entries.count > 1 else { return [] }
|
||||
|
||||
let cores = max(2, ProcessInfo.processInfo.activeProcessorCount)
|
||||
let chunkSize = max(64, entries.count / (cores * 4))
|
||||
let starts = stride(from: 0, to: entries.count, by: chunkSize).map { $0 }
|
||||
|
||||
let chunks = try await mapConcurrently(starts) { start -> [Pair] in
|
||||
var found: [Pair] = []
|
||||
let end = min(start + chunkSize, entries.count)
|
||||
for i in start..<end {
|
||||
let hash = entries[i].hash
|
||||
for j in (i + 1)..<entries.count
|
||||
where DHash.hammingDistance(hash, entries[j].hash) <= threshold {
|
||||
found.append(Pair(a: entries[i].index, b: entries[j].index))
|
||||
}
|
||||
}
|
||||
return found
|
||||
}
|
||||
return chunks.flatMap { $0 }
|
||||
}
|
||||
|
||||
// MARK: - Stage 4
|
||||
|
||||
/// グループの中に「代表とは別だが、互いにバイト単位で同じ」ファイルの組があれば印を付ける。
|
||||
///
|
||||
/// 代表と同一のものは相手が「完全同一」と表示されるので印は要らない。ここで拾うのは
|
||||
/// たとえば縮小版が2つあってその2つだけが同一、といったケース。
|
||||
private static func exactSetTags(
|
||||
in ranked: [Int],
|
||||
leaderOf: [Int],
|
||||
excluding representativeLeader: Int
|
||||
) -> [Int: String] {
|
||||
var membersPerLeader: [Int: Int] = [:]
|
||||
for index in ranked where leaderOf[index] != representativeLeader {
|
||||
membersPerLeader[leaderOf[index], default: 0] += 1
|
||||
}
|
||||
|
||||
let labels = Array("ABCDEFGHIJKLMNOPQRSTUVWXYZ")
|
||||
var tags: [Int: String] = [:]
|
||||
// ranked の順に印を振るので、同じ結果からは常に同じ印になる。
|
||||
for index in ranked {
|
||||
let leader = leaderOf[index]
|
||||
guard membersPerLeader[leader, default: 0] > 1, tags[leader] == nil else { continue }
|
||||
tags[leader] = tags.count < labels.count
|
||||
? String(labels[tags.count])
|
||||
: "\(tags.count + 1)"
|
||||
}
|
||||
return tags
|
||||
}
|
||||
|
||||
private static func buildGroups(
|
||||
components: [[Int]],
|
||||
files: [ScannedFile],
|
||||
leaderOf: [Int],
|
||||
basicsByURL: [URL: ImageBasics],
|
||||
printsByURL: [URL: FeaturePrintObservation]
|
||||
) -> [SimilarityGroup] {
|
||||
func makeAsset(_ index: Int) -> PhotoAsset {
|
||||
let file = files[index]
|
||||
// 完全同一なら中身は同じなので、代表の解像度・撮影日をそのまま使える。
|
||||
let basics = basicsByURL[files[leaderOf[index]].url]
|
||||
return PhotoAsset(
|
||||
url: file.url,
|
||||
fileSize: file.fileSize,
|
||||
modificationDate: file.modificationDate,
|
||||
captureDate: basics?.captureDate ?? file.creationDate,
|
||||
pixelWidth: basics?.pixelWidth ?? 0,
|
||||
pixelHeight: basics?.pixelHeight ?? 0
|
||||
)
|
||||
}
|
||||
|
||||
var groups: [SimilarityGroup] = []
|
||||
for component in components where component.count > 1 {
|
||||
let assets = Dictionary(uniqueKeysWithValues: component.map { ($0, makeAsset($0)) })
|
||||
|
||||
// 代表は「解像度が最大 → ファイルサイズが最大 → 撮影が最古」で選ぶ。
|
||||
let ranked = component.sorted { lhs, rhs in
|
||||
guard let a = assets[lhs], let b = assets[rhs] else { return lhs < rhs }
|
||||
if a.pixelCount != b.pixelCount { return a.pixelCount > b.pixelCount }
|
||||
if a.fileSize != b.fileSize { return a.fileSize > b.fileSize }
|
||||
let dateA = a.captureDate ?? .distantFuture
|
||||
let dateB = b.captureDate ?? .distantFuture
|
||||
if dateA != dateB { return dateA < dateB }
|
||||
return a.url.path < b.url.path
|
||||
}
|
||||
guard let representativeIndex = ranked.first,
|
||||
let representative = assets[representativeIndex]
|
||||
else { continue }
|
||||
let representativeLeader = leaderOf[representativeIndex]
|
||||
let exactSetTags = exactSetTags(
|
||||
in: ranked, leaderOf: leaderOf, excluding: representativeLeader
|
||||
)
|
||||
|
||||
let members = component.compactMap { index -> SimilarityGroup.Member? in
|
||||
guard let asset = assets[index] else { return nil }
|
||||
let relation: SimilarityGroup.Relation
|
||||
if leaderOf[index] == representativeLeader {
|
||||
relation = .exact
|
||||
} else {
|
||||
relation = .similar(
|
||||
distance: FeaturePrinter.distance(
|
||||
printsByURL[files[leaderOf[index]].url],
|
||||
printsByURL[files[representativeLeader].url]
|
||||
)
|
||||
)
|
||||
}
|
||||
return SimilarityGroup.Member(
|
||||
asset: asset,
|
||||
relation: relation,
|
||||
exactSetTag: exactSetTags[leaderOf[index]]
|
||||
)
|
||||
}
|
||||
|
||||
let ordered = members.sorted { lhs, rhs in
|
||||
if lhs.id == representative.url { return true }
|
||||
if rhs.id == representative.url { return false }
|
||||
if lhs.relation.sortKey != rhs.relation.sortKey {
|
||||
return lhs.relation.sortKey < rhs.relation.sortKey
|
||||
}
|
||||
return lhs.asset.url.path < rhs.asset.url.path
|
||||
}
|
||||
|
||||
groups.append(
|
||||
SimilarityGroup(members: ordered, representativeID: representative.url)
|
||||
)
|
||||
}
|
||||
|
||||
// 減らせる容量が大きいグループから見せる。
|
||||
return groups.sorted { lhs, rhs in
|
||||
if lhs.reclaimableBytes != rhs.reclaimableBytes {
|
||||
return lhs.reclaimableBytes > rhs.reclaimableBytes
|
||||
}
|
||||
return lhs.representativeID.path < rhs.representativeID.path
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user