<commit_before>// Copyright 2020 The SwiftFusion Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

import Datasets
import Foundation
import ModelSupport
import SwiftFusion
import TensorFlow

/// A video of bees, with tracking labels.
public struct BeeVideo {
  /// The frames of the video, as tensors of shape `[height, width, channels]`.
  public var frames: [Tensor<Double>]

  /// Labeled tracks for this video.
  public var tracks: [[BeeTrackPoint]]
}

/// The location of a bee within a frame of a `BeeVideo`.
public struct BeeTrackPoint {
  /// The frame within the video.
  public var frameIndex: Int

  /// The location of the bee.
  public var location: OrientedBoundingBox
}

extension BeeVideo {
  /// Creates an instance from the video named `videoName`.
  public init?(videoName: String) {
    let dataset = Self.downloadDatasetIfNotPresent()
    self.init(directory: dataset.appendingPathComponent(videoName))
  }

  /// Creates an instance from the data in the given `directory`.
  ///
  /// The directory must contain:
  /// - Frames named "frame0.jpeg", "frame1.jpeg", etc, consecutively.
  /// - Zero or more tracks named "track0.txt", "track1.txt", etc, consecutively.
  ///
  /// The format of a track file is:
  /// - A line "<height> <width>" specifying the size of the bounding boxes in the track.
  /// - Followed by arbitrarily many lines "<frame index> <x> <y> <theta>".
  public init?(directory: URL) {
    self.frames = []
    self.tracks = []

    func loadFrame(_ index: Int) -> Tensor<Double>? {
      let path = directory.appendingPathComponent("frame\(index).jpeg")
      guard FileManager.default.fileExists(atPath: path.path) else { return nil }
      return Tensor<Double>(Image(jpeg: path).tensor)
    }
    while let frame = loadFrame(self.frames.count) {
      self.frames.append(frame)
    }

    func loadTrack(_ index: Int) -> [BeeTrackPoint]? {
      let path = directory.appendingPathComponent("track\(index).txt")
      guard let track = try? String(contentsOf: path) else { return nil }
      let lines = track.split(separator: "\n")
      let bbSize = lines.first!.split(separator: " ")
      let bbHeight = Int(bbSize[0])!
      let bbWidth = Int(bbSize[1])!
      return lines.dropFirst().map { line in
        let split = line.split(separator: " ")
        return BeeTrackPoint(
          frameIndex: Int(split[0])!,
          location: OrientedBoundingBox(
            center: Pose2(
              Rot2(Double(split[3])!),
              Vector2(Double(split[1])!, Double(split[2])!)),
            rows: bbHeight,
            cols: bbWidth))
      }
    }
    while let track = loadTrack(self.tracks.count) {
      self.tracks.append(track)
    }
  }

