Create thumbnails after training classifier

Training/config_example.json

@@ -1,6 +1,5 @@
 {
-    "imageDirectory": "../Public/images",
+    "contentDirectory": "../Public",
-    "classifierModelPath": "../Public/classifier.mlmodel",
     "trainingIterations": 20,
     "serverPath": "https://mydomain.com/caps",
     "authenticationToken": "mysecretkey",

Training/train.swift

@@ -4,9 +4,7 @@ import CreateML
 
 struct Configuration: Codable {
     
-    let imageDirectory: String
+    let contentFolder: String
-    
-    let classifierModelPath: 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)
+        let contentDirectory = URL(fileURLWithPath: configuration.contentFolder)
-        self.classifierUrl = URL(fileURLWithPath: configuration.classifierModelPath)
+        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 {
+        createFolderIfMissing(imageDirectory)
-            return true
-        }
-        do {
-            try FileManager.default.createDirectory(at: imageDirectory, withIntermediateDirectories: true)
-            return true
-        } catch {
-            print("[ERROR] Failed to create image directory: \(error)")
-            return false
-        }
     }
     
     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)!,
+            body: classes.map(String.init).joined(separator: "\n").data(using: .utf8)!)
-            dateKey: lastUpdate)
+    }
+    
+    // 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()