learning_model 1.0.38 → 1.0.40

package/dist/utils/tf.js

@@ -0,0 +1,124 @@
+"use strict";
+var __createBinding = (this && this.__createBinding) || (Object.create ? (function(o, m, k, k2) {
+    if (k2 === undefined) k2 = k;
+    var desc = Object.getOwnPropertyDescriptor(m, k);
+    if (!desc || ("get" in desc ? !m.__esModule : desc.writable || desc.configurable)) {
+      desc = { enumerable: true, get: function() { return m[k]; } };
+    }
+    Object.defineProperty(o, k2, desc);
+}) : (function(o, m, k, k2) {
+    if (k2 === undefined) k2 = k;
+    o[k2] = m[k];
+}));
+var __setModuleDefault = (this && this.__setModuleDefault) || (Object.create ? (function(o, v) {
+    Object.defineProperty(o, "default", { enumerable: true, value: v });
+}) : function(o, v) {
+    o["default"] = v;
+});
+var __importStar = (this && this.__importStar) || function (mod) {
+    if (mod && mod.__esModule) return mod;
+    var result = {};
+    if (mod != null) for (var k in mod) if (k !== "default" && Object.prototype.hasOwnProperty.call(mod, k)) __createBinding(result, mod, k);
+    __setModuleDefault(result, mod);
+    return result;
+};
+var __awaiter = (this && this.__awaiter) || function (thisArg, _arguments, P, generator) {
+    function adopt(value) { return value instanceof P ? value : new P(function (resolve) { resolve(value); }); }
+    return new (P || (P = Promise))(function (resolve, reject) {
+        function fulfilled(value) { try { step(generator.next(value)); } catch (e) { reject(e); } }
+        function rejected(value) { try { step(generator["throw"](value)); } catch (e) { reject(e); } }
+        function step(result) { result.done ? resolve(result.value) : adopt(result.value).then(fulfilled, rejected); }
+        step((generator = generator.apply(thisArg, _arguments || [])).next());
+    });
+};
+Object.defineProperty(exports, "__esModule", { value: true });
+exports.cropTensor = exports.capture = exports.imageToTensor = exports.mobileNetURL = exports.loadModel = exports.isTensor = void 0;
+const tf = __importStar(require("@tensorflow/tfjs"));
+function isTensor(c) {
+    return typeof c.dataId === 'object' && typeof c.shape === 'object';
+}
+exports.isTensor = isTensor;
+function loadModel() {
+    return __awaiter(this, void 0, void 0, function* () {
+        const trainLayerV1 = 'conv_pw_13_relu';
+        const trainLayerV2 = 'out_relu';
+        var mobileNetVersion = 2;
+        const modelURL = mobileNetURL(mobileNetVersion);
+        const load_model = yield tf.loadLayersModel(modelURL);
+        if (mobileNetVersion == 1) {
+            const layer = load_model.getLayer(trainLayerV1);
+            const truncatedModel = tf.model({
+                inputs: load_model.inputs,
+                outputs: layer.output
+            });
+            const model = tf.sequential();
+            model.add(truncatedModel);
+            model.add(tf.layers.flatten());
+            return model;
+        }
+        else {
+            const layer = load_model.getLayer(trainLayerV2);
+            const truncatedModel = tf.model({
+                inputs: load_model.inputs,
+                outputs: layer.output
+            });
+            const model = tf.sequential();
+            model.add(truncatedModel);
+            model.add(tf.layers.globalAveragePooling2d({})); // go from shape [7, 7, 1280] to [1280]
+            return model;
+        }
+    });
+}
+exports.loadModel = loadModel;
+function mobileNetURL(version) {
+    if (version == 1) {
+        return "https://storage.googleapis.com/tfjs-models/tfjs/mobilenet_v1_1.0_224/model.json";
+    }
+    return "https://storage.googleapis.com/teachable-machine-models/mobilenet_v2_weights_tf_dim_ordering_tf_kernels_1.0_224_no_top/model.json";
+}
+exports.mobileNetURL = mobileNetURL;
+function imageToTensor(data) {
+    let tensor;
+    if (data instanceof tf.Tensor) {
+        tensor = data;
+    }
+    else {
+        // MobileNet 모델 로드
+        tensor = tf.browser.fromPixels(data);
+    }
+    return tensor;
+}
+exports.imageToTensor = imageToTensor;
+function capture(rasterElement, grayscale) {
+    return tf.tidy(() => {
+        const pixels = tf.browser.fromPixels(rasterElement);
+        // crop the image so we're using the center square
+        const cropped = cropTensor(pixels, grayscale);
+        // Expand the outer most dimension so we have a batch size of 1
+        const batchedImage = cropped.expandDims(0);
+        // Normalize the image between -1 and a1. The image comes in between 0-255
+        // so we divide by 127 and subtract 1.
+        return batchedImage.toFloat().div(tf.scalar(127)).sub(tf.scalar(1));
+    });
+}
+exports.capture = capture;
+function cropTensor(img, grayscaleModel, grayscaleInput) {
+    const size = Math.min(img.shape[0], img.shape[1]);
+    const centerHeight = img.shape[0] / 2;
+    const beginHeight = centerHeight - (size / 2);
+    const centerWidth = img.shape[1] / 2;
+    const beginWidth = centerWidth - (size / 2);
+    if (grayscaleModel && !grayscaleInput) {
+        //cropped rgb data
+        let grayscale_cropped = img.slice([beginHeight, beginWidth, 0], [size, size, 3]);
+        grayscale_cropped = grayscale_cropped.reshape([size * size, 1, 3]);
+        const rgb_weights = [0.2989, 0.5870, 0.1140];
+        grayscale_cropped = tf.mul(grayscale_cropped, rgb_weights);
+        grayscale_cropped = grayscale_cropped.reshape([size, size, 3]);
+        grayscale_cropped = tf.sum(grayscale_cropped, -1);
+        grayscale_cropped = tf.expandDims(grayscale_cropped, -1);
+        return grayscale_cropped;
+    }
+    return img.slice([beginHeight, beginWidth, 0], [size, size, 3]);
+}
+exports.cropTensor = cropTensor;

