Clean/plugins/classify_photo/ios/Classes/ClassifyPhoto.swift

import Photos
import Vision

class ClassifyPhoto {

    struct ClassifiedPhotos {
        var screenshots: [PHAsset] = []
        var locations: [String: [PHAsset]] = [:] // 按地点分组
        var people: [String: [PHAsset]] = [:]     // 按人物分组
        var similarPhotos: [[PHAsset]] = [] // 存储相似照片组
    }

    func classifyPhotos(
            assets: PHFetchResult<PHAsset>,
            progressHandler: @escaping (String, Float) -> Void,
            completion: @escaping (ClassifiedPhotos) -> Void
        ) {
            // 在后台队列处理
            DispatchQueue.global(qos: .userInitiated).async {
                var result = ClassifiedPhotos()
                let group = DispatchGroup()

                // 开始处理
                DispatchQueue.main.async {
                    progressHandler("正在加载照片...", 0.0)
                }

                // 1. 检测截图 (占总进度的 20%)
                group.enter()
                self.fetchScreenshots(from: assets) { screenshots in
                    result.screenshots = screenshots
                    DispatchQueue.main.async {
                        progressHandler("正在检测截图...", 0.2)
                    }
                    group.leave()
                }

                // 2. 检测相似照片 (占总进度的 80%)
                group.enter()
                self.detectSimilarPhotos(
                    assets: assets,
                    progressHandler: { stage, progress in
                        // 将相似照片检测的进度映射到 20%-100% 的范围
                        let mappedProgress = 0.2 + (progress * 0.6)
                        DispatchQueue.main.async {
                            progressHandler(stage, mappedProgress)
                        }
                    }
                ) { similarPhotos in
                    result.similarPhotos = similarPhotos
                    group.leave()
                }
                
                // 3. 按地点分类 (占总进度的 20%)
                group.enter()
                self.classifyByLocation(assets: assets) { locationGroups in
                    result.locations = locationGroups
                    DispatchQueue.main.async {
                        progressHandler("正在按地点分类...", 0.8)
                    }
                    group.leave()
                }

                // 4. 按人物分类 (占总进度的 20%)
                group.enter()
                self.classifyByPeople(assets: assets) { peopleGroups in
                    result.people = peopleGroups
                    DispatchQueue.main.async {
                        progressHandler("正在按人物分类...", 1.0)
                    }
                    group.leave()
                }

                // 等待所有处理完成
                group.notify(queue: .main) {
                    progressHandler("分类完成", 1.0)
                    completion(result)
                }
            }
        }

