62 lines
2.1 KiB
Swift
62 lines
2.1 KiB
Swift
|
import Cocoa
|
||
|
|
import CreateML
|
||
|
|
|
||
|
|
let defaultIterations = 17
|
||
|
|
let defaultWorkingDirectory = URL(fileURLWithPath: "/Users/imac/Development/CapCollectorData")
|
||
|
|
let defaultTrainDirectory = defaultWorkingDirectory.appendingPathComponent("images")
|
||
|
|
let defaultClassifierFile = defaultWorkingDirectory.appendingPathComponent("classifier.mlmodel")
|
||
|
|
|
||
|
|
func readArguments() -> (images: URL, classifier: URL, iterations: Int) {
|
||
|
|
let count = CommandLine.argc
|
||
|
|
guard count > 1 else {
|
||
|
|
// No arguments
|
||
|
|
return (defaultTrainDirectory, defaultClassifierFile, defaultIterations)
|
||
|
|
}
|
||
|
|
// First argument is the image directory
|
||
|
|
let imageDir = URL(fileURLWithPath: CommandLine.arguments[1])
|
||
|
|
let classifier = imageDir.deletingLastPathComponent().appendingPathComponent("classifier.mlmodel")
|
||
|
|
guard count > 2 else {
|
||
|
|
// Single argument is the image directory
|
||
|
|
return (imageDir, classifier, defaultIterations)
|
||
|
|
}
|
||
|
|
// Second argument is the iteration count
|
||
|
|
guard let iterations = Int(CommandLine.arguments[2]) else {
|
||
|
|
print("[ERROR] Invalid iterations argument '\(CommandLine.arguments[2])'")
|
||
|
|
exit(-1)
|
||
|
|
}
|
||
|
|
guard count > 3 else {
|
||
|
|
return (imageDir, classifier, iterations)
|
||
|
|
}
|
||
|
|
// Third argument is the classifier path
|
||
|
|
let classifierPath = URL(fileURLWithPath: CommandLine.arguments[3])
|
||
|
|
if count > 4 {
|
||
|
|
print("[WARNING] Ignoring additional arguments")
|
||
|
|
}
|
||
|
|
return (imageDir, classifierPath, iterations)
|
||
|
|
}
|
||
|
|
|
||
|
|
let arguments = readArguments()
|
||
|
|
|
||
|
|
print("[INFO] Using images in \(arguments.images.path)")
|
||
|
|
print("[INFO] Training for \(arguments.iterations) iterations")
|
||
|
|
print("[INFO] Classifier path set to \(arguments.classifier.path)")
|
||
|
|
|
||
|
|
var params = MLImageClassifier.ModelParameters(augmentation: [])
|
||
|
|
params.maxIterations = arguments.iterations
|
||
|
|
|
||
|
|
let model = try MLImageClassifier(
|
||
|
|
trainingData: .labeledDirectories(at: arguments.images),
|
||
|
|
parameters: params)
|
||
|
|
|
||
|
|
print("[INFO] Writing classifier...")
|
||
|
|
try model.write(to: arguments.classifier)
|
||
|
|
|
||
|
|
/*
|
||
|
|
let evaluation = model.evaluation(on: .labeledDirectories(at: trainDirectory))
|
||
|
|
print("Printing evaluation:")
|
||
|
|
print(evaluation)
|
||
|
|
print("Finished evaluation")
|
||
|
|
*/
|
||
|
|
|
||
|
|
|