package/lib/learning/base.ts

@@ -5,8 +5,6 @@ import { io } from '@tensorflow/tfjs-core';
 interface LearningInterface {
     // tf 모델
     model: tf.LayersModel | null;
-    // 라벨 array
-    labels: string[];
     // 이미 train중인지 여부 - 이미 진행중인경우 취소
     isRunning: boolean;
     // 준비 

package/lib/learning/mobilenet.ts

@@ -5,8 +5,14 @@
 ///////////////////////////////////////////////////////////////////////////
 
 import * as tf from '@tensorflow/tfjs';
+import { dispose } from '@tensorflow/tfjs';
 import { io } from '@tensorflow/tfjs-core';
 import LearningInterface from './base';
+import { capture, loadModel, isTensor } from '../utils/tf';
+import { cropTo } from '../utils/canvas';
+import { flatOneHot, fisherYates, Sample } from '../utils/dataset';
+import { Initializer } from '@tensorflow/tfjs-layers/dist/initializers';
+
 
 class LearningMobilenet implements LearningInterface {
   model: tf.LayersModel | null;
@@ -14,41 +20,39 @@ class LearningMobilenet implements LearningInterface {
   batchSize: number;
   learningRate: number;
   validateRate: number;
-  labels: string[];
-  modelURL: string;
   isRunning: boolean;
   isReady: boolean;
   isTrainedDone: boolean;
   limitSize: number;
   mobilenetModule: tf.LayersModel | null;
-  imageTensors: tf.Tensor<tf.Rank>[] = [];
+  imageExamples: Float32Array[][] = [];
+  classNumber: string[] = [];
   
   readonly MOBILE_NET_INPUT_WIDTH = 224;
   readonly MOBILE_NET_INPUT_HEIGHT = 224;
   readonly MOBILE_NET_INPUT_CHANNEL = 3;
   readonly IMAGE_NORMALIZATION_FACTOR = 255.0;
   
+
   constructor({ 
-    modelURL = 'https://storage.googleapis.com/tfjs-models/tfjs/mobilenet_v1_0.25_224/model.json', // 디폴트 mobilenet 이미지 
-    epochs = 10,
+    epochs = 50,
     batchSize = 16,
     limitSize = 2,
     learningRate = 0.001,
-    validateRate = 0.2,
-  }: { modelURL?: string, epochs?: number, batchSize?: number, limitSize?: number, learningRate?: number, validateRate?: number} = {}) {
+    validateRate = 0.15,
+  }: { epochs?: number, batchSize?: number, limitSize?: number, learningRate?: number, validateRate?: number} = {}) {
     this.model = null;
     this.epochs = epochs;
     this.batchSize = batchSize;
     this.learningRate = learningRate;
     this.validateRate = validateRate;
-    this.labels = [];
-    this.modelURL = modelURL;
     this.isRunning = false;
     this.isReady = false;
     this.isTrainedDone = false;
     this.limitSize = limitSize;
     this.mobilenetModule = null;
-    this.imageTensors = [];
+    this.classNumber = [];
+
   }
 
   // 진행 상태를 나타내는 이벤트를 정의합니다.
@@ -73,44 +77,93 @@ class LearningMobilenet implements LearningInterface {
 
     try {
       this.model = await tf.loadLayersModel(jsonURL);
-      this.labels = labels;
+      for(var i = 0; i < labels.length; i++) {
+        this.registerClassNumber(labels[i]);
+    }
       this.isReady = true;
       this.model.summary();
     } catch (error) {
       console.error('Model load failed', error);
       throw error;
     }
-    
   }
 
 
+  private registerClassNumber(value: string): number {
+    // 중복 값의 인덱스를 찾습니다.
+    const existingIndex = this.classNumber.indexOf(value);
+
+    // 중복 값이 있다면 해당 인덱스를 반환합니다.
+    if (existingIndex !== -1) {
+        return existingIndex;
+    }
+
+    // 중복 값이 없다면 새로운 항목을 추가하고 그 인덱스를 반환합니다.
+    this.classNumber.push(value);
+    return this.classNumber.length - 1;
+  }
+
+  private _convertToTfDataset() {
+    for (let i = 0; i < this.imageExamples.length; i++) {
+        this.imageExamples[i] = fisherYates(this.imageExamples[i]) as Float32Array[];
+    }
+
+    const trainDataset: Sample[] = [];
+    const validationDataset: Sample[] = [];
+    
+    for (let i = 0; i < this.imageExamples.length; i++) {
+        const classLength = this.imageExamples[i].length;
+
+        // 클래스의 전체 데이터 수를 사용하여 학습 및 검증 데이터 수 계산
+        const numValidation = Math.ceil(this.validateRate * classLength);
+        const numTrain = classLength - numValidation;
+
+        // One-Hot 인코딩을 사용하여 라벨 생성
+        const y = flatOneHot(i, this.classNumber.length);
+
+        // numTrain과 numValidation에 따라 데이터를 학습 및 검증 데이터로 분할
+        const classTrain: Sample[] = this.imageExamples[i].slice(0, numTrain).map(dataArray => ({ data: dataArray, label: y }));
+        trainDataset.push(...classTrain);
+        
+        const classValidation: Sample[] = this.imageExamples[i].slice(numTrain, numTrain + numValidation).map(dataArray => ({ data: dataArray, label: y }));
+        validationDataset.push(...classValidation);
+    }
+    
+    // Shuffle entire datasets
+    const shuffledTrainDataset = fisherYates(trainDataset) as Sample[];
+    const shuffledValidationDataset = fisherYates(validationDataset) as Sample[];
+
+    // Convert to tf.data.Dataset
+    const trainX = tf.data.array(shuffledTrainDataset.map(sample => sample.data));
+    const validationX = tf.data.array(shuffledValidationDataset.map(sample => sample.data));
+    const trainY = tf.data.array(shuffledTrainDataset.map(sample => sample.label));
+    const validationY = tf.data.array(shuffledValidationDataset.map(sample => sample.label));
+
+    return {
+        trainDataset: tf.data.zip({ xs: trainX, ys: trainY }),
+        validationDataset: tf.data.zip({ xs: validationX, ys: validationY })
+    };
+}
+
+
   // 학습 데이타 등록 
   public async addData(label: string, data: any): Promise<void> {
     try {
-      const imgTensor = tf.tidy(() => {
-        if (this.mobilenetModule !== null) {
-          console.log('predict before', data);
-          const t = tf.browser.fromPixels(data)
-          // Resize the image to match the model's input shape
-          const resizedTensor = tf.image.resizeBilinear(t, [224, 224]);
-          const normalizedTensor = resizedTensor.div(255.0);
-
-          // Add an extra dimension to the tensor
-          const expandedTensor = normalizedTensor.expandDims(0);
-          
-          console.log('predict extend', expandedTensor);
-          const predict = this.mobilenetModule.predict(expandedTensor);
-          console.log('predict', predict);
-          return predict;
-        } else {
-          throw new Error('mobilenetModule is null');
+      if (this.mobilenetModule !== null) {
+        const croppedImage = cropTo(data, 224, false);
+        const cap = isTensor(data) ? data : capture(croppedImage, false);
+        const predict = this.mobilenetModule.predict(cap) as tf.Tensor;
+        const activation = predict.dataSync() as Float32Array;
+        const classIndex = this.registerClassNumber(label);
+        if (!this.imageExamples[classIndex]) {
+          this.imageExamples[classIndex] = [];
         }
-      });
-      console.log('imgTensor', imgTensor);
-      this.imageTensors.push(imgTensor as tf.Tensor<tf.Rank>);
-      this.labels.push(label);
-      if(this.labels.length >= this.limitSize) {
-        this.isReady = true;
+        this.imageExamples[classIndex].push(activation);
+        if(this.classNumber.length >= this.limitSize) {
+          this.isReady = true;
+        }
+      } else {
+        throw new Error('mobilenetModule is null');
       }
       return Promise.resolve();
     } catch (error) {
@@ -121,13 +174,14 @@ class LearningMobilenet implements LearningInterface {
 
   public async init() {
     try {
-      this.mobilenetModule = await this.loadModel();
+      this.mobilenetModule = await loadModel();
     } catch(error) {
       console.log('init Error', error);
       throw error;
     } 
   } 
 
+
   // 모델 학습 처리
   public async train(): Promise<tf.History> {
     if (this.isRunning) {
@@ -139,27 +193,23 @@ class LearningMobilenet implements LearningInterface {
       onTrainBegin: (log: any) => {
         this.isTrainedDone = false;
         this.onTrainBegin(log);
-        console.log('Training has started.');
       },
       onTrainEnd: (log: any) => {
         this.isTrainedDone = true;
         this.onTrainEnd(log);
-        console.log('Training has ended.');
         this.isRunning = false;
       },
       onBatchBegin: (batch: any, logs: any) => {
-        console.log(`Batch ${batch} is starting.`);
+        //console.log(`Batch ${batch} is starting.`);
       },
       onBatchEnd: (batch: any, logs: any) => {
-        console.log(`Batch ${batch} has ended.`);
+        //console.log(`Batch ${batch} has ended.`);
       },
       onEpochBegin: (epoch: number, logs: any) => {
-        console.log(`Epoch ${epoch+1} is starting.`, logs);
+        //console.log(`Epoch ${epoch+1} is starting.`, logs);
       },
       onEpochEnd: (epoch: number, logs: any) => {
         this.onEpochEnd(epoch, logs);
-        console.log(`Epoch ${epoch+1} has ended.`);
-        console.log('Loss:', logs);
         this.onLoss(logs.loss);
         this.onProgress(epoch+1);
         this.onEvents({
@@ -171,22 +221,24 @@ class LearningMobilenet implements LearningInterface {
 
     try {
       this.isRunning = true;
-      if (this.labels.length < this.limitSize) {
+      if (this.classNumber.length < this.limitSize) {
         return Promise.reject(new Error('Please train Data need over 2 data length'));
       }
-      this.model = await this._createModel(this.labels.length);
-      const inputData = this._preprocessedInputData(this.model);
-
-      console.log('this.imageTensors', this.imageTensors, inputData);
-      const targetData = this._preprocessedTargetData();
-      console.log('targetData', targetData);
-      const history = await this.model.fit(inputData, targetData, {
+      const datasets = this._convertToTfDataset();
+      const trainData = datasets.trainDataset.batch(this.batchSize);
+      const validationData = datasets.validationDataset.batch(this.batchSize);
+      const optimizer = tf.train.adam(this.learningRate); 
+      const trainModel = await this._createModel(optimizer);
+      const history = await trainModel.fitDataset(trainData, {
         epochs: this.epochs,
-        batchSize: this.batchSize,
-        validationSplit: this.validateRate, // 검증 데이터의 비율 설정
+        validationData: validationData,
         callbacks: customCallback
       });
-      console.log('Model training completed', history);
+      const jointModel = tf.sequential();
+      jointModel.add(this.mobilenetModule!);
+      jointModel.add(trainModel);
+      this.model = jointModel;
+      optimizer.dispose();
       return history;
     } catch (error) {
       this.isRunning = false;
@@ -202,38 +254,19 @@ class LearningMobilenet implements LearningInterface {
     }
   
     try {
-      const imgTensor = tf.tidy(() => {
-        if (this.mobilenetModule !== null) {
-          console.log('predict before', data);
-          const t = tf.browser.fromPixels(data)
-          // Resize the image to match the model's input shape
-          const resizedTensor = tf.image.resizeBilinear(t, [224, 224]);
-          const normalizedTensor = resizedTensor.div(255.0);
-
-          // Add an extra dimension to the tensor
-          const expandedTensor = normalizedTensor.expandDims(0);
-          
-          console.log('predict extend', expandedTensor);
-          const predict = this.mobilenetModule.predict(expandedTensor);
-          console.log('predict', predict);
-          return predict;
-        } else {
-          throw new Error('mobilenetModule is null');
-        }
+      const classProbabilities = new Map<string, number>();
+      const croppedImage = cropTo(data, 224, false);
+
+      const logits = tf.tidy(() => {
+        const captured = capture(croppedImage, false);
+        return this.model!.predict(captured);
       });
-      const predictions = this.model.predict(imgTensor) as tf.Tensor;
-      const predictedClass = predictions.as1D().argMax();
-      const classId = (await predictedClass.data())[0];
-      console.log('classId', classId, this.labels[classId]);
-      const predictionsData = await predictions.data(); // 예측 텐서의 데이터를 비동기로 가져옴
-      const classProbabilities = new Map<string, number>(); // 클래스별 확률 누적값을 저장할 맵
-      
-      console.log('predictionsData', predictionsData);
+      const values = await (logits as tf.Tensor<tf.Rank>).data();
       const EPSILON = 1e-6; // 매우 작은 값을 표현하기 위한 엡실론
-      for (let i = 0; i < predictionsData.length; i++) {
-        let probability = Math.max(0, Math.min(1, predictionsData[i])); // 확률 값을 0과 1 사이로 조정
+      for (let i = 0; i < values.length; i++) {
+        let probability = Math.max(0, Math.min(1, values[i])); // 확률 값을 0과 1 사이로 조정
         probability = probability < EPSILON ? 0 : probability; // 매우 작은 확률 값을 0으로 간주
-        const className = this.labels[i]; // 클래스 이름
+        const className = this.classNumber[i]; // 클래스 이름
         const existingProbability = classProbabilities.get(className);
         if (existingProbability !== undefined) {
           classProbabilities.set(className, existingProbability + probability);
@@ -241,9 +274,8 @@ class LearningMobilenet implements LearningInterface {
           classProbabilities.set(className, probability);
         }
       }
-      console.log('Class Probabilities:', classProbabilities);
-      predictedClass.dispose();
-      predictions.dispose();
+      console.log('classProbabilities', classProbabilities);
+      dispose(logits);
       return classProbabilities;
     } catch (error) {
       throw error;
@@ -253,15 +285,12 @@ class LearningMobilenet implements LearningInterface {
   
   // 모델 저장
   public async saveModel(handlerOrURL: io.IOHandler | string, config?: io.SaveConfig): Promise<void> {
-    console.log('saved model');
     if (!this.isTrainedDone) {
       return Promise.reject(new Error('Train is not done status'));
     }
     await this.model?.save(handlerOrURL, config);
   }
 
-
-
   // 진행중 여부
   public running(): boolean {
     return this.isRunning;
@@ -272,83 +301,41 @@ class LearningMobilenet implements LearningInterface {
     return this.isReady;
   }
 
-
-  // target 라벨 데이타
-  private _preprocessedTargetData(): tf.Tensor<tf.Rank> {
-    // 라벨 unique 처리 & 배열 리턴
-    console.log('uniqueLabels.length', this.labels, this.labels.length);
-    const labelIndices = this.labels.map((label) => this.labels.indexOf(label));
-    console.log('labelIndices', labelIndices);
-    const oneHotEncode = tf.oneHot(tf.tensor1d(labelIndices, 'int32'),this.labels.length);
-    console.log('oneHotEncode', oneHotEncode);
-    return oneHotEncode
-  }
-  
-
-  // 입력 이미지 데이타 
-  private _preprocessedInputData(model: tf.LayersModel): tf.Tensor<tf.Rank> {
-    // 이미지 배열을 배치로 변환 - [null, 224, 224, 3]
-    const inputShape = model.inputs[0].shape;
-    console.log('inputShape', inputShape);
-    // inputShape를 이와 같이 포멧 맞춘다. for reshape to [224, 224, 3]
-    const inputShapeArray = inputShape.slice(1) as number[];
-    console.log('inputShapeArray', inputShapeArray);
-    const inputBatch = tf.stack(this.imageTensors.map((image) => {
-      return tf.reshape(image, inputShapeArray);
-    }));
-    return inputBatch
-  }
-
-
-  private async loadModel(): Promise<tf.LayersModel> {
-    const load_model = await tf.loadLayersModel(this.modelURL);
-    load_model.summary();
-    const layer = load_model.getLayer('conv_pw_13_relu');
-    const truncatedModel = tf.model({
-      inputs: load_model.inputs, 
-      outputs: layer.output
-    })
-    return truncatedModel
-  }
-
   // 모델 저장
-  private async _createModel(numClasses: number): Promise<tf.Sequential> {
+  private async _createModel(optimizer: tf.AdamOptimizer): Promise<tf.Sequential> {
     try {
-        const truncatedModel = await this.loadModel()
-        // 입력 이미지 크기에 맞게 모델 구조 수정
-        const model = tf.sequential();
-        // 추가적인 합성곱 층
-        model.add(tf.layers.conv2d({
-          filters: 64,
-          kernelSize: 3,
-          activation: 'relu',
-          padding: 'same',
-          inputShape: truncatedModel.outputs[0].shape.slice(1)
-        }));
-
-        model.add(tf.layers.flatten());
-
-        // 더 깊은 밀집층
-        model.add(tf.layers.dense({ units: 100, activation: 'relu' }));
-        model.add(tf.layers.dense({ units: 100, activation: 'relu' }));
-
-        // 드롭아웃 층
-        model.add(tf.layers.dropout({ rate: 0.5 }));
-
-        model.add(tf.layers.dense({ units: this.labels.length, activation: 'softmax' }));
-
-
-        const optimizer = tf.train.adam(this.learningRate); // Optimizer를 생성하고 학습률을 설정합니다.
-        model.compile({
-            loss: (numClasses === 2) ? 'binaryCrossentropy': 'categoricalCrossentropy',
-            optimizer: optimizer,
-            metrics: ['accuracy', 'acc']
-        });
-        model.summary();
-        return model;
+      // 입력 이미지 크기에 맞게 모델 구조 수정
+      let varianceScaling: Initializer;
+      varianceScaling = tf.initializers.varianceScaling({});
+
+      const trainModel = tf.sequential({
+        layers: [
+          tf.layers.dense({
+            inputShape: this.mobilenetModule!.outputs[0].shape.slice(1),
+            units: 128,
+            activation: 'relu',
+            kernelInitializer: varianceScaling, // 'varianceScaling'
+            useBias: true
+          }),
+          tf.layers.dense({
+            kernelInitializer: varianceScaling, // 'varianceScaling'
+            useBias: false,
+            activation: 'softmax',
+            units: this.classNumber.length, 
+          })
+        ]
+      });
+      
+      trainModel.compile({
+        loss: 'categoricalCrossentropy',
+        optimizer: optimizer,
+        metrics: ['accuracy']
+      });
+      
+      return trainModel;
     } catch (error) {
-        console.error('Failed to load model', error);
-        throw error;
+      console.error('Failed to load model', error);
+      throw error;
     }
   }
 }

package/lib/utils/canvas.ts

@@ -0,0 +1,49 @@
+type Drawable = ImageData | HTMLImageElement | HTMLCanvasElement | HTMLVideoElement;
+
+const newCanvas = () => document.createElement('canvas');
+
+export function cropTo(image: Drawable, size: number, flipped = false, canvas: HTMLCanvasElement = newCanvas()) {
+  let width: number;
+  let height: number;
+
+  // If ImageData
+  if (image instanceof ImageData) {
+    width = image.width;
+    height = image.height;
+  } 
+  // If image, bitmap, or canvas
+  else if (image instanceof HTMLImageElement || image instanceof HTMLCanvasElement) {
+    width = image.width;
+    height = image.height;
+  } 
+  // If video element
+  else if (image instanceof HTMLVideoElement) {
+    width = image.videoWidth;
+    height = image.videoHeight;
+  } else {
+    throw new Error("Unsupported Drawable type");
+  }
+
+  const min = Math.min(width, height);
+  const scale = size / min;
+  const scaledW = Math.ceil(width * scale);
+  const scaledH = Math.ceil(height * scale);
+  const dx = scaledW - size;
+  const dy = scaledH - size;
+  canvas.width = canvas.height = size;
+  const ctx = canvas.getContext('2d') as CanvasRenderingContext2D;
+
+  // Handle ImageData separately
+  if (image instanceof ImageData) {
+    ctx.putImageData(image, ~~(dx / 2) * -1, ~~(dy / 2) * -1); // Adjust this if needed
+  } else {
+    ctx.drawImage(image, ~~(dx / 2) * -1, ~~(dy / 2) * -1, scaledW, scaledH);
+  }
+
+  if (flipped) {
+    ctx.scale(-1, 1);
+    ctx.drawImage(canvas, size * -1, 0);
+  }
+
+  return canvas;
+}

package/lib/utils/dataset.ts

@@ -0,0 +1,24 @@
+export interface Sample {
+  data: Float32Array;
+  label: number[];
+}
+
+export function flatOneHot(label: number, numClasses: number) {
+  const labelOneHot = new Array(numClasses).fill(0) as number[];
+  labelOneHot[label] = 1;
+
+  return labelOneHot;
+}
+
+export function fisherYates(array: Float32Array[] | Sample[]) {
+  const length = array.length;
+
+  // need to clone array or we'd be editing original as we goo
+  const shuffled = array.slice();
+  for (let i = (length - 1); i > 0; i -= 1) {
+    let randomIndex ;
+    randomIndex = Math.floor(Math.random() * (i + 1));
+    [shuffled[i], shuffled[randomIndex]] = [shuffled[randomIndex],shuffled[i]];
+  }
+  return shuffled;
+}