From 793cf9d5a1cb2ae7eb1ae3e6420ae0b5baf2aa65 Mon Sep 17 00:00:00 2001
From: Rushikesh Tote <rushi.tote@gmail.com>
Date: Wed, 1 Jun 2022 18:33:27 +0530
Subject: [PATCH] refactor

---
 src/constants/machineLearning/config.ts       |  2 ++
 .../machineLearning/imageSceneService.ts      | 20 ++++++++++++-------
 2 files changed, 15 insertions(+), 7 deletions(-)

diff --git a/src/constants/machineLearning/config.ts b/src/constants/machineLearning/config.ts
index 099bf2a17..454599fae 100644
--- a/src/constants/machineLearning/config.ts
+++ b/src/constants/machineLearning/config.ts
@@ -89,3 +89,5 @@ export const MOBILEFACENET_FACE_SIZE = 112;
 export const TESSERACT_MIN_IMAGE_WIDTH = 44;
 export const TESSERACT_MIN_IMAGE_HEIGHT = 20;
 export const TESSERACT_MAX_IMAGE_DIMENSION = 720;
+
+export const SCENE_DETECTION_IMAGE_SIZE = [224, 224];
diff --git a/src/services/machineLearning/imageSceneService.ts b/src/services/machineLearning/imageSceneService.ts
index bc20acc39..2e66f7140 100644
--- a/src/services/machineLearning/imageSceneService.ts
+++ b/src/services/machineLearning/imageSceneService.ts
@@ -6,6 +6,7 @@ import {
     SceneDetectionService,
     Versioned,
 } from 'types/machineLearning';
+import { SCENE_DETECTION_IMAGE_SIZE } from 'constants/machineLearning/config';
 
 class ImageScene implements SceneDetectionService {
     method: Versioned<SceneDetectionMethod>;
@@ -35,9 +36,11 @@ class ImageScene implements SceneDetectionService {
         this.model = model;
 
         // warmup the model
-        const warmupResult = this.model.predict(tf.zeros([1, 224, 224, 3]));
-        await (warmupResult as tf.Tensor).data();
-        (warmupResult as tf.Tensor).dispose();
+        const warmupResult = this.model.predict(
+            tf.zeros([1, 224, 224, 3])
+        ) as tf.Tensor;
+        await warmupResult.data();
+        warmupResult.dispose();
     }
 
     async detectScenes(image: ImageBitmap, minScore: number) {
@@ -52,18 +55,21 @@ class ImageScene implements SceneDetectionService {
 
             // This model takes fixed-shaped (224x224) inputs
             // https://tfhub.dev/sayannath/lite-model/image-scene/1
-            let resizedTensor = tf.image.resizeBilinear(tensor, [224, 224]);
+            let resizedTensor = tf.image.resizeBilinear(
+                tensor,
+                SCENE_DETECTION_IMAGE_SIZE as [number, number]
+            );
 
             resizedTensor = tf.expandDims(resizedTensor);
             resizedTensor = tf.cast(resizedTensor, 'float32');
 
-            const output = this.model.predict(resizedTensor);
+            const output = this.model.predict(resizedTensor) as tf.Tensor;
 
             return output;
         });
 
-        const data = await (output as tf.Tensor).data();
-        (output as tf.Tensor).dispose();
+        const data = await output.data();
+        output.dispose();
 
         const scenes = this.getScenes(
             data as Float32Array,