11. On-Device Training & Personalisasi

// Core ML: on-device fine-tuning dengan MLUpdateTask (iOS 13+)

import CoreML

final class PersonalizedClassifier {
    private var modelURL: URL
    private var updatedModelURL: URL

    init(baseModelURL: URL, userModelURL: URL) {
        self.modelURL = baseModelURL
        self.updatedModelURL = userModelURL
    }

    // Train ulang model dengan data baru dari user
    func updateModel(with newSamples: [(image: UIImage, label: String)]) async throws {
        var batchProvider: [MLFeatureProvider] = []

        for (image, label) in newSamples {
            guard let cgImage = image.cgImage else { continue }

            // Resize gambar ke ukuran yang diharapkan model (misal 224x224)
            let resizedImage = resizeCGImage(cgImage, to: CGSize(width: 224, height: 224))

            guard let pixelBuffer = resizedImage?.toPixelBuffer() else { continue }

            let imageFeature = try MLFeatureValue(pixelBuffer: pixelBuffer)
            let labelFeature = MLFeatureValue(string: label)

            let provider = try MLDictionaryFeatureProvider(dictionary: [
                "image": imageFeature,
                "label": labelFeature
            ])
            batchProvider.append(provider)
        }

        let trainingData = MLArrayBatchProvider(array: batchProvider)

        let config = MLModelConfiguration()
        config.computeUnits = .cpuAndNeuralEngine

        // Output ke file terpisah — jangan overwrite model original
        let updateTask = try MLUpdateTask(
            forModelAt: modelURL,
            trainingData: trainingData,
            configuration: config,
            progressHandlers: MLUpdateProgressHandlers(
                forEvents: [.trainingBegin, .epochEnd],
                progressHandler: { context in
                    if case .epochEnd = context.event {
                        print("Epoch \(context.metrics[.epochIndex] ?? 0): loss = \(context.metrics[.lossValue] ?? 0)")
                    }
                },
                completionHandler: { [weak self] context in
                    if context.task.error == nil {
                        do {
                            try context.model.write(to: self!.updatedModelURL)
                            print("Model updated successfully")
                        } catch {
                            print("Failed to save updated model: \(error)")
                        }
                    }
                }
            )
        )

        updateTask.resume()
    }

    private func resizeCGImage(_ image: CGImage, to size: CGSize) -> CGImage? {
        let context = CGContext(
            data: nil,
            width: Int(size.width),
            height: Int(size.height),
            bitsPerComponent: image.bitsPerComponent,
            bytesPerRow: 0,
            space: image.colorSpace ?? CGColorSpaceCreateDeviceRGB(),
            bitmapInfo: image.bitmapInfo.rawValue
        )
        context?.draw(image, in: CGRect(origin: .zero, size: size))
        return context?.makeImage()
    }
}