  private static func downloadDatasetIfNotPresent() -> URL {
    let downloadDir = DatasetUtilities.defaultDirectory.appendingPathComponent(
      "bee_videos", isDirectory: true)
    let directoryExists = FileManager.default.fileExists(atPath: downloadDir.path)
    let contentsOfDir = try? FileManager.default.contentsOfDirectory(atPath: downloadDir.path)
    let directoryEmpty = (contentsOfDir == nil) || (contentsOfDir!.isEmpty)

    guard !directoryExists || directoryEmpty else { return downloadDir }

    let remoteRoot = URL(
      string: "https://storage.googleapis.com/swift-tensorflow-misc-files/beetracking")!

    let _ = DatasetUtilities.downloadResource(
      filename: "bee_videos", fileExtension: "tar.gz",
      remoteRoot: remoteRoot, localStorageDirectory: downloadDir
    )<commit_msg>Add defered video loading to avoid S4TF bug<commit_after>// Copyright 2020 The SwiftFusion Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

import Datasets
import Foundation
import ModelSupport
import SwiftFusion
import TensorFlow

/// A video of bees, with tracking labels.
public struct BeeVideo {
  /// The frames of the video, as tensors of shape `[height, width, channels]`.
  public var frames: [Tensor<Double>]

  /// Labeled tracks for this video.
  public var tracks: [[BeeTrackPoint]]

  var directory: URL?
}

/// The location of a bee within a frame of a `BeeVideo`.
public struct BeeTrackPoint {
  /// The frame within the video.
  public var frameIndex: Int

  /// The location of the bee.
  public var location: OrientedBoundingBox
}

extension BeeVideo {
  /// Creates an instance from the video named `videoName`.
  public init?(videoName: String, deferLoadingFrames: Bool = false) {
    let dataset = Self.downloadDatasetIfNotPresent()
    self.init(directory: dataset.appendingPathComponent(videoName), deferLoadingFrames: deferLoadingFrames)
  }

  /// Creates an instance from the data in the given `directory`.
  ///
  /// The directory must contain:
  /// - Frames named "frame0.jpeg", "frame1.jpeg", etc, consecutively.
  /// - Zero or more tracks named "track0.txt", "track1.txt", etc, consecutively.
  ///
  /// The format of a track file is:
  /// - A line "<height> <width>" specifying the size of the bounding boxes in the track.
  /// - Followed by arbitrarily many lines "<frame index> <x> <y> <theta>".
  public init?(directory: URL, deferLoadingFrames: Bool = false) {
    self.frames = []
    self.tracks = []
    self.directory = directory

    while let frame = loadFrame(self.frames.count) {
      self.frames.append(frame)
    }

    func loadTrack(_ index: Int) -> [BeeTrackPoint]? {
      let path = directory.appendingPathComponent("track\(index).txt")
      guard let track = try? String(contentsOf: path) else { return nil }
      let lines = track.split(separator: "\n")
      let bbSize = lines.first!.split(separator: " ")
      let bbHeight = Int(bbSize[0])!
      let bbWidth = Int(bbSize[1])!
      return lines.dropFirst().map { line in
        let split = line.split(separator: " ")
        return BeeTrackPoint(
          frameIndex: Int(split[0])!,
          location: OrientedBoundingBox(
            center: Pose2(
              Rot2(Double(split[3])!),
              Vector2(Double(split[1])!, Double(split[2])!)),
            rows: bbHeight,
            cols: bbWidth))
      }
    }
    while let track = loadTrack(self.tracks.count) {
      self.tracks.append(track)
    }
  }

  func loadFrame(_ index: Int) -> Tensor<Double>? {
    let path = self.directory!.appendingPathComponent("frame\(index).jpeg")
    guard FileManager.default.fileExists(atPath: path.path) else { return nil }
    return Tensor<Double>(Image(jpeg: path).tensor)
  }
  
  private static func downloadDatasetIfNotPresent() -> URL {
    let downloadDir = DatasetUtilities.defaultDirectory.appendingPathComponent(
      "bee_videos", isDirectory: true)
    let directoryExists = FileManager.default.fileExists(atPath: downloadDir.path)
    let contentsOfDir = try? FileManager.default.contentsOfDirectory(atPath: downloadDir.path)
    let directoryEmpty = (contentsOfDir == nil) || (contentsOfDir!.isEmpty)

    guard !directoryExists || directoryEmpty else { return downloadDir }

    let remoteRoot = URL(
      string: "https://storage.googleapis.com/swift-tensorflow-misc-files/beetracking")!

    let _ = DatasetUtilities.downloadResource(
      filename: "bee_videos", fileExtension: "tar.gz",
      remoteRoot: remoteRoot, localStorageDirectory: downloadDir
    )
{'fraction_non_alphanumeric': 0.11313472439995176, 'fraction_numerical': 0.005186346640935955, 'mean_word_length': 3.2262996941896023, 'pattern_counts': {'":': 0, '<': 21, '<?xml version=': 0, '>': 27, 'https://': 2, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 16, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1535542', 'n_tokens_mistral': 2404, 'n_tokens_neox': 2295, 'n_words': 764}