Caps-Train/Sources/CapTrain.swift

67 lines
2.8 KiB
Swift
Raw Normal View History

2023-10-23 12:28:35 +02:00
import ArgumentParser
import Foundation
@main
struct CapTrain: AsyncParsableCommand {
2023-10-24 11:31:08 +02:00
private static let defaultIterations = 10
2024-10-28 19:06:45 +01:00
private static let defaultMinimumImageCount = 10
2023-10-24 11:16:22 +02:00
@Flag(name: .shortAndLong, help: "Resume the previous training session (default: false)")
var resume: Bool = false
2023-10-24 11:31:08 +02:00
@Option(name: .shortAndLong, help: "The path to the configuration file. The file must be a json object containing command line arguments. Command line options take precedence over configuration file options")
var configuration: String?
2023-10-23 14:58:58 +02:00
2023-10-24 11:31:08 +02:00
@Option(name: .shortAndLong, help: "The number of iterations to train (default: 10)")
2023-10-23 14:58:58 +02:00
var iterations: Int?
2024-10-28 19:06:45 +01:00
@Option(name: .shortAndLong, help: "The minimum number of images for a cap to be included in training (default: 10)")
var minimumImageCount: Int?
2023-10-24 11:31:08 +02:00
@Option(name: .shortAndLong, help: "The url of the caps server to retrieve images and upload the classifier")
var server: String?
2023-10-23 14:58:58 +02:00
@Option(name: .shortAndLong, help: "The authentication token for the server")
var authentication: String?
2023-10-24 11:31:08 +02:00
@Option(name: .shortAndLong, help: "The path to the folder where the content (images, classifier, thumbnails) is stored")
2023-10-23 14:58:58 +02:00
var folder: String?
2024-06-12 16:13:42 +02:00
@Flag(name: .long, help: "Only upload the previous classifier")
var skipTraining: Bool = false
2023-10-23 12:28:35 +02:00
func run() async throws {
2023-10-23 14:58:58 +02:00
let configurationFile = try configurationFile()
2023-10-24 11:31:08 +02:00
let iterations = iterations ?? configurationFile?.iterations ?? CapTrain.defaultIterations
2024-10-28 19:06:45 +01:00
let minimumCount = minimumImageCount ?? configurationFile?.minimumImagesPerCap ?? CapTrain.defaultMinimumImageCount
2023-10-24 11:31:08 +02:00
guard let contentFolder = folder ?? configurationFile?.folder else {
throw TrainingError.missingArguments("folder")
}
guard let serverPath = server ?? configurationFile?.server else {
throw TrainingError.missingArguments("server")
}
guard let authentication = authentication ?? configurationFile?.authentication else {
throw TrainingError.missingArguments("authentication")
2023-10-23 14:58:58 +02:00
}
let configuration = Configuration(
contentFolder: contentFolder,
2023-10-24 11:31:08 +02:00
trainingIterations: iterations,
2023-10-23 14:58:58 +02:00
serverPath: serverPath,
2023-10-23 12:28:35 +02:00
2024-10-28 19:06:45 +01:00
authenticationToken: authentication,
minimumImagesPerCap: minimumCount)
2023-10-24 11:16:22 +02:00
let creator = try ClassifierCreator(configuration: configuration, resume: resume)
2024-06-12 16:13:42 +02:00
try await creator.run(skipTraining: skipTraining)
2023-10-23 12:28:35 +02:00
}
2023-10-23 14:58:58 +02:00
private func configurationFile() throws -> ConfigurationFile? {
2023-10-24 11:31:08 +02:00
guard let configuration else {
2023-10-23 14:58:58 +02:00
return nil
}
2023-10-24 11:31:08 +02:00
let configurationFileUrl = URL(fileURLWithPath: configuration)
2023-10-23 14:58:58 +02:00
return try ConfigurationFile(at: configurationFileUrl)
}
2023-10-23 12:28:35 +02:00
}