Migrating OxfordIIITPets dataset to Epochs (#604)

* Initial conversion of the OxfordIIITPets dataset.

* Formatting.

* Adding convenience initializers.

* Documentation fix.

Author: BradLarson

Datasets/ImageSegmentationDataset.swift

@@ -12,13 +12,32 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-import Batcher
+import ModelSupport
 import TensorFlow
 
-public protocol ImageSegmentationDataset {
-    associatedtype SourceDataSet: Collection
-    where SourceDataSet.Element == TensorPair<Float, Int32>, SourceDataSet.Index == Int
-    init(batchSize: Int)
-    var training: Batcher<SourceDataSet> { get }
-    var test: Batcher<SourceDataSet> { get }
+/// An image with a label.
+public typealias SegmentedImage = LabeledData<Tensor<Float>, Tensor<Int32>>
+
+/// Types whose elements represent an image segmentation dataset (with both
+/// training and validation data).
+public protocol ImageSegmentationData {
+  /// The type of the training data, represented as a sequence of epochs, which
+  /// are collection of batches.
+  associatedtype Training: Sequence
+  where Training.Element: Collection, Training.Element.Element == SegmentedImage
+  /// The type of the validation data, represented as a collection of batches.
+  associatedtype Validation: Collection where Validation.Element == SegmentedImage
+  /// Creates an instance from a given `batchSize`.
+  init(batchSize: Int, on device: Device)
+  /// The `training` epochs.
+  var training: Training { get }
+  /// The `validation` batches.
+  var validation: Validation { get }
+
+  // The following is probably going to be necessary since we can't extract that
+  // information from `Epochs` or `Batches`.
+  /// The number of samples in the `training` set.
+  //var trainingSampleCount: Int {get}
+  /// The number of samples in the `validation` set.
+  //var validationSampleCount: Int {get}
 }

Datasets/OxfordIIITPets/OxfordIIITPets.swift

@@ -17,48 +17,98 @@
 // Omkar M Parkhi, Andrea Vedaldi, Andrew Zisserman and C. V. Jawahar
 // https://www.robots.ox.ac.uk/~vgg/data/pets/
 
-import Batcher
 import Foundation
 import ModelSupport
 import TensorFlow
 
-public typealias LazyDataSet = LazyMapSequence<[URL], TensorPair<Float, Int32>>
-
-public struct OxfordIIITPets: ImageSegmentationDataset {
-  public typealias SourceDataSet = LazyDataSet
-  public let training: Batcher<SourceDataSet>
-  public let test: Batcher<SourceDataSet>
-
-  public init(batchSize: Int) {
-    self.init(batchSize: batchSize, imageSize: 224)
+public struct OxfordIIITPets<Entropy: RandomNumberGenerator> {
+  /// Type of the collection of non-collated batches.
+  public typealias Batches = Slices<Sampling<[(file: URL, annotation: URL)], ArraySlice<Int>>>
+  /// The type of the training data, represented as a sequence of epochs, which
+  /// are collection of batches.
+  public typealias Training = LazyMapSequence<
+    TrainingEpochs<[(file: URL, annotation: URL)], Entropy>,
+    LazyMapSequence<Batches, SegmentedImage>
+  >
+  /// The type of the validation data, represented as a collection of batches.
+  public typealias Validation = LazyMapSequence<
+    Slices<[(file: URL, annotation: URL)]>, LabeledImage
+  >
+  /// The training epochs.
+  public let training: Training
+  /// The validation batches.
+  public let validation: Validation
+
+  /// Creates an instance with `batchSize`.
+  ///
+  /// - Parameters:
+  ///   - batchSize: Number of images provided per batch.
+  ///   - entropy: A source of randomness used to shuffle sample
+  ///     ordering.  It  will be stored in `self`, so if it is only pseudorandom
+  ///     and has value semantics, the sequence of epochs is deterministic and not
+  ///     dependent on other operations.
+  ///   - device: The Device on which resulting Tensors from this dataset will be placed, as well
+  ///     as where the latter stages of any conversion calculations will be performed.
+  public init(batchSize: Int, entropy: Entropy, device: Device) {
+    self.init(
+      batchSize: batchSize, entropy: entropy, device: device, imageSize: 224)
   }
 
+  /// Creates an instance with `batchSize` on `device` using `remoteBinaryArchiveLocation`.
+  ///
+  /// - Parameters:
+  ///   - batchSize: Number of images provided per batch.
+  ///   - entropy: A source of randomness used to shuffle sample ordering.  It
+  ///     will be stored in `self`, so if it is only pseudorandom and has value
+  ///     semantics, the sequence of epochs is deterministic and not dependent
+  ///     on other operations.
+  ///   - device: The Device on which resulting Tensors from this dataset will be placed, as well
+  ///     as where the latter stages of any conversion calculations will be performed.
+  ///   - imageSize: The square width and height of the images returned from this dataset.
+  ///   - localStorageDirectory: Where to place the downloaded and unarchived dataset.
   public init(
-    batchSize: Int,
+    batchSize: Int, entropy: Entropy, device: Device, imageSize: Int,
     localStorageDirectory: URL = DatasetUtilities.defaultDirectory
-      .appendingPathComponent("OxfordIIITPets", isDirectory: true),
-    imageSize: Int
+      .appendingPathComponent("OxfordIIITPets", isDirectory: true)
   ) {
     do {
-      training = Batcher<SourceDataSet>(
-        on: try loadOxfordIITPetsTraining(
-          imageSize: imageSize,
-          localStorageDirectory: localStorageDirectory
-        ),
-        batchSize: batchSize,
-        shuffle: true)
-      test = Batcher<SourceDataSet>(
-        on: try loadOxfordIIITPetsValidation(
-          imageSize: imageSize,
-          localStorageDirectory: localStorageDirectory
-        ),
-        batchSize: batchSize)
+      let trainingSamples = try loadOxfordIITPetsTraining(
+        localStorageDirectory: localStorageDirectory)
+
+      training = TrainingEpochs(samples: trainingSamples, batchSize: batchSize, entropy: entropy)
+        .lazy.map { (batches: Batches) -> LazyMapSequence<Batches, LabeledImage> in
+          return batches.lazy.map {
+            makeBatch(samples: $0, imageSize: imageSize, device: device)
+          }
+        }
+
+      let validationSamples = try loadOxfordIITPetsTraining(
+        localStorageDirectory: localStorageDirectory)
+
+      validation = validationSamples.inBatches(of: batchSize).lazy.map {
+        makeBatch(samples: $0, imageSize: imageSize, device: device)
+      }
     } catch {
-      fatalError("Could not load Oxford IIIT Pets dataset: \(error)")
+      fatalError("Could not load the Oxford IIIT Pets dataset: \(error)")
     }
   }
 }
 
+extension OxfordIIITPets: ImageSegmentationData where Entropy == SystemRandomNumberGenerator {
+  /// Creates an instance with `batchSize`, using the SystemRandomNumberGenerator.
+  public init(batchSize: Int, on device: Device = Device.default) {
+    self.init(batchSize: batchSize, entropy: SystemRandomNumberGenerator(), device: device)
+  }
+
+  /// Creates an instance with `batchSize`, `inputSize`, and `outputSize`, using the
+  /// SystemRandomNumberGenerator.
+  public init(batchSize: Int, imageSize: Int, on device: Device = Device.default) {
+    self.init(
+      batchSize: batchSize, entropy: SystemRandomNumberGenerator(), device: device,
+      imageSize: imageSize)
+  }
+}
+
 func downloadOxfordIIITPetsIfNotPresent(to directory: URL) {
   let downloadPath = directory.appendingPathComponent("images", isDirectory: true).path
   let directoryExists = FileManager.default.fileExists(atPath: downloadPath)
@@ -80,21 +130,13 @@ func downloadOxfordIIITPetsIfNotPresent(to directory: URL) {
   )
 }
 
-func loadOxfordIIITPets(filename: String, in directory: URL, imageSize: Int) throws -> LazyDataSet {
+func loadOxfordIIITPets(filename: String, in directory: URL) throws -> [(
+  file: URL, annotation: URL
+)] {
   downloadOxfordIIITPetsIfNotPresent(to: directory)
   let imageURLs = getImageURLs(filename: filename, directory: directory)
-  return imageURLs.lazy.map { (imageURL: URL) -> TensorPair<Float, Int32> in
-    TensorPair<Float, Int32>(
-      first:
-        Image(jpeg: imageURL).resized(to: (imageSize, imageSize)).tensor[0..., 0..., 0..<3]
-        / 255.0,
-      second: Tensor<Int32>(
-        Image(jpeg: makeAnnotationURL(imageURL: imageURL, directory: directory)).resized(
-          to: (imageSize, imageSize)
-        ).tensor[0..., 0..., 0...0] - 1
-      )
-    )
-
+  return imageURLs.lazy.map { (imageURL: URL) -> (file: URL, annotation: URL) in
+    (file: imageURL, annotation: makeAnnotationURL(imageURL: imageURL, directory: directory))
   }
 }
 
@@ -114,20 +156,36 @@ func getImageURLs(filename: String, directory: URL) -> [URL] {
   }
 }
 
-func loadOxfordIITPetsTraining(
-  imageSize: Int, localStorageDirectory: URL
-) throws
-  -> LazyDataSet
+func loadOxfordIITPetsTraining(localStorageDirectory: URL) throws -> [(file: URL, annotation: URL)]
 {
   return try loadOxfordIIITPets(
-    filename: "trainval.txt", in: localStorageDirectory, imageSize: imageSize)
+    filename: "trainval.txt", in: localStorageDirectory)
 }
 
-func loadOxfordIIITPetsValidation(
-  imageSize: Int, localStorageDirectory: URL
-) throws
-  -> LazyDataSet
-{
+func loadOxfordIIITPetsValidation(localStorageDirectory: URL) throws -> [(
+  file: URL, annotation: URL
+)] {
   return try loadOxfordIIITPets(
-    filename: "test.txt", in: localStorageDirectory, imageSize: imageSize)
+    filename: "test.txt", in: localStorageDirectory)
+}
+
+fileprivate func makeBatch<BatchSamples: Collection>(
+  samples: BatchSamples, imageSize: Int, device: Device
+) -> SegmentedImage where BatchSamples.Element == (file: URL, annotation: URL) {
+  let images = samples.map(\.file).map { url -> Tensor<Float> in
+    Image(jpeg: url).resized(to: (imageSize, imageSize)).tensor[0..., 0..., 0..<3]
+  }
+
+  var imageTensor = Tensor(stacking: images)
+  imageTensor = Tensor(copying: imageTensor, to: device)
+  imageTensor /= 255.0
+
+  let annotations = samples.map(\.annotation).map { url -> Tensor<Int32> in
+    Tensor<Int32>(
+      Image(jpeg: url).resized(to: (imageSize, imageSize)).tensor[0..., 0..., 0...0] - 1)
+  }
+  var annotationTensor = Tensor(stacking: annotations)
+  annotationTensor = Tensor(copying: annotationTensor, to: device)
+
+  return SegmentedImage(data: imageTensor, label: annotationTensor)
 }

Tests/DatasetsTests/OxfordIIITPets/OxfordIIITPetsTests.swift

@@ -7,20 +7,22 @@ final class OxfordIIITPetsTests: XCTestCase {
     let dataset = OxfordIIITPets(batchSize: 64)
 
     var batchCount = 0
-    for batch in dataset.training.sequenced() {
-      batchCount += 1
-      /// 3680 samples make 57 batches of size 64 and one batch of size 32
-      let expectedBS = batchCount <= 57 ? 64 : 32
+    for epochBatches in dataset.training.prefix(1) {
+      for batch in epochBatches {
+        batchCount += 1
+        /// 3680 samples make 57 batches of size 64 and one batch of size 32
+        let expectedBS = batchCount <= 57 ? 64 : 32
 
-      XCTAssertEqual(batch.first.shape, [expectedBS, 224, 224, 3])
-      XCTAssertEqual(batch.second.shape, [expectedBS, 224, 224, 1])
+        XCTAssertEqual(batch.data.shape, [expectedBS, 224, 224, 3])
+        XCTAssertEqual(batch.label.shape, [expectedBS, 224, 224, 1])
+      }
     }
-    XCTAssertEqual(batchCount, dataset.training.count)
+    XCTAssertEqual(batchCount, 57)
   }
 }
 
 extension OxfordIIITPetsTests {
   static var allTests = [
-    ("testCreateOxfordIIITPets", testCreateOxfordIIITPets),
+    ("testCreateOxfordIIITPets", testCreateOxfordIIITPets)
   ]
 }