From 6bda1dbe06c8ceb9cb99d639689cb4a36dd73a24 Mon Sep 17 00:00:00 2001
From: Denis-Cosmin NUTIU <dnutiu@bitdefender.com>
Date: Sat, 23 Mar 2024 21:36:55 +0200
Subject: [PATCH] add log when model or categories file is not found

---
 .../tag_prediction/ImageTagsPrediction.kt     | 29 +++++++++++++------
 1 file changed, 20 insertions(+), 9 deletions(-)

diff --git a/src/main/kotlin/com/nuculabs/dev/imagetagger/tag_prediction/ImageTagsPrediction.kt b/src/main/kotlin/com/nuculabs/dev/imagetagger/tag_prediction/ImageTagsPrediction.kt
index b44787d..a77a04a 100644
--- a/src/main/kotlin/com/nuculabs/dev/imagetagger/tag_prediction/ImageTagsPrediction.kt
+++ b/src/main/kotlin/com/nuculabs/dev/imagetagger/tag_prediction/ImageTagsPrediction.kt
@@ -19,17 +19,28 @@ class ImageTagsPrediction private constructor() {
     private var modelClasses: MutableList<String> = mutableListOf()
 
     init {
-        logger.info("Loaded ML model.")
-        ImageTagsPrediction::class.java.getResourceAsStream("/AIModels/prediction.onnx").let { modelFile ->
-            ortSession = ortEnv.createSession(
-                modelFile!!.readBytes(),
-                OrtSession.SessionOptions()
+        try {
+            logger.info("Loaded ML model.")
+            ImageTagsPrediction::class.java.getResourceAsStream("/AIModels/prediction.onnx").let { modelFile ->
+                ortSession = ortEnv.createSession(
+                    modelFile!!.readBytes(),
+                    OrtSession.SessionOptions()
+                )
+            }
+            ImageTagsPrediction::class.java.getResourceAsStream("/AIModels/prediction_categories.txt")
+                .let { classesFile ->
+                    modelClasses.addAll(0, classesFile!!.bufferedReader().readLines())
+                }
+            logger.info("Loaded ${modelClasses.size} model classes.")
+        } catch (e: NullPointerException) {
+            logger.severe(
+                "Failed to load model file or categories file. If you're building the project from " +
+                        "source, please follow the instructions from the README.md: " +
+                        "https://github.com/dnutiu/ImageTagger." +
+                        "Exception ${e.message}"
             )
+            throw e
         }
-        ImageTagsPrediction::class.java.getResourceAsStream("/AIModels/prediction_categories.txt").let { classesFile ->
-            modelClasses.addAll(0, classesFile!!.bufferedReader().readLines())
-        }
-        logger.info("Loaded ${modelClasses.size} model classes.")
     }
 
     private fun processImage(bufferedImage: BufferedImage): Array<Array<Array<FloatArray>>> {