Tweaks

litan · Oct 14, 2024 · 177c823 · 177c823
1 parent 1ad5f24
commit 177c823
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 10 deletions.
diff --git a/src/main/scala/net/kogics/kojo/aiapp/FaceEmbedder.scala b/src/main/scala/net/kogics/kojo/aiapp/FaceEmbedder.scala
@@ -75,15 +75,14 @@ class FaceEmbedder(modelDir: String) {
     faceFeatureModel.toString
   }
 
-  def faceEmbedding(image: Mat): Array[Float] = {
+  def faceEmbedding(imageMat: Mat): Array[Float] = {
     import ai.djl.pytorch.jni.JniUtils
     JniUtils.setGraphExecutorOptimize(false)
     Using.Manager { use =>
       println("Calculating embedding")
-      val src = Java2DFrameUtils.toBufferedImage(image)
-      val faceImage = bufferedImageToDjlImage(src)
-      println(s"Feeding into net, img(${faceImage.getWidth}, ${faceImage.getHeight})")
-      val ret = faceFeaturePredictor.predict(faceImage)
+      val faceDjlImage = matToDjlImage(imageMat)
+      println(s"Feeding into net, img(${faceDjlImage.getWidth}, ${faceDjlImage.getHeight})")
+      val ret = faceFeaturePredictor.predict(faceDjlImage)
       println("Done")
       ret
     }.get

diff --git a/src/main/scala/net/kogics/kojo/aiapp/ObjectDetector.scala b/src/main/scala/net/kogics/kojo/aiapp/ObjectDetector.scala
@@ -22,7 +22,6 @@ class ObjectDetector(modelDir: String) {
       "resize" -> Boolean.box(true),
       "rescale" -> Boolean.box(true),
       "optApplyRatio" -> Boolean.box(true),
-      "toTensor" -> Boolean.box(true),
       "threshold" -> Double.box(0.4),
     )
 
@@ -32,20 +31,19 @@ class ObjectDetector(modelDir: String) {
         .setTypes(classOf[Image], classOf[DetectedObjects])
         .optModelPath(Paths.get(modelDir))
         .optTranslatorFactory(new ai.djl.modality.cv.translator.YoloV5TranslatorFactory())
-//        .optArguments(args.asJava)
+        .optArguments(args.asJava)
         .build()
 
     val mdl = criteria.loadModel()
     println("Done")
     (mdl, mdl.newPredictor())
   }
 
-  def findObjects(image: Mat): (DetectedObjects, BufferedImage) = {
+  def findObjects(imageMat: Mat): (DetectedObjects, BufferedImage) = {
     import ai.djl.pytorch.jni.JniUtils
     JniUtils.setGraphExecutorOptimize(false)
     Using.Manager { use =>
-      val src = Java2DFrameUtils.toBufferedImage(image)
-      val djlImage = bufferedImageToDjlImage(src)
+      val djlImage = matToDjlImage(imageMat)
       val detectedObjects = predictor.predict(djlImage)
       djlImage.drawBoundingBoxes(detectedObjects)
       (detectedObjects, djlImage.getWrappedImage.asInstanceOf[BufferedImage])

diff --git a/src/main/scala/net/kogics/kojo/aiapp/package.scala b/src/main/scala/net/kogics/kojo/aiapp/package.scala
@@ -17,4 +17,7 @@ package object aiapp {
     Java2DFrameUtils.toBufferedImage(imageMat)
   }
 
+  def matToDjlImage(imageMat: Mat): Image = {
+    bufferedImageToDjlImage(matToBufferedImage(imageMat))
+  }
 }