Create thumbnails after training classifier

This commit is contained in:
Christoph Hagen 2023-01-15 14:54:31 +01:00
parent ba40cb0a17
commit 59fa3966a4
2 changed files with 129 additions and 25 deletions

View File

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

View File

@ -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()