    private func detectSimilarPhotos(
            assets: PHFetchResult<PHAsset>,
            progressHandler: @escaping (String, Float) -> Void,
            completion: @escaping ([[PHAsset]]) -> Void
        ) {
            var similarGroups: [[PHAsset]] = []
            let group = DispatchGroup()
            
            if #available(iOS 13.0, *) {
                var imageFeatures: [(asset: PHAsset, feature: VNFeaturePrintObservation)] = []
                
                // 创建处理队列
                let processingQueue = DispatchQueue(label: "com.app.similarPhotos", qos: .userInitiated)
                let semaphore = DispatchSemaphore(value: 5)
                
                // 1. 提取所有图片的特征
                let totalAssets = assets.count
                var processedAssets = 0
                
                progressHandler("正在加载照片...", 0.0)
                
                for i in 0..<assets.count {
                    let asset = assets[i]
                    group.enter()
                    semaphore.wait()
                    
                    let options = PHImageRequestOptions()
                    options.deliveryMode = .highQualityFormat
                    options.isSynchronous = false
                    options.resizeMode = .exact
                    
                    PHImageManager.default().requestImage(
                        for: asset,
                        targetSize: CGSize(width: 448, height: 448),
                        contentMode: .aspectFit,
                        options: options
                    ) { image, _ in
                        defer {
                            semaphore.signal()
                        }
                        
                        guard let image = image,
                              let cgImage = image.cgImage else {
                            group.leave()
                            return
                        }
                        
                        processingQueue.async {
                            do {
                                let requestHandler = VNImageRequestHandler(cgImage: cgImage, options: [:])
                                let request = VNGenerateImageFeaturePrintRequest()
                                try requestHandler.perform([request])
                                
                                if let result = request.results?.first as? VNFeaturePrintObservation {
                                    imageFeatures.append((asset, result))
                                    
                                    // 更新特征提取进度
                                    processedAssets += 1
                                    let progress = Float(processedAssets) / Float(totalAssets)
                                    progressHandler("正在提取特征...", progress * 0.6)
                                }
                            } catch {
                                print("特征提取失败: \(error)")
                            }
                            group.leave()
                        }
                    }
                }
                
                // 2. 比较特征相似度并分组
                group.notify(queue: processingQueue) {
                    progressHandler("正在比较相似度...", 0.6)
                    
                    // 近似度
                    let similarityThreshold: Float = 0.7
                    var processedComparisons = 0
                    let totalComparisons = (imageFeatures.count * (imageFeatures.count - 1)) / 2
                    var processedIndices = Set<Int>()
                    
                    for i in 0..<imageFeatures.count {
                        if processedIndices.contains(i) { continue }
                        
                        var similarGroup: [PHAsset] = [imageFeatures[i].asset]
                        processedIndices.insert(i)
                        
                        for j in (i + 1)..<imageFeatures.count {
                            if processedIndices.contains(j) { continue }
                            
                            do {
                                var distance: Float = 0
                                try imageFeatures[i].feature.computeDistance(&distance, to: imageFeatures[j].feature)
                                
                                let similarity = 1 - distance
                                if similarity >= similarityThreshold {
                                    similarGroup.append(imageFeatures[j].asset)
                                    processedIndices.insert(j)
                                }
                                
                                // 更新比较进度
                                processedComparisons += 1
                                let compareProgress = Float(processedComparisons) / Float(totalComparisons)
                                progressHandler("正在比较相似度...", 0.6 + compareProgress * 0.4)
                            } catch {
                                print("相似度计算失败: \(error)")
                            }
                        }
                        
                        if similarGroup.count > 1 {
                            similarGroups.append(similarGroup)
                        }
                    }
                    
                    // 按照照片数量降序排序
                    similarGroups.sort { $0.count > $1.count }
                    
                    DispatchQueue.main.async {
                        completion(similarGroups)
                    }
                }
            }
        }

    // 按地点分类
    private func classifyByLocation(assets: PHFetchResult<PHAsset>,
                                  completion: @escaping ([String: [PHAsset]]) -> Void) {
        var locationGroups: [String: [PHAsset]] = [:]
        let group = DispatchGroup()
        let geocodeQueue = DispatchQueue(label: "com.app.geocoding")
        let semaphore = DispatchSemaphore(value: 10) // 限制并发请求数

        assets.enumerateObjects { asset, _, _ in
            if let location = asset.location {
                group.enter()
                semaphore.wait()

                geocodeQueue.async {
                    let geocoder = CLGeocoder()
                    geocoder.reverseGeocodeLocation(location) { placemarks, error in
                        defer {
                            semaphore.signal()
                            group.leave()
                        }

                        if let placemark = placemarks?.first {
                            let locationName = self.formatLocationName(placemark)
                            DispatchQueue.main.async {
                                if locationGroups[locationName] == nil {
                                    locationGroups[locationName] = []
                                }
                                locationGroups[locationName]?.append(asset)
                            }
                        }
                    }
                }
            }
        }

        // 等待所有地理编码完成后回调
        group.notify(queue: .main) {
            completion(locationGroups)
        }
    }

    // 格式化地点名称（只返回城市名）
    private func formatLocationName(_ placemark: CLPlacemark) -> String {
        if let city = placemark.locality {
            return city
        }
        return "其他"
    }

    // 按人物分类
    private func classifyByPeople(assets: PHFetchResult<PHAsset>,
                                completion: @escaping ([String: [PHAsset]]) -> Void) {
        var peopleGroups: [String: [PHAsset]] = [:]
        let group = DispatchGroup()

        // 创建一个数组来存储检测到人脸的照片
        var facesArray: [PHAsset] = []

        // 遍历所有照片
        assets.enumerateObjects { asset, _, _ in
            group.enter()

            // 获取照片的缩略图进行人脸检测
            let options = PHImageRequestOptions()
            options.isSynchronous = false
            options.deliveryMode = .fastFormat

            PHImageManager.default().requestImage(
                for: asset,
                targetSize: CGSize(width: 500, height: 500), // 使用较小的尺寸提高性能
                contentMode: .aspectFit,
                options: options
            ) { image, _ in
                guard let image = image else {
                    group.leave()
                    return
                }

                // 使用 Vision 框架检测人脸
                guard let ciImage = CIImage(image: image) else {
                    group.leave()
                    return
                }

                let request = VNDetectFaceRectanglesRequest()
                let handler = VNImageRequestHandler(ciImage: ciImage)

                do {
                    try handler.perform([request])
                    if let results = request.results, !results.isEmpty {
                        // 检测到人脸，添加到数组
                        DispatchQueue.main.async {
                            facesArray.append(asset)
                        }
                    }
                } catch {
                    print("人脸检测失败: \(error)")
                }

                group.leave()
            }
        }

        // 等待所有检测完成后更新结果
        group.notify(queue: .main) {
            if !facesArray.isEmpty {
                peopleGroups["包含人脸的照片"] = facesArray
            }
            completion(peopleGroups)
        }
    }

    // 识别截图
    private func fetchScreenshots(from assets: PHFetchResult<PHAsset>,
                                completion: @escaping ([PHAsset]) -> Void) {
        var screenshots: [PHAsset] = []

        // 获取系统的截图智能相册
        let screenshotAlbums = PHAssetCollection.fetchAssetCollections(
            with: .smartAlbum,
            subtype: .smartAlbumScreenshots,
            options: nil
        )

        // 从截图相册中获取所有截图
        screenshotAlbums.enumerateObjects { collection, _, _ in
            let fetchOptions = PHFetchOptions()
            let screenshotAssets = PHAsset.fetchAssets(in: collection, options: fetchOptions)

            screenshotAssets.enumerateObjects { asset, _, _ in
                screenshots.append(asset)
            }
        }

        completion(screenshots)
    }
}