Data block foundations¶

// export
import Path
import TensorFlow
import Python

let plt = Python.import("matplotlib.pyplot")

//export
let dataPath = Path.home/".fastai"/"data"

//export
func downloadImagette(path: Path = dataPath) -> Path {
    let url = "https://s3.amazonaws.com/fast-ai-imageclas/imagenette-160.tgz"
    let fname = "imagenette-160"
    let file = path/fname
    try! path.mkdir(.p)
    if !file.exists {
        downloadFile(url, dest:(path/"\(fname).tgz").string)
        _ = shellCommand("/bin/tar", ["-xzf", (path/"\(fname).tgz").string, "-C", path.string])
    }
    return file
}

let path = downloadImagette()

let pathTench = path/"val"/"n01440764"

let imgBytes = Raw.readFile(filename: StringTensor(imgFn.string))
let decodedImg = Raw.decodeJpeg(contents: imgBytes, channels: 3, dctMethod: "") 

let np = Python.import("numpy")

//export
public func fetchFiles(path: Path, recurse: Bool = false, extensions: [String]? = nil) -> [Path]{
    var res: [Path] = []
    for p in try! path.ls(){
        if p.kind == .directory && recurse { 
            res += fetchFiles(path: p.path, recurse: recurse, extensions: extensions)
        } else if extensions == nil || extensions!.contains(p.path.extension.lowercased) {
            res.append(p.path)
        }
    }
    return res
}

let fNames = fetchFiles(path: path, recurse: true, extensions: ["jpeg", "jpg"])

// export
public struct SplitData<Item>{
    public let train: [Item]
    public let valid: [Item]
    
    public init(train: [Item], valid: [Item]){
        (self.train, self.valid) = (train, valid)
    }
    
    public init(items: [Item], fromFunc: (Item) -> Bool){
        var (trn, val): ([Item], [Item]) = ([], [])
        for x in items {
            if fromFunc(x) { val.append(x) }
            else           { trn.append(x) }
        }
        (self.train, self.valid) = (trn, val)
    }
}

// export
public func grandParentSplitter(fName: Path, valid: String = "valid") -> Bool{
    return fName.parent.parent.basename() == valid
}

let sd = SplitData(items: fNames, fromFunc: {grandParentSplitter(fName: $0, valid: "val")})

// export
protocol Processor {
    associatedtype Input
    associatedtype Output
    
    mutating func initState(items: [Input])
    func process1(item: Input) -> Output
    func deprocess1(item: Output) -> Input
}

// export
extension Processor {
    func process(items: [Input]) -> [Output]{
        return items.map(){process1(item: $0)}
    }
    
    func deprocess(items: [Output]) -> [Input]{
        return items.map(){deprocess1(item: $0)}
    }
}

// export
public struct CategoryProcessor: Processor {
    public typealias Input = String
    public typealias Output = Int32
    public var vocab: [Input]? = nil
    public var reverseMap: [Input: Output]? = nil
    
    public mutating func initState(items: [Input]){
        vocab = Array(Set(items)).sorted()
        reverseMap = [:]
        for (i,x) in vocab!.enumerated(){ reverseMap![x] = Output(i) }
    }
    
    public func process1  (item: Input)  -> Output { return reverseMap![item]! }
    public func deprocess1(item: Output) -> Input  { return vocab![Int(item)] }
}

//export
public struct SplitLabelData<Item,Label>{
    public let train: ([Item], [Label])
    public let valid: ([Item], [Label])
    
    public init(train: ([Item], [Label]), valid: ([Item], [Label])){
        (self.train, self.valid) = (train, valid)
    }
    
    public init(splitData: SplitData<Item>, labelWith: (Item) -> (Label)){
        train = (splitData.train, splitData.train.map { labelWith($0) })
        valid = (splitData.valid, splitData.valid.map { labelWith($0) })
    }
}

//export
public func labelWithParent(_ fName: Path) -> String { return fName.parent.basename() }

let ld = SplitLabelData(splitData: sd, labelWith: labelWithParent)

//TODO: figure out the generics to put this in a function
var proc = CategoryProcessor()
proc.initState(items: ld.train.1)
let ld1 = SplitLabelData(train: (ld.train.0, proc.process(items: ld.train.1)),
                         valid: (ld.valid.0, proc.process(items: ld.valid.1)))

struct LabeledElement: TensorGroup {
    let fnames: StringTensor
    let labels: Tensor<Int32>
}

let trainDs: Dataset<LabeledElement> = Dataset(elements: LabeledElement(
    fnames: StringTensor(ld1.train.0.map{ $0.string }), 
    labels: Tensor<Int32>(ld1.train.1)))

func openAndResize(fname: StringTensor, size: Int) -> Tensor<Float>{
    let imgBytes = Raw.readFile(filename: fname)
    let decodedImg = Raw.decodeJpeg(contents: imgBytes, channels: 3, dctMethod: "")
    let resizedImg = Raw.resizeBicubic(
        images: Tensor<UInt8>([decodedImg]), 
        size: Tensor<Int32>([Int32(size), Int32(size)]))
    return resizedImg.reshaped(to: TensorShape(Int32(size), Int32(size), 3))
}

let myresize = {openAndResize(fname: $0, size: 128)}

var trainDs: Dataset<DataBatch<Tensor<Float>, Tensor<Int32>>> = Dataset<LabeledElement>(elements: LabeledElement(
    fnames: StringTensor(ld1.train.0.map{ $0.string }), 
    labels: Tensor<Int32>(ld1.train.1))).map{ DataBatch(
        xb: openAndResize(fname: $0.fnames, size: 128),
        yb: $0.labels)}

trainDs = trainDs.batched(Int64(64))

var firstBatch: DataBatch<Tensor<Float>, Tensor<Int32>>? = nil
for batch in trainDs.shuffled() {
    firstBatch = batch
    break
}