622 lines
20 KiB
Swift
622 lines
20 KiB
Swift
//
|
|
// Database.swift
|
|
// CapCollector
|
|
//
|
|
// Created by Christoph on 14.04.20.
|
|
// Copyright © 2020 CH. All rights reserved.
|
|
//
|
|
|
|
import Foundation
|
|
import UIKit
|
|
import CoreML
|
|
import SQLite
|
|
|
|
protocol DatabaseDelegate: class {
|
|
|
|
func database(didChangeCap cap: Int)
|
|
|
|
func database(didAddCap cap: Cap)
|
|
|
|
func databaseRequiresFullRefresh()
|
|
}
|
|
|
|
struct Weak {
|
|
|
|
weak var value : DatabaseDelegate?
|
|
|
|
init (_ value: DatabaseDelegate) {
|
|
self.value = value
|
|
}
|
|
}
|
|
|
|
extension Array where Element == Weak {
|
|
|
|
mutating func reap () {
|
|
self = self.filter { $0.value != nil }
|
|
}
|
|
}
|
|
|
|
final class Database {
|
|
|
|
// MARK: Variables
|
|
|
|
let db: Connection
|
|
|
|
let upload: Upload
|
|
|
|
let download: Download
|
|
|
|
private var listeners = [Weak]()
|
|
|
|
// MARK: Listeners
|
|
|
|
func add(listener: DatabaseDelegate) {
|
|
listeners.append(Weak(listener))
|
|
}
|
|
|
|
init?(url: URL, server: URL) {
|
|
guard let db = try? Connection(url.path) else {
|
|
return nil
|
|
}
|
|
|
|
let upload = Upload(server: server)
|
|
let download = Download(server: server)
|
|
|
|
do {
|
|
try db.run(Cap.createQuery)
|
|
try db.run(upload.createQuery)
|
|
} catch {
|
|
return nil
|
|
}
|
|
|
|
self.db = db
|
|
self.upload = upload
|
|
self.download = download
|
|
log("Database loaded with \(capCount) caps")
|
|
}
|
|
|
|
// MARK: Computed properties
|
|
|
|
/// All caps currently in the database
|
|
var caps: [Cap] {
|
|
(try? db.prepare(Cap.table))?.map(Cap.init) ?? []
|
|
}
|
|
|
|
/// The ids of the caps which weren't included in the last classification
|
|
var unmatchedCaps: [Int] {
|
|
let query = Cap.table.select(Cap.rowId).filter(Cap.rowMatched == false)
|
|
return (try? db.prepare(query).map { $0[Cap.rowId] }) ?? []
|
|
}
|
|
|
|
/// The number of caps which could be recognized during the last classification
|
|
var recognizedCapCount: Int {
|
|
(try? db.scalar(Cap.table.filter(Cap.rowMatched == true).count)) ?? 0
|
|
}
|
|
|
|
/// The number of caps currently in the database
|
|
var capCount: Int {
|
|
(try? db.scalar(Cap.table.count)) ?? 0
|
|
}
|
|
|
|
/// The total number of images for all caps
|
|
var imageCount: Int {
|
|
(try? db.prepare(Cap.table).reduce(0) { $0 + $1[Cap.rowCount] }) ?? 0
|
|
}
|
|
|
|
/// The number of caps without a downloaded image
|
|
var capsWithoutImages: Int {
|
|
caps.filter({ !$0.hasImage }).count
|
|
}
|
|
|
|
|
|
|
|
var pendingUploads: [(cap: Int, version: Int)] {
|
|
do {
|
|
return try db.prepare(upload.table).map { row in
|
|
(cap: row[upload.rowCapId], version: row[upload.rowCapVersion])
|
|
}
|
|
} catch {
|
|
log("Failed to get pending uploads")
|
|
return []
|
|
}
|
|
}
|
|
|
|
/// Indicate if there are any unfinished uploads
|
|
var hasPendingUploads: Bool {
|
|
((try? db.scalar(upload.table.count)) ?? 0) > 0
|
|
}
|
|
|
|
var classifierVersion: Int {
|
|
set {
|
|
UserDefaults.standard.set(newValue, forKey: Classifier.userDefaultsKey)
|
|
log("Classifier version set to \(newValue)")
|
|
}
|
|
get {
|
|
UserDefaults.standard.integer(forKey: Classifier.userDefaultsKey)
|
|
}
|
|
}
|
|
|
|
// MARK: Data updates
|
|
|
|
/**
|
|
Create a new cap with an image.
|
|
|
|
The cap is inserted into the database, and the name and image will be uploaded to the server.
|
|
|
|
- parameter image: The main image of the cap
|
|
- parameter name: The name of the cap
|
|
- note: Must be called on the main queue.
|
|
- note: The registered delegates will be informed about the added cap through `database(didAddCap:)`
|
|
- returns: `true`, if the cap was created.
|
|
*/
|
|
func createCap(image: UIImage, name: String) -> Bool {
|
|
guard let color = image.averageColor else {
|
|
return false
|
|
}
|
|
let cap = Cap(name: name, id: capCount, color: color)
|
|
guard insert(cap: cap) else {
|
|
return false
|
|
}
|
|
guard app.storage.save(image: image, for: cap.id) else {
|
|
return false
|
|
}
|
|
listeners.forEach { $0.value?.database(didAddCap: cap) }
|
|
upload.upload(name: name, for: cap.id) { success in
|
|
guard success else {
|
|
return
|
|
}
|
|
self.update(uploaded: true, for: cap.id)
|
|
self.upload.uploadImage(for: cap.id, version: 0) { count in
|
|
guard let count = count else {
|
|
self.log("Failed to upload first image for cap \(cap.id)")
|
|
return
|
|
}
|
|
self.log("Uploaded first image for cap \(cap.id)")
|
|
self.update(count: count, for: cap.id)
|
|
}
|
|
}
|
|
return true
|
|
}
|
|
|
|
/**
|
|
Insert a new cap.
|
|
|
|
Only inserts the cap into the database, and optionally notifies the delegates.
|
|
- note: When a new cap is created, use `createCap(image:name:)` instead
|
|
*/
|
|
@discardableResult
|
|
private func insert(cap: Cap, notifyDelegate: Bool = true) -> Bool {
|
|
do {
|
|
try db.run(cap.insertQuery)
|
|
if notifyDelegate {
|
|
listeners.forEach { $0.value?.database(didAddCap: cap) }
|
|
}
|
|
return true
|
|
} catch {
|
|
log("Failed to insert cap \(cap.id): \(error)")
|
|
return false
|
|
}
|
|
}
|
|
|
|
func add(image: UIImage, for cap: Int) -> Bool {
|
|
guard let version = count(for: cap) else {
|
|
log("Failed to get count for cap \(cap)")
|
|
return false
|
|
}
|
|
guard app.storage.save(image: image, for: cap, version: version) else {
|
|
log("Failed to save image \(version) for cap \(cap) to disk")
|
|
return false
|
|
}
|
|
guard update(count: version + 1, for: cap) else {
|
|
log("Failed update count \(version) for cap \(cap)")
|
|
return false
|
|
}
|
|
listeners.forEach { $0.value?.database(didChangeCap: cap) }
|
|
|
|
guard addPendingUpload(for: cap, version: version) else {
|
|
log("Failed to add cap \(cap) version \(version) to upload queue")
|
|
return false
|
|
}
|
|
upload.uploadImage(for: cap, version: version) { count in
|
|
guard let _ = count else {
|
|
self.log("Failed to upload image \(version) for cap \(cap)")
|
|
return
|
|
}
|
|
guard self.removePendingUpload(of: cap, version: version) else {
|
|
self.log("Failed to remove version \(version) for cap \(cap) from upload queue")
|
|
return
|
|
}
|
|
self.log("Uploaded version \(version) for cap \(cap)")
|
|
}
|
|
return true
|
|
}
|
|
|
|
// MARK: Updating cap properties
|
|
|
|
private func update(_ property: String, for cap: Int, setter: Setter...) -> Bool {
|
|
do {
|
|
let query = updateQuery(for: cap).update(setter)
|
|
try db.run(query)
|
|
listeners.forEach { $0.value?.database(didChangeCap: cap) }
|
|
return true
|
|
} catch {
|
|
log("Failed to update \(property) for cap \(cap): \(error)")
|
|
return false
|
|
}
|
|
}
|
|
|
|
@discardableResult
|
|
private func update(uploaded: Bool, for cap: Int) -> Bool {
|
|
update("uploaded", for: cap, setter: Cap.rowUploaded <- uploaded)
|
|
}
|
|
|
|
@discardableResult
|
|
func update(name: String, for cap: Int) -> Bool {
|
|
update("name", for: cap, setter: Cap.rowName <- name)
|
|
}
|
|
|
|
@discardableResult
|
|
func update(color: UIColor, for cap: Int) -> Bool {
|
|
let (red, green, blue) = color.rgb
|
|
return update("color", for: cap, setter: Cap.rowRed <- red, Cap.rowGreen <- green, Cap.rowBlue <- blue)
|
|
}
|
|
|
|
@discardableResult
|
|
func update(tile: Int, for cap: Int) -> Bool {
|
|
update("tile", for: cap, setter: Cap.rowTile <- tile)
|
|
}
|
|
|
|
@discardableResult
|
|
func update(count: Int, for cap: Int) -> Bool {
|
|
update("count", for: cap, setter: Cap.rowCount <- count)
|
|
}
|
|
|
|
@discardableResult
|
|
func update(matched: Bool, for cap: Int) -> Bool {
|
|
update("matched", for: cap, setter: Cap.rowMatched <- matched)
|
|
}
|
|
|
|
func update(recognizedCaps: Set<Int>) {
|
|
let unrecognized = self.unmatchedCaps
|
|
// Update caps which haven't been recognized before
|
|
let newlyRecognized = recognizedCaps.intersection(unrecognized)
|
|
let logIndividualMessages = newlyRecognized.count < 10
|
|
if !logIndividualMessages {
|
|
log("Marking \(newlyRecognized.count) caps as matched")
|
|
}
|
|
for cap in newlyRecognized {
|
|
if logIndividualMessages {
|
|
log("Marking cap \(cap) as matched")
|
|
}
|
|
update(matched: true, for: cap)
|
|
}
|
|
// Update caps which are no longer recognized
|
|
let missing = Set(1...capCount).subtracting(recognizedCaps).subtracting(unrecognized)
|
|
for cap in missing {
|
|
log("Marking cap \(cap) as not matched")
|
|
update(matched: false, for: cap)
|
|
}
|
|
}
|
|
|
|
func addPendingUpload(for cap: Int, version: Int) -> Bool {
|
|
do {
|
|
try db.run(upload.insertQuery(for: cap, version: version))
|
|
return true
|
|
} catch {
|
|
log("Failed to add pending upload of cap \(cap) version \(version): \(error)")
|
|
return false
|
|
}
|
|
}
|
|
|
|
func removePendingUpload(for cap: Int, version: Int) -> Bool {
|
|
do {
|
|
try db.run(upload.deleteQuery(for: cap, version: version))
|
|
return true
|
|
} catch {
|
|
log("Failed to remove pending upload of cap \(cap) version \(version): \(error)")
|
|
return false
|
|
}
|
|
}
|
|
|
|
// MARK: Information retrieval
|
|
|
|
func cap(for id: Int) -> Cap? {
|
|
do {
|
|
guard let row = try db.pluck(updateQuery(for: id)) else {
|
|
log("No cap with id \(id) in database")
|
|
return nil
|
|
}
|
|
return Cap(row: row)
|
|
} catch {
|
|
log("Failed to get cap \(id): \(error)")
|
|
return nil
|
|
}
|
|
}
|
|
|
|
private func count(for cap: Int) -> Int? {
|
|
do {
|
|
let row = try db.pluck(updateQuery(for: cap).select(Cap.rowCount))
|
|
return row?[Cap.rowCount]
|
|
} catch {
|
|
log("Failed to get count for cap \(cap)")
|
|
return nil
|
|
}
|
|
}
|
|
|
|
func countOfCaps(withImageCountLessThan limit: Int) -> Int {
|
|
do {
|
|
return try db.scalar(Cap.table.filter(Cap.rowCount < limit).count)
|
|
} catch {
|
|
log("Failed to get caps with less than \(limit) images")
|
|
return 0
|
|
}
|
|
}
|
|
|
|
func lowestImageCountForCaps(startingAt start: Int) -> (count: Int, numberOfCaps: Int) {
|
|
do {
|
|
var currentCount = start - 1
|
|
var capsFound = 0
|
|
repeat {
|
|
currentCount += 1
|
|
capsFound = try db.scalar(Cap.table.filter(Cap.rowCount == currentCount).count)
|
|
} while capsFound == 0
|
|
|
|
return (currentCount, capsFound)
|
|
} catch {
|
|
return (0,0)
|
|
}
|
|
}
|
|
|
|
func updateQuery(for cap: Int) -> Table {
|
|
Cap.table.filter(Cap.rowId == cap)
|
|
}
|
|
|
|
// MARK: Downloads
|
|
|
|
@discardableResult
|
|
func downloadMainImage(for cap: Int, completion: @escaping (_ image: UIImage?) -> Void) -> Bool {
|
|
return download.image(for: cap, version: 0) { image in
|
|
// Guaranteed to be on the main queue
|
|
guard let image = image else {
|
|
completion(nil)
|
|
return
|
|
}
|
|
|
|
defer {
|
|
completion(image)
|
|
}
|
|
if !app.storage.save(thumbnail: Cap.thumbnail(for: image), for: cap) {
|
|
self.log("Failed to save thumbnail for cap \(cap)")
|
|
}
|
|
guard let color = image.averageColor else {
|
|
self.log("Failed to calculate color for cap \(cap)")
|
|
return
|
|
}
|
|
self.update(color: color, for: cap)
|
|
}
|
|
}
|
|
|
|
@discardableResult
|
|
func downloadImage(for cap: Int, version: Int, completion: @escaping (_ image: UIImage?) -> Void) -> Bool {
|
|
return download.image(for: cap, version: version, completion: completion)
|
|
}
|
|
|
|
func getServerDatabaseSize(completion: @escaping (_ size: Int64?) -> Void) {
|
|
download.databaseSize(completion: completion)
|
|
}
|
|
|
|
func downloadServerDatabase(progress: Download.Delegate.ProgressHandler? = nil, completion: @escaping (_ success: Bool) -> Void, processed: (() -> Void)? = nil) {
|
|
download.database(progress: progress) { tempUrl in
|
|
guard let url = tempUrl else {
|
|
self.log("Failed to download database")
|
|
completion(false)
|
|
return
|
|
}
|
|
completion(true)
|
|
self.processServerDatabase(at: url)
|
|
processed?()
|
|
}
|
|
}
|
|
|
|
func downloadMainCapImages(progress: @escaping (_ current: Int, _ total: Int) -> Void) {
|
|
let caps = self.caps.filter { !$0.hasImage }.map { $0.id }
|
|
|
|
var downloaded = 0
|
|
let total = caps.count
|
|
|
|
func update() {
|
|
DispatchQueue.main.async {
|
|
progress(downloaded, total)
|
|
}
|
|
}
|
|
update()
|
|
|
|
guard total > 0 else {
|
|
log("No images to download")
|
|
return
|
|
}
|
|
log("Starting to download \(total) images")
|
|
|
|
let group = DispatchGroup()
|
|
let split = 50
|
|
DispatchQueue.global(qos: .userInitiated).async {
|
|
for part in caps.split(intoPartsOf: split) {
|
|
for id in part {
|
|
let downloading = self.downloadMainImage(for: id) { _ in
|
|
group.leave()
|
|
}
|
|
if downloading {
|
|
group.enter()
|
|
}
|
|
}
|
|
if group.wait(timeout: .now() + .seconds(30)) != .success {
|
|
self.log("Timed out waiting for images to be downloaded")
|
|
}
|
|
downloaded += part.count
|
|
self.log("Finished \(downloaded) of \(total) image downloads")
|
|
update()
|
|
}
|
|
self.log("Finished all image downloads")
|
|
}
|
|
|
|
}
|
|
|
|
func hasNewClassifier(completion: @escaping (_ version: Int?, _ size: Int64?) -> Void) {
|
|
download.classifierVersion { version in
|
|
guard let version = version else {
|
|
self.log("Failed to download server model version")
|
|
completion(nil, nil)
|
|
return
|
|
}
|
|
let ownVersion = self.classifierVersion
|
|
guard ownVersion < version else {
|
|
self.log("Not updating classifier: Own version \(ownVersion), server version \(version)")
|
|
completion(nil, nil)
|
|
return
|
|
}
|
|
self.log("Getting classifier size: Own version \(ownVersion), server version \(version)")
|
|
self.download.classifierSize { size in
|
|
completion(version, size)
|
|
}
|
|
}
|
|
}
|
|
|
|
func downloadClassifier(progress: Download.Delegate.ProgressHandler? = nil, completion: @escaping (_ success: Bool) -> Void) {
|
|
download.classifier(progress: progress) { url in
|
|
guard let url = url else {
|
|
self.log("Failed to download classifier")
|
|
completion(false)
|
|
return
|
|
}
|
|
let compiledUrl: URL
|
|
do {
|
|
compiledUrl = try MLModel.compileModel(at: url)
|
|
} catch {
|
|
self.log("Failed to compile downloaded classifier: \(error)")
|
|
completion(false)
|
|
return
|
|
}
|
|
|
|
guard app.storage.save(recognitionModelAt: compiledUrl) else {
|
|
self.log("Failed to save classifier")
|
|
completion(false)
|
|
return
|
|
}
|
|
completion(true)
|
|
self.download.classifierVersion { version in
|
|
guard let version = version else {
|
|
self.log("Failed to download classifier version")
|
|
return
|
|
}
|
|
self.classifierVersion = version
|
|
}
|
|
}
|
|
}
|
|
|
|
func downloadImageCounts() {
|
|
guard !hasPendingUploads else {
|
|
log("Waiting to refresh server image counts (uploads pending)")
|
|
return
|
|
}
|
|
log("Refreshing server image counts")
|
|
app.database.download.imageCounts { counts in
|
|
guard let counts = counts else {
|
|
self.log("Failed to download server image counts")
|
|
return
|
|
}
|
|
self.didDownload(imageCounts: counts)
|
|
|
|
}
|
|
}
|
|
|
|
private func didDownload(imageCounts newCounts: [(cap: Int, count: Int)]) {
|
|
let capsCounts = self.caps.reduce(into: [:]) { $0[$1.id] = $1.count }
|
|
if newCounts.count != capsCounts.count {
|
|
log("Downloaded \(newCounts.count) image counts, but \(app.database.capCount) caps stored locally")
|
|
return
|
|
}
|
|
let changed = newCounts.compactMap { id, newCount -> Int? in
|
|
guard let oldCount = capsCounts[id] else {
|
|
log("Received count \(newCount) for unknown cap \(id)")
|
|
return nil
|
|
}
|
|
guard oldCount != newCount else {
|
|
return nil
|
|
}
|
|
app.database.update(count: newCount, for: id)
|
|
return id
|
|
}
|
|
switch changed.count {
|
|
case 0:
|
|
log("Refreshed image counts for all caps without changes")
|
|
case 1:
|
|
log("Refreshed image counts for caps, changed cap \(changed[0])")
|
|
case 2...10:
|
|
log("Refreshed image counts for caps \(changed.map(String.init).joined(separator: ", ")).")
|
|
default:
|
|
log("Refreshed image counts for all caps (\(changed.count) changed)")
|
|
}
|
|
}
|
|
|
|
private func processServerDatabase(at url: URL) {
|
|
guard let db = ServerDatabase(downloadedTo: url) else {
|
|
log("Failed to open downloaded server database")
|
|
return
|
|
}
|
|
for (id, count, name) in db.caps {
|
|
let cap = Cap(id: id, name: name, count: count)
|
|
insert(cap: cap, notifyDelegate: false)
|
|
}
|
|
listeners.forEach { $0.value?.databaseRequiresFullRefresh() }
|
|
}
|
|
|
|
func uploadRemainingImages() {
|
|
guard pendingUploads.count > 0 else {
|
|
log("No pending uploads")
|
|
return
|
|
}
|
|
log("\(pendingUploads.count) image uploads pending")
|
|
|
|
for (cap, version) in pendingUploads {
|
|
upload.uploadImage(for: cap, version: version) { count in
|
|
guard let _ = count else {
|
|
self.log("Failed to upload version \(version) of cap \(cap)")
|
|
return
|
|
}
|
|
self.log("Uploaded version \(version) of cap \(cap)")
|
|
self.removePendingUpload(of: cap, version: version)
|
|
}
|
|
}
|
|
}
|
|
|
|
@discardableResult
|
|
func removePendingUpload(of cap: Int, version: Int) -> Bool {
|
|
do {
|
|
let query = upload.table.filter(upload.rowCapId == cap && upload.rowCapVersion == version).delete()
|
|
try db.run(query)
|
|
log("Deleted pending upload of cap \(cap) version \(version)")
|
|
return true
|
|
} catch {
|
|
log("Failed to delete pending upload of cap \(cap) version \(version)")
|
|
return false
|
|
}
|
|
}
|
|
|
|
func setMainImage(of cap: Int, to version: Int) {
|
|
guard version != 0 else {
|
|
log("No need to switch main image with itself for cap \(cap)")
|
|
return
|
|
}
|
|
upload.setMainImage(for: cap, to: version) { color in
|
|
guard let color = color else {
|
|
self.log("Could not make \(version) the main image for cap \(cap)")
|
|
return
|
|
}
|
|
self.update(color: color, for: cap)
|
|
self.listeners.forEach { $0.value?.database(didChangeCap: cap) }
|
|
}
|
|
}
|
|
}
|
|
|
|
extension Database: Logger { }
|