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 = 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>> {