diff --git a/Sources/App/Cap.swift b/Sources/App/Cap.swift index 67006f6..f1e7141 100644 --- a/Sources/App/Cap.swift +++ b/Sources/App/Cap.swift @@ -10,9 +10,6 @@ struct Cap: Codable { var mainImage: Int - /// The version of the first classifier trained on this cap - var classifierVersion: Int? - var color: Color? enum CodingKeys: String, CodingKey { @@ -20,7 +17,6 @@ struct Cap: Codable { case name = "n" case count = "c" case mainImage = "m" - case classifierVersion = "v" case color = "f" } diff --git a/Sources/App/CapServer+Routes.swift b/Sources/App/CapServer+Routes.swift index 2f3ae66..20955ba 100755 --- a/Sources/App/CapServer+Routes.swift +++ b/Sources/App/CapServer+Routes.swift @@ -79,7 +79,7 @@ extension CapServer { try authenticator.authorize(request) let body = try request.getStringBody(request: "/classes/:date") - self.updateTrainedClasses(content: body) + try self.saveTrainedClasses(content: body) self.removeAllEntriesInImageChangeList(before: date) } diff --git a/Sources/App/CapServer.swift b/Sources/App/CapServer.swift index 92f2b8d..e94f883 100644 --- a/Sources/App/CapServer.swift +++ b/Sources/App/CapServer.swift @@ -23,6 +23,8 @@ final class CapServer { private let classifierFile: URL + private let classifierClassesFile: URL + private let changedImagesFile: URL private let fm = FileManager.default @@ -87,6 +89,7 @@ final class CapServer { self.htmlFile = folder.appendingPathComponent("count.html") self.classifierVersionFile = folder.appendingPathComponent("classifier.version") self.classifierFile = folder.appendingPathComponent("classifier.mlmodel") + self.classifierClassesFile = folder.appendingPathComponent("classifier.classes") self.changedImagesFile = folder.appendingPathComponent("changes.txt") self.changedImageEntryDateFormatter = DateFormatter() changedImageEntryDateFormatter.dateFormat = "yy-MM-dd-HH-mm-ss" @@ -465,7 +468,6 @@ final class CapServer { } var cap = cap cap.count = 0 - cap.classifierVersion = nil caps[cap.id] = cap saveCapCountHTML() updateGridCapCount() @@ -537,18 +539,33 @@ final class CapServer { // 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 + func saveTrainedClasses(content: String) throws { + let classes = content.components(separatedBy: ",") + + // Validate input + try classes.forEach { s in + guard let id = Int(s) else { + log("Invalid id '\(s)' in uploaded id list") + throw Abort(.badRequest) + } + guard caps[id] != nil else { + log("Unknown id '\(id)' in uploaded id list") + throw Abort(.badRequest) } } - log("Updated \(trainedCaps.count) classifier classes") + + guard let data = content.data(using: .utf8) else { + log("Failed to get classes data for writing") + throw Abort(.internalServerError) + } + + do { + try data.write(to: classifierClassesFile) + log("Updated \(classes.count) classifier classes") + } catch { + log("Failed to write classifier classes: \(error)") + throw Abort(.internalServerError) + } } func save(classifier: Data, version: Int) throws { diff --git a/Training/train.swift b/Training/train.swift index 8398ea0..8241d1e 100644 --- a/Training/train.swift +++ b/Training/train.swift @@ -156,6 +156,12 @@ final class ClassifierCreator { } let changedMainImages = changedImageList.filter { $0.image == 0 }.map { $0.cap } let classes = imageCounts.keys.sorted() + + // Delete any image folders not present as caps + guard deleteUnnecessaryImageFolders(caps: classes) else { + return nil + } + return (classes, missingImageList.count + changedImageList.count, changedMainImages) } @@ -176,6 +182,33 @@ final class ClassifierCreator { } } + private func deleteUnnecessaryImageFolders(caps: [Int]) -> Bool { + let validNames = caps.map { String(format: "%04d", $0) } + let folders: [String] + do { + folders = try FileManager.default.contentsOfDirectory(atPath: imageDirectory.path) + } catch { + print("[ERROR] Failed to get list of image folders: \(error)") + return false + } + for folder in folders { + if validNames.contains(folder) { + continue + } + + // Not a valid cap folder + let url = imageDirectory.appendingPathComponent(folder) + do { + try FileManager.default.removeItem(at: url) + print("[ERROR] Removed unused image folder '\(folder)'") + } catch { + print("[ERROR] Failed to delete unused image folder \(folder): \(error)") + return false + } + } + return true + } + private func imageUrl(base: URL, cap: Int, image: Int) -> URL { base.appendingPathComponent(String(format: "%04d/%04d-%02d.jpg", cap, cap, image)) } @@ -329,7 +362,7 @@ final class ClassifierCreator { let dateString = df.string(from: lastUpdate) return await post( url: server.appendingPathComponent("classes/\(dateString)"), - body: classes.map(String.init).joined(separator: "\n").data(using: .utf8)!) + body: classes.map(String.init).joined(separator: ",").data(using: .utf8)!) } // MARK: Step 7: Create thumbnails