From a6cc8121cee0fc2faf4b416087b89e1345b00d7c Mon Sep 17 00:00:00 2001 From: yoshikazu Date: Mon, 27 Jul 2026 15:40:14 +0000 Subject: [PATCH] Add similarity grouping pipeline --- Sources/Photochk/Core/SimilarityGrouper.swift | 351 ++++++++++++++++++ 1 file changed, 351 insertions(+) create mode 100644 Sources/Photochk/Core/SimilarityGrouper.swift diff --git a/Sources/Photochk/Core/SimilarityGrouper.swift b/Sources/Photochk/Core/SimilarityGrouper.swift new file mode 100644 index 0000000..f22a725 --- /dev/null +++ b/Sources/Photochk/Core/SimilarityGrouper.swift @@ -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() + 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.. [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 + } + } +}