diff --git a/Training/config_example.json b/Training/config_example.json index 3a401ec..af6a6fe 100644 --- a/Training/config_example.json +++ b/Training/config_example.json @@ -1,6 +1,5 @@ { - "imageDirectory": "../Public/images", - "classifierModelPath": "../Public/classifier.mlmodel", + "contentDirectory": "../Public", "trainingIterations": 20, "serverPath": "https://mydomain.com/caps", "authenticationToken": "mysecretkey", diff --git a/Training/train.swift b/Training/train.swift index f8c966a..37b114d 100644 --- a/Training/train.swift +++ b/Training/train.swift @@ -4,9 +4,7 @@ import CreateML struct Configuration: Codable { - let imageDirectory: String - - let classifierModelPath: String + let contentFolder: String let trainingIterations: Int @@ -47,6 +45,8 @@ final class ClassifierCreator { let imageDirectory: URL + let thumbnailDirectory: URL + let classifierUrl: URL let df = DateFormatter() @@ -64,8 +64,10 @@ final class ClassifierCreator { return nil } self.server = serverUrl - self.imageDirectory = URL(fileURLWithPath: configuration.imageDirectory) - self.classifierUrl = URL(fileURLWithPath: configuration.classifierModelPath) + let contentDirectory = URL(fileURLWithPath: configuration.contentFolder) + self.imageDirectory = contentDirectory.appendingPathComponent("images") + self.classifierUrl = contentDirectory.appendingPathComponent("classifier.mlmodel") + self.thumbnailDirectory = contentDirectory.appendingPathComponent("thumbnails") df.dateFormat = "yy-MM-dd-HH-mm-ss" } @@ -73,7 +75,7 @@ final class ClassifierCreator { func run() async { let imagesSnapshotDate = Date() - guard let (classes, changedImageCount) = await loadImages() else { + guard let (classes, changedImageCount, changedMainImages) = await loadImages() else { return } @@ -83,6 +85,8 @@ final class ClassifierCreator { guard changedImageCount > 0 else { print("[INFO] No changed images, so no new classifier trained") + await createThumbnails(changed: changedMainImages) + print("[INFO] Done") return } @@ -92,7 +96,7 @@ final class ClassifierCreator { let newVersion = classifierVersion + 1 print("[INFO] Image directory: \(imageDirectory.absoluteURL.path)") - print("[INFO] Model path: \(configuration.classifierModelPath)") + print("[INFO] Model path: \(classifierUrl.path)") print("[INFO] Version: \(newVersion)") print("[INFO] Classes: \(classes.count)") print("[INFO] Iterations: \(configuration.trainingIterations)") @@ -108,13 +112,13 @@ final class ClassifierCreator { guard await upload(classes: classes, lastUpdate: imagesSnapshotDate) else { return } - + await createThumbnails(changed: changedMainImages) print("[INFO] Done") } // MARK: Step 2: Load changed images - func loadImages() async -> (classes: [Int], changedImages: Int)? { + func loadImages() async -> (classes: [Int], changedImageCount: Int, changedMainImages: [Int])? { guard createImageFolderIfMissing() else { return nil } @@ -149,21 +153,13 @@ final class ClassifierCreator { guard await loadImages(changedImageList) else { return nil } + let changedMainImages = changedImageList.filter { $0.image == 0 }.map { $0.cap } let classes = imageCounts.keys.sorted() - return (classes, missingImageList.count + changedImageList.count) + return (classes, missingImageList.count + changedImageList.count, changedMainImages) } private func createImageFolderIfMissing() -> Bool { - guard !FileManager.default.fileExists(atPath: imageDirectory.path) else { - return true - } - do { - try FileManager.default.createDirectory(at: imageDirectory, withIntermediateDirectories: true) - return true - } catch { - print("[ERROR] Failed to create image directory: \(error)") - return false - } + createFolderIfMissing(imageDirectory) } private func getImageCounts() async -> [Int : Int] { @@ -329,8 +325,117 @@ 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)!, - dateKey: lastUpdate) + body: classes.map(String.init).joined(separator: "\n").data(using: .utf8)!) + } + + // MARK: Step 7: Create thumbnails + + func createThumbnails(changed: [Int]) async { + guard checkMagickAvailability() else { + return + } + guard createFolderIfMissing(thumbnailDirectory) else { + print("[ERROR] Failed to create folder for thumbnails") + return + } + let capIdsOfMissingThumbnails = await getMissingThumbnailIds() + let all = Set(capIdsOfMissingThumbnails).union(changed) + print("[INFO] Creating \(all.count) thumbnails...") + for cap in all { + await createThumbnail(for: cap) + } + } + + func checkMagickAvailability() -> Bool { + do { + let (code, output) = try safeShell("magick --version") + guard code == 0, let version = output.components(separatedBy: "ImageMagick ").dropFirst().first? + .components(separatedBy: " ").first else { + print("[ERROR] Magick not found, install using 'brew install imagemagick'") + return false + } + print("[INFO] Using magick \(version)") + } catch { + print("[ERROR] Failed to get version of magick: (\(error))") + return false + } + return true + } + + private func getMissingThumbnailIds() async -> [Int] { + guard let string: String = await get(server.appendingPathComponent("thumbnails/missing")) else { + print("[ERROR] Failed to get missing thumbnails") + return [] + } + return string.components(separatedBy: ",").compactMap(Int.init) + } + + private func createThumbnail(for cap: Int) async { + let inputUrl = imageUrl(base: imageDirectory, cap: cap, image: 0) + guard FileManager.default.fileExists(atPath: inputUrl.path) else { + print("[ERROR] Local main image not found for cap \(cap): \(inputUrl.path)") + return + } + + let output = thumbnailDirectory.appendingPathComponent(String(format: "%04d.jpg", cap)) + do { + let command = "magick convert \(inputUrl.path) -quality 70% -resize 100x100 \(output.path)" + let (code, output) = try safeShell(command) + if code != 0 { + print("Failed to create thumbnail for cap \(cap): \(output)") + return + } + } catch { + print("Failed to read created thumbnail for cap \(cap): \(error)") + return + } + + let data: Data + do { + data = try Data(contentsOf: output) + } catch { + print("Failed to read created thumbnail for cap \(cap): \(error)") + return + } + guard await post(url: server.appendingPathComponent("thumbnails/\(cap)"), body: data) else { + print("Failed to upload thumbnail for cap \(cap)") + return + } + } + + // MARK: Helper + + @discardableResult + private func safeShell(_ command: String) throws -> (code: Int32, output: String) { + let task = Process() + let pipe = Pipe() + + task.standardOutput = pipe + task.standardError = pipe + task.arguments = ["-cl", command] + task.executableURL = URL(fileURLWithPath: "/bin/zsh") + task.standardInput = nil + + try task.run() + task.waitUntilExit() + + let data = pipe.fileHandleForReading.readDataToEndOfFile() + let output = String(data: data, encoding: .utf8)! + return (task.terminationStatus, output) + } + + + private func createFolderIfMissing(_ folder: URL) -> Bool { + guard !FileManager.default.fileExists(atPath: folder.path) else { + return true + } + do { + try FileManager.default.createDirectory(at: folder, withIntermediateDirectories: true) + return true + } catch { + print("[ERROR] Failed to create directory \(folder.path): \(error)") + return false + } } // MARK: Requests @@ -376,4 +481,4 @@ final class ClassifierCreator { } } -await ClassifierCreator().run() +await ClassifierCreator()?.run()