637 lines
18 KiB
Swift
637 lines
18 KiB
Swift
import Foundation
|
|
import SwiftUI
|
|
import Vision
|
|
import CryptoKit
|
|
|
|
final class Database: ObservableObject {
|
|
|
|
static let imageCacheMemory = 10_000_000
|
|
|
|
static let imageCacheStorage = 200_000_000
|
|
|
|
private let imageCompressionQuality: CGFloat = 0.3
|
|
|
|
private static var documentDirectory: URL {
|
|
try! FileManager.default.url(
|
|
for: .documentDirectory,
|
|
in: .userDomainMask,
|
|
appropriateFor: nil, create: true)
|
|
}
|
|
|
|
private var fm: FileManager {
|
|
.default
|
|
}
|
|
|
|
private var localDbUrl: URL {
|
|
Database.documentDirectory.appendingPathComponent("db.json")
|
|
}
|
|
|
|
private var localClassifierUrl: URL {
|
|
Database.documentDirectory.appendingPathComponent("classifier.mlmodel")
|
|
}
|
|
|
|
private var imageUploadFolderUrl: URL {
|
|
Database.documentDirectory.appendingPathComponent("uploads")
|
|
}
|
|
|
|
private var serverDbUrl: URL {
|
|
serverUrl.appendingPathComponent("caps.json")
|
|
}
|
|
|
|
private var serverClassifierUrl: URL {
|
|
serverUrl.appendingPathComponent("classifier.mlmodel")
|
|
}
|
|
|
|
private var serverClassifierVersionUrl: URL {
|
|
serverUrl.appendingPathComponent("classifier.version")
|
|
}
|
|
|
|
private let encoder = JSONEncoder()
|
|
private let decoder = JSONDecoder()
|
|
|
|
let serverUrl: URL
|
|
|
|
@AppStorage("authKey")
|
|
private var serverAuthenticationKey: String?
|
|
|
|
var hasServerAuthentication: Bool {
|
|
serverAuthenticationKey != nil
|
|
}
|
|
|
|
@Published
|
|
private(set) var caps: [Int : Cap] {
|
|
didSet { scheduleSave() }
|
|
}
|
|
|
|
var nextCapId: Int {
|
|
(caps.values.max()?.id ?? 0) + 1
|
|
}
|
|
|
|
@AppStorage("changed")
|
|
private var changedCapStorage: String = ""
|
|
|
|
private(set) var changedCaps: Set<Int> {
|
|
get {
|
|
Set(changedCapStorage.components(separatedBy: ",").compactMap(Int.init))
|
|
}
|
|
set {
|
|
changedCapStorage = newValue.map { "\($0)" }.joined(separator: ",")
|
|
}
|
|
}
|
|
|
|
private lazy var imageUploads: [Int: Int] = loadImageUploadCounts()
|
|
|
|
private var uploadTimer: Timer?
|
|
|
|
/// The classifications for all caps from the classifier
|
|
@Published
|
|
var matches = [Int : Float]()
|
|
|
|
@Published
|
|
var image: UIImage? = nil {
|
|
didSet { classifyImage() }
|
|
}
|
|
|
|
private var classifier: Classifier?
|
|
|
|
/**
|
|
The time to wait for changes to be written to disk.
|
|
|
|
This delay is used to prevent file writes for each small update to the caps.
|
|
*/
|
|
private let saveDelay: TimeInterval = 1
|
|
|
|
/**
|
|
The time when a save should occur.
|
|
|
|
No save is necessary if this property is `nil`.
|
|
*/
|
|
private var nextSaveTime: Date?
|
|
|
|
let imageCache: URLCache
|
|
|
|
init(server: URL) {
|
|
self.serverUrl = server
|
|
self.caps = [:]
|
|
|
|
let cacheDirectory = Database.documentDirectory.appendingPathComponent("images")
|
|
self.imageCache = URLCache(
|
|
memoryCapacity: Database.imageCacheMemory,
|
|
diskCapacity: Database.imageCacheStorage,
|
|
directory: cacheDirectory)
|
|
loadCaps()
|
|
|
|
}
|
|
|
|
@Published
|
|
var isUploading = false
|
|
|
|
// MARK: Disk storage
|
|
|
|
private func loadCaps() {
|
|
guard fm.fileExists(atPath: localDbUrl.path) else {
|
|
return
|
|
}
|
|
let data: Data
|
|
do {
|
|
data = try Data(contentsOf: localDbUrl)
|
|
} catch {
|
|
print("Failed to read database file: \(error)")
|
|
return
|
|
}
|
|
do {
|
|
let array = try JSONDecoder().decode([Cap].self, from: data)
|
|
self.caps = array.reduce(into: [:]) { $0[$1.id] = $1 }
|
|
// Prevent immediate save after modifying caps
|
|
nextSaveTime = nil
|
|
} catch {
|
|
print("Failed to decode database file: \(error)")
|
|
return
|
|
}
|
|
}
|
|
|
|
private func scheduleSave() {
|
|
nextSaveTime = Date.now.addingTimeInterval(saveDelay)
|
|
DispatchQueue.main.asyncAfter(deadline: .now() + saveDelay) {
|
|
self.performScheduledSave()
|
|
}
|
|
}
|
|
|
|
private func performScheduledSave() {
|
|
guard let date = nextSaveTime else {
|
|
// No save necessary, or already saved
|
|
return
|
|
}
|
|
guard date < .now else {
|
|
// Save pushed to future
|
|
return
|
|
}
|
|
save()
|
|
nextSaveTime = nil
|
|
}
|
|
|
|
private func save() {
|
|
let data: Data
|
|
do {
|
|
data = try encoder.encode(caps.values.sorted())
|
|
} catch {
|
|
print("Failed to encode database: \(error)")
|
|
return
|
|
}
|
|
do {
|
|
try data.write(to: localDbUrl)
|
|
} catch {
|
|
print("Failed to save database: \(error)")
|
|
}
|
|
print("Database saved")
|
|
}
|
|
|
|
private func ensureFolderExistence(_ url: URL) -> Bool {
|
|
guard !fm.fileExists(atPath: url.path) else {
|
|
return true
|
|
}
|
|
do {
|
|
try fm.createDirectory(at: url, withIntermediateDirectories: true)
|
|
return true
|
|
} catch {
|
|
log("Failed to create folder \(url.path): \(error)")
|
|
return false
|
|
}
|
|
}
|
|
|
|
// MARK: Downloads
|
|
|
|
@discardableResult
|
|
func downloadCaps() async -> Bool {
|
|
print("Downloading cap data")
|
|
let data: Data
|
|
let response: URLResponse
|
|
do {
|
|
(data, response) = try await URLSession.shared.data(from: serverDbUrl)
|
|
} catch {
|
|
print("Failed to download classifier version: \(error)")
|
|
return false
|
|
}
|
|
guard (response as? HTTPURLResponse)?.statusCode == 200 else {
|
|
return false
|
|
}
|
|
|
|
let capData: [CapData]
|
|
do {
|
|
capData = try decoder.decode([CapData].self, from: data)
|
|
} catch {
|
|
print("Failed to decode server database: \(error)")
|
|
return false
|
|
}
|
|
var inserts = 0
|
|
var updates = 0
|
|
for cap in capData {
|
|
guard var oldCap = caps[cap.id] else {
|
|
caps[cap.id] = Cap(data: cap)
|
|
inserts += 1
|
|
continue
|
|
}
|
|
guard oldCap != cap else {
|
|
continue
|
|
}
|
|
if changedCaps.contains(oldCap.id) {
|
|
#warning("Merge changed caps with server updates")
|
|
} else {
|
|
oldCap.update(with: cap)
|
|
caps[cap.id] = oldCap
|
|
updates += 1
|
|
}
|
|
}
|
|
print("Updated database from server (\(inserts) added, \(updates) updated)")
|
|
return true
|
|
}
|
|
|
|
@discardableResult
|
|
func serverHasNewClassifier() async -> Bool {
|
|
let data: Data
|
|
let response: URLResponse
|
|
do {
|
|
(data, response) = try await URLSession.shared.data(from: serverClassifierVersionUrl)
|
|
} catch {
|
|
print("Failed to download classifier version: \(error)")
|
|
return false
|
|
}
|
|
guard (response as? HTTPURLResponse)?.statusCode == 200 else {
|
|
return false
|
|
}
|
|
|
|
guard let string = String(data: data, encoding: .utf8) else {
|
|
log("Classifier version is invalid data (not a string)")
|
|
return false
|
|
}
|
|
guard let serverVersion = Int(string.trimmingCharacters(in: .whitespacesAndNewlines)) else {
|
|
log("Classifier version has an invalid value '\(string)'")
|
|
return false
|
|
}
|
|
DispatchQueue.main.async {
|
|
self.serverClassifierVersion = serverVersion
|
|
}
|
|
guard serverVersion > self.classifierVersion else {
|
|
print("No new classifier available")
|
|
return false
|
|
}
|
|
print("New classifier \(serverVersion) available")
|
|
return true
|
|
}
|
|
|
|
@discardableResult
|
|
func downloadClassifier() async -> Bool {
|
|
print("Downloading classifier")
|
|
let tempUrl: URL
|
|
let response: URLResponse
|
|
do {
|
|
(tempUrl, response) = try await URLSession.shared.download(from: serverClassifierUrl)
|
|
} catch {
|
|
print("Failed to download classifier version: \(error)")
|
|
return false
|
|
}
|
|
guard (response as? HTTPURLResponse)?.statusCode == 200 else {
|
|
return false
|
|
}
|
|
do {
|
|
let url = self.localClassifierUrl
|
|
if fm.fileExists(atPath: url.path) {
|
|
try self.fm.removeItem(at: url)
|
|
}
|
|
try self.fm.moveItem(at: tempUrl, to: url)
|
|
} catch {
|
|
print("Failed to replace classifier: \(error)")
|
|
return false
|
|
}
|
|
DispatchQueue.main.async {
|
|
self.classifierVersion = self.serverClassifierVersion
|
|
self.classifier = nil
|
|
}
|
|
print("Downloaded classifier \(classifierVersion)")
|
|
return true
|
|
}
|
|
|
|
/**
|
|
Indicate that the cap has pending operations, such as determining the color or a thumbnail
|
|
*/
|
|
func hasPendingOperations(for cap: Int) -> Bool {
|
|
return false
|
|
}
|
|
|
|
// MARK: Adding new data
|
|
|
|
func save(newCap name: String) -> Cap {
|
|
let cap = Cap(id: nextCapId, name: name, classifier: serverClassifierVersion)
|
|
caps[cap.id] = cap
|
|
#warning("Upload new cap")
|
|
return cap
|
|
}
|
|
|
|
@discardableResult
|
|
func save(_ image: UIImage, for capId: Int) -> Bool {
|
|
guard caps[capId] != nil else {
|
|
log("Failed to save image for missing cap \(capId)")
|
|
return false
|
|
}
|
|
guard ensureFolderExistence(imageUploadFolderUrl) else {
|
|
return false
|
|
}
|
|
guard let data = image.jpegData(compressionQuality: imageCompressionQuality) else {
|
|
log("Failed to compress image for cap: \(capId)")
|
|
return false
|
|
}
|
|
let hash = Data(SHA256.hash(data: data)).hexEncoded.prefix(16)
|
|
let url = imageUploadFolderUrl.appendingPathComponent("\(capId)-\(hash).jpg")
|
|
do {
|
|
try data.write(to: url)
|
|
} catch {
|
|
log("Failed to save \(url.lastPathComponent): \(error)")
|
|
return false
|
|
}
|
|
log("Saved \(url.lastPathComponent) for upload")
|
|
caps[capId]?.imageCount += 1
|
|
updateImageUploadCounts()
|
|
return true
|
|
}
|
|
|
|
private func updateImageUploadCounts() {
|
|
|
|
}
|
|
|
|
private func loadImageUploadCounts() -> [Int : Int] {
|
|
var result = [Int : Int]()
|
|
pendingImageUploads.forEach { url in
|
|
guard let capId = capId(from: url) else {
|
|
return
|
|
}
|
|
if let old = result[capId] {
|
|
result[capId] = old + 1
|
|
} else {
|
|
result[capId] = 1
|
|
}
|
|
}
|
|
return result
|
|
}
|
|
|
|
// MARK: Uploads
|
|
|
|
func startRegularUploads() {
|
|
guard uploadTimer != nil else {
|
|
return
|
|
}
|
|
DispatchQueue.main.async {
|
|
self.uploadTimer = Timer.scheduledTimer(withTimeInterval: 5, repeats: true, block: self.uploadTimerElapsed)
|
|
}
|
|
}
|
|
|
|
private func uploadTimerElapsed(timer: Timer) {
|
|
Task {
|
|
await uploadAll()
|
|
}
|
|
}
|
|
|
|
private func uploadAll() async {
|
|
guard !isUploading else {
|
|
return
|
|
}
|
|
DispatchQueue.main.async {
|
|
self.isUploading = true
|
|
}
|
|
await uploadAllChangedCaps()
|
|
await uploadAllImages()
|
|
DispatchQueue.main.async {
|
|
self.isUploading = false
|
|
}
|
|
}
|
|
|
|
/**
|
|
Indicate that the cap has pending uploads, either changes or images
|
|
*/
|
|
func hasPendingUpdates(for cap: Int) -> Bool {
|
|
changedCaps.contains(cap) || imageUploads[cap] != nil
|
|
}
|
|
|
|
private var pendingImageUploads: [URL] {
|
|
(try? fm.contentsOfDirectory(at: imageUploadFolderUrl, includingPropertiesForKeys: nil)) ?? []
|
|
}
|
|
|
|
var pendingImageUploadCount: Int {
|
|
pendingImageUploads.count
|
|
}
|
|
|
|
private func capId(from url: URL) -> Int? {
|
|
Int(url.lastPathComponent.components(separatedBy: "-").first!)
|
|
}
|
|
|
|
private func uploadAllImages() async {
|
|
guard hasServerAuthentication else {
|
|
log("No server authentication to upload to server")
|
|
return
|
|
}
|
|
updateImageUploadCounts()
|
|
for url in pendingImageUploads {
|
|
guard let capId = capId(from: url) else {
|
|
log("Unexpected image \(url.lastPathComponent) in upload folder")
|
|
continue
|
|
}
|
|
guard await upload(imageAt: url, for: capId) else {
|
|
continue
|
|
}
|
|
do {
|
|
try fm.removeItem(at: url)
|
|
updateImageUploadCounts()
|
|
} catch {
|
|
log("Failed to remove uploaded image \(url.lastPathComponent): \(error)")
|
|
}
|
|
}
|
|
}
|
|
|
|
@discardableResult
|
|
private func upload(imageAt url: URL, for cap: Int) async -> Bool {
|
|
guard let key = serverAuthenticationKey else {
|
|
return false
|
|
}
|
|
let url = serverUrl
|
|
.appendingPathComponent("images")
|
|
.appendingPathComponent("\(cap)?key=\(key)")
|
|
var request = URLRequest(url: url)
|
|
request.httpMethod = "POST"
|
|
do {
|
|
let (_, response) = try await URLSession.shared.upload(for: request, fromFile: url)
|
|
guard let httpResponse = response as? HTTPURLResponse else {
|
|
log("Unexpected response for upload of image \(url.lastPathComponent): \(response)")
|
|
return false
|
|
}
|
|
guard httpResponse.statusCode == 200 else {
|
|
log("Failed to upload image \(url.lastPathComponent): Response \(httpResponse.statusCode)")
|
|
return false
|
|
}
|
|
return true
|
|
} catch {
|
|
log("Failed to upload image \(url.lastPathComponent): \(error)")
|
|
return false
|
|
}
|
|
}
|
|
|
|
var pendingCapUploadCount: Int {
|
|
changedCaps.count
|
|
}
|
|
|
|
private func uploadAllChangedCaps() async {
|
|
guard hasServerAuthentication else {
|
|
log("No server authentication to upload to server")
|
|
return
|
|
}
|
|
var uploaded = Set<Int>()
|
|
for capId in changedCaps {
|
|
guard let cap = caps[capId] else {
|
|
uploaded.insert(capId)
|
|
continue
|
|
}
|
|
guard await upload(cap: cap) else {
|
|
continue
|
|
}
|
|
uploaded.insert(capId)
|
|
}
|
|
changedCaps.subtract(uploaded)
|
|
}
|
|
|
|
@discardableResult
|
|
private func upload(cap: Cap) async -> Bool {
|
|
guard let key = serverAuthenticationKey else {
|
|
return false
|
|
}
|
|
let data: Data
|
|
do {
|
|
/// `Cap` and `CapData` have equivalent JSON layout
|
|
data = try encoder.encode(cap)
|
|
} catch {
|
|
log("Failed to encode cap \(cap.id) for upload: \(error)")
|
|
return false
|
|
}
|
|
let url = serverUrl
|
|
.appendingPathComponent("images")
|
|
.appendingPathComponent("\(cap)?key=\(key)")
|
|
var request = URLRequest(url: url)
|
|
request.httpMethod = "POST"
|
|
do {
|
|
let (_, response) = try await URLSession.shared.upload(for: request, from: data)
|
|
guard let httpResponse = response as? HTTPURLResponse else {
|
|
log("Unexpected response for upload of cap \(cap.id): \(response)")
|
|
return false
|
|
}
|
|
guard httpResponse.statusCode == 200 else {
|
|
log("Failed to upload cap \(cap.id): Response \(httpResponse.statusCode)")
|
|
return false
|
|
}
|
|
changedCaps.remove(cap.id)
|
|
return true
|
|
} catch {
|
|
log("Failed to upload cap \(cap.id): \(error)")
|
|
return false
|
|
}
|
|
}
|
|
|
|
// MARK: Classification
|
|
|
|
/// The compiled recognition model on disk
|
|
private var recognitionModel: VNCoreMLModel? {
|
|
guard fm.fileExists(atPath: localClassifierUrl.path) else {
|
|
log("No recognition model to load from disk")
|
|
return nil
|
|
}
|
|
do {
|
|
log("Loading model from disk")
|
|
let newUrl = try MLModel.compileModel(at: localClassifierUrl)
|
|
let model = try MLModel(contentsOf: newUrl)
|
|
return try VNCoreMLModel(for: model)
|
|
} catch {
|
|
log("Failed to load recognition model: \(error)")
|
|
return nil
|
|
}
|
|
}
|
|
|
|
private func classifyImage() {
|
|
guard let image = image?.cgImage else {
|
|
matches.removeAll()
|
|
log("Image removed")
|
|
return
|
|
}
|
|
DispatchQueue.global().async {
|
|
guard let classifier = self.getClassifier() else {
|
|
return
|
|
}
|
|
log("Image classification started")
|
|
classifier.recognize(image: image) { matches in
|
|
DispatchQueue.main.async {
|
|
self.matches = matches ?? [:]
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
private func getClassifier() -> Classifier? {
|
|
if let classifier = classifier {
|
|
return classifier
|
|
}
|
|
guard let model = recognitionModel else {
|
|
return nil
|
|
}
|
|
return Classifier(model: model)
|
|
}
|
|
|
|
// MARK: Statistics
|
|
|
|
var numberOfCaps: Int {
|
|
caps.count
|
|
}
|
|
|
|
var numberOfImages: Int {
|
|
caps.values.reduce(0) { $0 + $1.imageCount }
|
|
}
|
|
|
|
var averageImageCount: Float {
|
|
Float(numberOfImages) / Float(numberOfCaps)
|
|
}
|
|
|
|
@AppStorage("classifier")
|
|
private(set) var classifierVersion = 0
|
|
|
|
@AppStorage("serverClassifier")
|
|
private(set) var serverClassifierVersion = 0
|
|
|
|
var classifierClassCount: Int {
|
|
let version = classifierVersion
|
|
return caps.values.filter { $0.classifiable(by: version) }.count
|
|
}
|
|
|
|
var imageCacheSize: Int {
|
|
imageCache.currentDiskUsage
|
|
}
|
|
|
|
var databaseSize: Int {
|
|
localDbUrl.fileSize
|
|
}
|
|
|
|
var classifierSize: Int {
|
|
localClassifierUrl.fileSize
|
|
}
|
|
}
|
|
|
|
extension Database {
|
|
|
|
static var mock: Database {
|
|
let db = Database(server: URL(string: "https://christophhagen.de/caps")!)
|
|
db.caps = [
|
|
Cap(id: 123, name: "My new cap"),
|
|
Cap(id: 234, name: "My favorite cap"),
|
|
Cap(id: 345, name: "My oldest cap"),
|
|
Cap(id: 456, name: "My new cap"),
|
|
Cap(id: 567, name: "My favorite cap"),
|
|
Cap(id: 678, name: "My oldest cap"),
|
|
].reduce(into: [:]) { $0[$1.id] = $1 }
|
|
db.image = UIImage(systemSymbol: .photo)
|
|
return db
|
|
}
|
|
}
|