Code from F# Advent 2020 Image Classifier

ImageClassifierFsAdvent2020.fsx

#r "nuget:Microsoft.ML"
#r "nuget:Microsoft.ML.Vision"
#r "nuget:Microsoft.ML.ImageAnalytics"
#r "nuget:SciSharp.TensorFlow.Redist" 

open System
open System.IO
open Microsoft.ML
open Microsoft.ML.Data
open Microsoft.ML.Vision

[<CLIMutable>]
type ImageData = {
    ImagePath: string
    Label: string
}

[<CLIMutable>]
type ImagePrediction = {
    ImagePath: string
    PredictedLabel: string
}

// Load images from directory and assign directory name as labels
let loadImagesFromDirectory (path:string) (useDirectoryAsLabel:bool) = 

    let files = Directory.GetFiles(path, "*",searchOption=SearchOption.AllDirectories)

    files
    |> Array.filter(fun file -> 
        (Path.GetExtension(file) = ".jpg") ||
        (Path.GetExtension(file) = ".png"))
    |> Array.map(fun file -> 
        let mutable label = Path.GetFileName(file)
        if useDirectoryAsLabel then
            label <-  Directory.GetParent(file).Name
        else
            let mutable brk = false
            for index in 0..label.Length do
                while not brk do
                    if not (label.[index] |> Char.IsLetter) then
                        label <- label.Substring(0,index)
                        brk <- true

        {ImagePath=file; Label=label}
    )

// Initialize MLContext 
let ctx = MLContext()

// Load images 
let imageData = loadImagesFromDirectory "C:/Datasets/fsadvent2020/Train" true

// Createa an IDataView for the images. 
let imageIdv = ctx.Data.LoadFromEnumerable<ImageData>(imageData)

// Set image classifier options
let classifierOptions = ImageClassificationTrainer.Options()
classifierOptions.FeatureColumnName <- "Image" 
classifierOptions.LabelColumnName <- "LabelKey" 
classifierOptions.TestOnTrainSet <- true  
classifierOptions.Arch <- ImageClassificationTrainer.Architecture.ResnetV2101
classifierOptions.MetricsCallback <- Action<ImageClassificationTrainer.ImageClassificationMetrics>(fun x -> printfn "%s" (x.ToString()))

// Define training / consumption pipeline
let pipeline = 
    EstimatorChain()
        .Append(ctx.Transforms.LoadRawImageBytes("Image",null,"ImagePath"))
        .Append(ctx.Transforms.Conversion.MapValueToKey("LabelKey","Label"))
        .Append(ctx.MulticlassClassification.Trainers.ImageClassification(classifierOptions))
        .Append(ctx.Transforms.Conversion.MapKeyToValue("PredictedLabel"))

// Train the model
let model = pipeline.Fit(imageIdv)

// (Optional) Save the model
ctx.Model.Save(model,imageIdv.Schema,"fsadvent2020-model.zip")

// Load the model
//let (model,schema) = ctx.Model.Load("fsadvent2020-model.zip")

// Load test images
let testImages = 
    Directory.GetFiles("C:/Datasets/fsadvent2020/Test")
    |> Array.map(fun file -> {ImagePath=file; Label=""})

// Create IDataView for test images
let testImageIdv = ctx.Data.LoadFromEnumerable<ImageData>(testImages)

// Make predictions
let predictionIdv = model.Transform(testImageIdv)

// Display predictions
let predictions = ctx.Data.CreateEnumerable<ImagePrediction>(predictionIdv,false)

predictions |> Seq.iter(fun pred -> 
    printfn "%s is %s" (Path.GetFileNameWithoutExtension(pred.ImagePath)) pred.PredictedLabel)