Cyclegan tweaks (#441)

* make cyclegan demo standalone

* add readme

* formatting, use let

Author: brettkoonce

CycleGAN/CLI.swift

@@ -15,7 +15,7 @@
 import ArgumentParser
 
 struct Options: ParsableArguments {
-    @Option(default: "./dataset", help: ArgumentHelp("Path to the dataset folder", valueName: "dataset-path"))
+    @Option(default: "", help: ArgumentHelp("Path to the dataset folder", valueName: "dataset-path"))
     var datasetPath: String
 
     @Option(default: 50, help: ArgumentHelp("Number of epochs", valueName: "epochs"))

CycleGAN/README.md

@@ -0,0 +1,18 @@
+# CycleGAN
+
+**Unpaired Image-to-Image Translation using Cycle-Consistent Adversarial Networks**
+website: https://junyanz.github.io/CycleGAN/
+
+arXiv: https://arxiv.org/abs/1703.10593
+
+## Setup
+
+To begin, you'll need the [latest version of Swift for
+TensorFlow](https://github.com/tensorflow/swift/blob/master/Installation.md)
+installed. Make sure you've added the correct version of `swift` to your path.
+
+To train the model, run:
+
+```sh
+swift run CycleGAN
+```

CycleGAN/main.swift

@@ -15,14 +15,45 @@
 import Foundation
 import ModelSupport
 import TensorFlow
+import Datasets
 
 let options = Options.parseOrExit()
 
-let datasetFolder = URL(fileURLWithPath: options.datasetPath, isDirectory: true)
-let trainFolderA = datasetFolder.appendingPathComponent("trainA")
-let trainFolderB = datasetFolder.appendingPathComponent("trainB")
-let testFolderA = datasetFolder.appendingPathComponent("testA")
-let testFolderB = datasetFolder.appendingPathComponent("testB")
+let datasetFolder: URL
+let trainFolderA: URL
+let trainFolderB: URL
+let testFolderA: URL
+let testFolderB: URL
+
+if options.datasetPath.length != 0 {
+    datasetFolder = URL(fileURLWithPath: options.datasetPath, isDirectory: true)
+    trainFolderA = datasetFolder.appendingPathComponent("trainA")
+    trainFolderB = datasetFolder.appendingPathComponent("trainB")
+    testFolderA = datasetFolder.appendingPathComponent("testA")
+    testFolderB = datasetFolder.appendingPathComponent("testB")
+} else {
+    func downloadZebraDataSetIfNotPresent(to directory: URL) {
+        let downloadPath = directory.appendingPathComponent("horse2zebra").path
+        let directoryExists = FileManager.default.fileExists(atPath: downloadPath)
+        let contentsOfDir = try? FileManager.default.contentsOfDirectory(atPath: downloadPath)
+        let directoryEmpty = (contentsOfDir == nil) || (contentsOfDir!.isEmpty)
+
+        guard !directoryExists || directoryEmpty else { return }
+
+        let location = URL(
+            string: "https://people.eecs.berkeley.edu/~taesung_park/CycleGAN/datasets/horse2zebra.zip")!
+        let _ = DatasetUtilities.downloadResource(
+            filename: "horse2zebra", fileExtension: "zip",
+            remoteRoot: location.deletingLastPathComponent(), localStorageDirectory: directory)
+    }
+
+    datasetFolder = DatasetUtilities.defaultDirectory.appendingPathComponent("CycleGAN", isDirectory: true)
+    downloadZebraDataSetIfNotPresent(to: datasetFolder)
+    trainFolderA = datasetFolder.appendingPathComponent("horse2zebra/trainA")
+    trainFolderB = datasetFolder.appendingPathComponent("horse2zebra/trainB")
+    testFolderA = datasetFolder.appendingPathComponent("horse2zebra/testA")
+    testFolderB = datasetFolder.appendingPathComponent("horse2zebra/testB")
+}
 
 let trainDatasetA = try Images(folderURL: trainFolderA)
 let trainDatasetB = try Images(folderURL: trainFolderB)

Package.swift

@@ -121,7 +121,7 @@ let package = Package(
         .testTarget(name: "SupportTests", dependencies: ["ModelSupport"]),
         .target(
             name: "CycleGAN",
-            dependencies: ["Batcher", .product(name: "ArgumentParser", package: "swift-argument-parser"), "ModelSupport"],
+            dependencies: ["Batcher", .product(name: "ArgumentParser", package: "swift-argument-parser"), "ModelSupport", "Datasets"],
             path: "CycleGAN"
         )
     ]