598 lines
18 KiB
Swift
598 lines
18 KiB
Swift
import Foundation
|
||
import Vapor
|
||
import Clairvoyant
|
||
|
||
final class CapServer {
|
||
|
||
// MARK: Paths
|
||
|
||
private let imageFolder: URL
|
||
|
||
private let thumbnailFolder: URL
|
||
|
||
/// The file where the cap count is stored for the grid webpage
|
||
private let gridCountFile: URL
|
||
|
||
/// The file where the database of caps is stored
|
||
private let dbFile: URL
|
||
|
||
/// The file to store the HTML info of the cap count
|
||
private let htmlFile: URL
|
||
|
||
private let classifierVersionFile: URL
|
||
|
||
private let classifierFile: URL
|
||
|
||
private let changedImagesFile: URL
|
||
|
||
private let fm = FileManager.default
|
||
|
||
private let changedImageEntryDateFormatter: DateFormatter
|
||
|
||
/// Indicates that the data is loaded
|
||
private(set) var isOperational = false
|
||
|
||
// MARK: Caps
|
||
|
||
|
||
/// The changed images not yet written to disk
|
||
private var unwrittenImageChanges: [(cap: Int, image: Int)] = []
|
||
|
||
var classifierVersion: Int = 0 {
|
||
didSet {
|
||
writeClassifierVersion()
|
||
Task {
|
||
try? await classifierMetric.update(classifierVersion)
|
||
}
|
||
}
|
||
}
|
||
|
||
/**
|
||
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?
|
||
|
||
private var caps = [Int: Cap]() {
|
||
didSet {
|
||
scheduleSave()
|
||
Task {
|
||
try? await capCountMetric.update(caps.count)
|
||
try? await imageCountMetric.update(imageCount)
|
||
}
|
||
}
|
||
}
|
||
|
||
var capCount: Int {
|
||
caps.count
|
||
}
|
||
|
||
var imageCount: Int {
|
||
caps.reduce(0) { $0 + $1.value.count }
|
||
}
|
||
|
||
init(in folder: URL) async {
|
||
self.imageFolder = folder.appendingPathComponent("images")
|
||
self.thumbnailFolder = folder.appendingPathComponent("thumbnails")
|
||
self.gridCountFile = folder.appendingPathComponent("count.js")
|
||
self.dbFile = folder.appendingPathComponent("caps.json")
|
||
self.htmlFile = folder.appendingPathComponent("count.html")
|
||
self.classifierVersionFile = folder.appendingPathComponent("classifier.version")
|
||
self.classifierFile = folder.appendingPathComponent("classifier.mlmodel")
|
||
self.changedImagesFile = folder.appendingPathComponent("changes.txt")
|
||
self.changedImageEntryDateFormatter = DateFormatter()
|
||
changedImageEntryDateFormatter.dateFormat = "yy-MM-dd-HH-mm-ss"
|
||
|
||
// Metric initializers only fail if observer is missing or ID is duplicate
|
||
self.capCountMetric = try! await .init("caps.count",
|
||
name: "Number of caps",
|
||
description: "The total number of caps in the database")
|
||
|
||
self.imageCountMetric = try! await .init("caps.images",
|
||
name: "Total images",
|
||
description: "The total number of images for all caps")
|
||
|
||
self.classifierMetric = try! await .init("caps.classifier",
|
||
name: "Classifier Version",
|
||
description: "The current version of the image classifier")
|
||
}
|
||
|
||
func loadData() throws {
|
||
loadClassifierVersion(at: classifierVersionFile)
|
||
try loadCaps()
|
||
saveCapCountHTML()
|
||
updateGridCapCount()
|
||
try ensureExistenceOfChangedImagesFile()
|
||
organizeImages()
|
||
isOperational = true
|
||
}
|
||
|
||
private func loadClassifierVersion(at url: URL) {
|
||
guard exists(url) else {
|
||
return
|
||
}
|
||
let content: String
|
||
do {
|
||
content = try String(contentsOf: url)
|
||
.trimmingCharacters(in: .whitespacesAndNewlines)
|
||
} catch {
|
||
log("Failed to read classifier version file: \(error)")
|
||
return
|
||
}
|
||
|
||
guard let value = Int(content) else {
|
||
log("Invalid classifier version: \(content)")
|
||
return
|
||
}
|
||
self.classifierVersion = value
|
||
}
|
||
|
||
private func writeClassifierVersion() {
|
||
do {
|
||
try "\(classifierVersion)".data(using: .utf8)!
|
||
.write(to: classifierVersionFile)
|
||
} catch {
|
||
log("Failed to save classifier version: \(error)")
|
||
}
|
||
}
|
||
|
||
private func loadCaps() throws {
|
||
guard exists(dbFile) else {
|
||
log("No cap database found")
|
||
return
|
||
}
|
||
do {
|
||
let data = try Data(contentsOf: dbFile)
|
||
caps = try JSONDecoder().decode([Cap].self, from: data)
|
||
.reduce(into: [:]) { $0[$1.id] = $1 }
|
||
} catch {
|
||
log("Failed to load caps: \(error)")
|
||
throw error
|
||
}
|
||
log("\(caps.count) caps loaded")
|
||
}
|
||
|
||
private func scheduleSave() {
|
||
nextSaveTime = Date().addingTimeInterval(saveDelay)
|
||
DispatchQueue.global().asyncAfter(deadline: .now() + saveDelay) {
|
||
self.performScheduledSave()
|
||
}
|
||
}
|
||
|
||
private func performScheduledSave() {
|
||
guard let date = nextSaveTime else {
|
||
// No save necessary, or already saved
|
||
return
|
||
}
|
||
guard date < Date() else {
|
||
// Save pushed to future
|
||
return
|
||
}
|
||
do {
|
||
try saveCaps()
|
||
nextSaveTime = nil
|
||
} catch {
|
||
// Attempt save again
|
||
scheduleSave()
|
||
}
|
||
}
|
||
|
||
private func saveCaps() throws {
|
||
let data = try JSONEncoder().encode(caps.values.sorted())
|
||
try data.write(to: dbFile)
|
||
}
|
||
|
||
private func saveCapCountHTML() {
|
||
let count = caps.count
|
||
let content =
|
||
"""
|
||
<body style="margin: 0;overflow: hidden">
|
||
<div style="display: flex; justify-content: center;">
|
||
<div style="font-size: 60px; font-family: 'SF Pro Display',-apple-system,BlinkMacSystemFont,Helvetica,sans-serif; -webkit-font-smoothing: antialiased;">\(count)</div>
|
||
</div>
|
||
</body>
|
||
"""
|
||
try? content.data(using: .utf8)!.write(to: htmlFile)
|
||
}
|
||
|
||
private func organizeImages() {
|
||
caps.values.sorted().forEach(organizeImages)
|
||
}
|
||
|
||
private func createImageFolder(for cap: Int) throws {
|
||
let folderUrl = folder(of: cap)
|
||
do {
|
||
try fm.createDirectory(at: folderUrl, withIntermediateDirectories: true)
|
||
} catch {
|
||
log("Failed to create folder for cap \(cap): \(error)")
|
||
throw error
|
||
}
|
||
}
|
||
|
||
/**
|
||
Rearrange images of a cap to ensure that an image exists for each number from 0 to `image count - 1`.
|
||
This is done by using the last images to fill in possible gaps in the sequence.
|
||
E.g. If there are images `0`, `2`, `3`, then `3` will be renamed to `1`.
|
||
- Note: The main image is also changed, if the main image is renamed.
|
||
*/
|
||
private func organizeImages(for cap: Cap) {
|
||
var cap = cap
|
||
let folderUrl = folder(of: cap.id)
|
||
guard exists(folderUrl) else {
|
||
try? createImageFolder(for: cap.id)
|
||
cap.count = 0
|
||
caps[cap.id] = cap
|
||
log("Found cap \(cap.id) without image folder")
|
||
return
|
||
}
|
||
|
||
guard let images = try? images(in: folderUrl) else {
|
||
log("Failed to get image urls for cap \(cap.id)")
|
||
return
|
||
}
|
||
|
||
if images.count != cap.count {
|
||
log("\(images.count) instead of \(cap.count) images for cap \(cap.id)")
|
||
}
|
||
|
||
// Get list of existing images
|
||
var sorted: [(id: Int, url: URL)] = images.compactMap {
|
||
guard let id = Int($0.deletingPathExtension().lastPathComponent.components(separatedBy: "-").last!) else {
|
||
return nil
|
||
}
|
||
return (id, $0)
|
||
}.sorted { $0.id < $1.id }
|
||
|
||
// Check that all images are available
|
||
for version in 0..<images.count {
|
||
guard version != sorted[version].id else {
|
||
continue
|
||
}
|
||
let lastImage = sorted.popLast()!
|
||
let newUrl = imageUrl(of: cap.id, version: version)
|
||
do {
|
||
try fm.moveItem(at: lastImage.url, to: newUrl)
|
||
log("Moved image \(lastImage.id) to \(version) for cap \(cap.id)")
|
||
} catch {
|
||
log("Failed to move file \(lastImage.url.path) to \(newUrl.path): \(error)")
|
||
return
|
||
}
|
||
if cap.mainImage == lastImage.id {
|
||
cap.mainImage = version
|
||
}
|
||
sorted.insert((version, newUrl), at: version)
|
||
}
|
||
|
||
cap.count = sorted.count
|
||
|
||
// Fix invalid main image
|
||
if cap.mainImage >= cap.count || cap.mainImage < 0 {
|
||
cap.mainImage = 0
|
||
}
|
||
|
||
caps[cap.id] = cap
|
||
}
|
||
|
||
// MARK: Paths
|
||
|
||
func folder(of cap: Int) -> URL {
|
||
imageFolder.appendingPathComponent(String(format: "%04d", cap))
|
||
}
|
||
|
||
func thumbnail(of cap: Int) -> URL {
|
||
thumbnailFolder.appendingPathComponent(String(format: "%04d.jpg", cap))
|
||
}
|
||
|
||
func imageUrl(of cap: Int, version: Int) -> URL {
|
||
folder(of: cap).appendingPathComponent(String(format: "%04d-%02d.jpg", cap, version))
|
||
}
|
||
|
||
private func exists(_ url: URL) -> Bool {
|
||
fm.fileExists(atPath: url.path)
|
||
}
|
||
|
||
// MARK: Counts
|
||
|
||
private func images(in folder: URL) throws -> [URL] {
|
||
do {
|
||
return try fm.contentsOfDirectory(at: folder, includingPropertiesForKeys: nil)
|
||
.filter { $0.pathExtension == "jpg" }
|
||
} catch {
|
||
log("Failed to get image urls for cap \(folder.lastPathComponent): \(error)")
|
||
throw error
|
||
}
|
||
}
|
||
|
||
/**
|
||
Get the image count of a cap.
|
||
*/
|
||
func count(of cap: Int) throws -> Int {
|
||
let capImageFolder = folder(of: cap)
|
||
guard exists(capImageFolder) else {
|
||
return 0
|
||
}
|
||
return try images(in: capImageFolder).count
|
||
}
|
||
|
||
// MARK: Images
|
||
|
||
/**
|
||
Save a cap image to disk.
|
||
|
||
Automatically creates the image name with the current image count.
|
||
- Parameter data: The image data
|
||
- Parameter cap: The id of the cap.
|
||
- Throws: `CapError.unknownId`, if the cap doesn't exist. `CapError.dataInconsistency` if an image already exists for the current count.
|
||
*/
|
||
func save(image data: Data, for cap: Int) throws {
|
||
guard caps[cap] != nil else {
|
||
log("Tried to save image for unknown cap \(cap)")
|
||
throw CapError.unknownId
|
||
}
|
||
var id = 0
|
||
let capFolder = folder(of: cap)
|
||
var capImageUrl = imageUrl(of: cap, version: id)
|
||
if exists(capFolder) {
|
||
while exists(capImageUrl) {
|
||
id += 1
|
||
capImageUrl = imageUrl(of: cap, version: id)
|
||
}
|
||
} else {
|
||
try createImageFolder(for: cap)
|
||
}
|
||
do {
|
||
try data.write(to: capImageUrl)
|
||
} catch {
|
||
log("Failed to write image \(id) for cap \(cap): \(error)")
|
||
throw CapError.invalidFile
|
||
}
|
||
let count = try count(of: cap)
|
||
caps[cap]!.count = count
|
||
addChangedImageToLog(cap: cap, image: id)
|
||
log("Added image \(id) for cap \(cap) (\(count) total)")
|
||
}
|
||
|
||
private func writeChangedImagesToDisk() throws {
|
||
guard !unwrittenImageChanges.isEmpty else {
|
||
return
|
||
}
|
||
|
||
try ensureExistenceOfChangedImagesFile()
|
||
|
||
let handle: FileHandle
|
||
do {
|
||
handle = try FileHandle(forWritingTo: changedImagesFile)
|
||
try handle.seekToEnd()
|
||
} catch {
|
||
log("Failed to open changed images file for writing: \(error)")
|
||
throw error
|
||
}
|
||
|
||
var entries = unwrittenImageChanges
|
||
defer {
|
||
unwrittenImageChanges = entries
|
||
try? handle.close()
|
||
}
|
||
let dateString = changedImageEntryDateFormatter.string(from: Date())
|
||
while let entry = entries.popLast() {
|
||
let content = "\(dateString):\(entry.cap):\(entry.image)\n".data(using: .utf8)!
|
||
do {
|
||
try handle.write(contentsOf: content)
|
||
} catch {
|
||
log("Failed to write entry to changed images file: \(error)")
|
||
throw error
|
||
}
|
||
}
|
||
}
|
||
|
||
private func addChangedImageToLog(cap: Int, image: Int) {
|
||
unwrittenImageChanges.append((cap, image))
|
||
do {
|
||
try writeChangedImagesToDisk()
|
||
} catch {
|
||
log("Failed to save changed image list: \(error)")
|
||
}
|
||
}
|
||
|
||
private func ensureExistenceOfChangedImagesFile() throws {
|
||
if exists(changedImagesFile) {
|
||
return
|
||
}
|
||
do {
|
||
try Data().write(to: changedImagesFile)
|
||
} catch {
|
||
log("Failed to create changed images file: \(error)")
|
||
throw error
|
||
}
|
||
}
|
||
|
||
func removeAllEntriesInImageChangeList(before date: Date) {
|
||
guard exists(changedImagesFile) else {
|
||
log("No file for changed images to update")
|
||
return
|
||
}
|
||
do {
|
||
try String(contentsOf: changedImagesFile)
|
||
.components(separatedBy: "\n")
|
||
.filter { $0 != "" }
|
||
.compactMap { line -> String? in
|
||
guard let entryDate = changedImageEntryDateFormatter.date(from: line.components(separatedBy: ":").first!) else {
|
||
return nil
|
||
}
|
||
guard entryDate > date else {
|
||
return nil
|
||
}
|
||
return line
|
||
}
|
||
.joined(separator: "\n")
|
||
.data(using: .utf8)!
|
||
.write(to: changedImagesFile)
|
||
} catch {
|
||
log("Failed to update changed images file: \(error)")
|
||
}
|
||
}
|
||
|
||
func switchMainImage(to version: Int, for cap: Int) throws {
|
||
let capImageUrl = imageUrl(of: cap, version: version)
|
||
guard exists(capImageUrl) else {
|
||
log("No image \(version) for cap \(cap)")
|
||
throw CapError.invalidFile
|
||
}
|
||
caps[cap]?.mainImage = version
|
||
log("Switched cap \(cap) to version \(version)")
|
||
}
|
||
|
||
func addOrUpdate(_ cap: Cap) throws {
|
||
if let existingCap = caps[cap.id] {
|
||
update(existingCap, with: cap)
|
||
} else {
|
||
try add(cap)
|
||
}
|
||
}
|
||
|
||
private func add(_ cap: Cap) throws {
|
||
guard cap.mainImage == 0 else {
|
||
log("Attempting to add cap \(cap.id) with main image \(cap.mainImage)")
|
||
throw CapError.invalidData
|
||
}
|
||
var cap = cap
|
||
cap.count = 0
|
||
cap.classifierVersion = nil
|
||
caps[cap.id] = cap
|
||
saveCapCountHTML()
|
||
updateGridCapCount()
|
||
log("Added cap \(cap.id) '\(cap.name)'")
|
||
}
|
||
|
||
private func update(_ existingCap: Cap, with cap: Cap) {
|
||
var updatedCap = existingCap
|
||
if cap.name != "" {
|
||
updatedCap.name = cap.name
|
||
}
|
||
let capImageUrl = imageUrl(of: existingCap.id, version: cap.mainImage)
|
||
if exists(capImageUrl) {
|
||
updatedCap.mainImage = cap.mainImage
|
||
}
|
||
if let color = cap.color {
|
||
updatedCap.color = color
|
||
}
|
||
caps[existingCap.id] = updatedCap
|
||
log("Updated cap \(existingCap.id)")
|
||
}
|
||
|
||
func deleteImage(version: Int, for capId: Int) -> Cap? {
|
||
guard let cap = caps[capId] else {
|
||
log("Attempting to delete image \(version) of unknown cap \(capId)")
|
||
return nil
|
||
}
|
||
let capImageUrl = imageUrl(of: capId, version: version)
|
||
guard exists(capImageUrl) else {
|
||
log("Attempting to delete missing image \(version) of cap \(capId)")
|
||
return nil
|
||
}
|
||
organizeImages(for: cap)
|
||
return caps[capId]!
|
||
}
|
||
|
||
func delete(cap capId: Int) -> Bool {
|
||
guard caps[capId] != nil else {
|
||
log("Attempting to delete unknown cap \(capId)")
|
||
return false
|
||
}
|
||
// 1. Remove all images
|
||
do {
|
||
let imageFolderUrl = folder(of: capId)
|
||
if exists(imageFolderUrl) {
|
||
try fm.removeItem(at: imageFolderUrl)
|
||
}
|
||
} catch {
|
||
log("Failed to delete image folder of cap \(capId): \(error)")
|
||
return false
|
||
}
|
||
// 2. Remove thumbnail
|
||
do {
|
||
let url = thumbnail(of: capId)
|
||
if exists(url) {
|
||
try fm.removeItem(at: url)
|
||
}
|
||
} catch {
|
||
log("Failed to delete thumbnail of cap \(capId): \(error)")
|
||
return false
|
||
}
|
||
|
||
// 3. Remove cap
|
||
caps[capId] = nil
|
||
saveCapCountHTML()
|
||
updateGridCapCount()
|
||
return true
|
||
}
|
||
|
||
// MARK: Classifier
|
||
|
||
func updateTrainedClasses(content: String) {
|
||
let trainedCaps = content
|
||
.components(separatedBy: "\n")
|
||
.compactMap(Int.init)
|
||
let version = classifierVersion
|
||
for cap in trainedCaps {
|
||
// Set current classifier
|
||
if caps[cap]?.classifierVersion == nil {
|
||
caps[cap]?.classifierVersion = version
|
||
}
|
||
}
|
||
log("Updated \(trainedCaps.count) classifier classes")
|
||
}
|
||
|
||
func save(classifier: Data, version: Int) throws {
|
||
do {
|
||
try classifier.write(to: classifierFile)
|
||
} catch {
|
||
log("Failed to write classifier \(version): \(error)")
|
||
throw Abort(.internalServerError)
|
||
}
|
||
classifierVersion = version
|
||
log("Updated classifier to version \(version)")
|
||
}
|
||
|
||
// MARK: Grid
|
||
|
||
func getListOfMissingThumbnails() -> [Int] {
|
||
caps.keys.filter { !exists(thumbnail(of: $0)) }
|
||
}
|
||
|
||
func saveThumbnail(_ data: Data, for cap: Int) {
|
||
let url = thumbnail(of: cap)
|
||
do {
|
||
try data.write(to: url)
|
||
} catch {
|
||
log("Failed to save thumbnail \(cap): \(error)")
|
||
}
|
||
}
|
||
|
||
private func updateGridCapCount() {
|
||
do {
|
||
try "const numberOfCaps = \(capCount);"
|
||
.data(using: .utf8)!
|
||
.write(to: gridCountFile)
|
||
} catch {
|
||
log("Failed to save grid cap count: \(error)")
|
||
}
|
||
}
|
||
|
||
// MARK: Monitoring
|
||
|
||
private let capCountMetric: Metric<Int>
|
||
|
||
private let imageCountMetric: Metric<Int>
|
||
|
||
private let classifierMetric: Metric<Int>
|
||
}
|