Add confidence level

nipunru · Mar 29, 2020 · e49cee0 · e49cee0
1 parent 53b1ee0
commit e49cee0
Showing 1 changed file with 9 additions and 9 deletions.
diff --git a/nsfwdetector/src/main/java/com/nipunru/nsfwdetector/NSFWDetector.kt b/nsfwdetector/src/main/java/com/nipunru/nsfwdetector/NSFWDetector.kt
@@ -26,12 +26,12 @@ object NSFWDetector {
      * This function return weather the bitmap is NSFW or not
      * @param bitmap: Bitmap Image
      * @param confidenceThreshold: Float 0 to 1 (Default is 0.7)
-     * @return callback with Boolean and Bitmap
+     * @return callback with isNSFW(Boolean), confidence(Float), and image(Bitmap)
      */
     fun isNSFW(
         bitmap: Bitmap,
         confidenceThreshold: Float = CONFIDENCE_THRESHOLD,
-        callback: (Boolean, Bitmap) -> Unit
+        callback: (Boolean, Float, Bitmap) -> Unit
     ) {
         var threshold = confidenceThreshold
 
@@ -45,29 +45,29 @@ object NSFWDetector {
                 when (label.text) {
                     LABEL_SFW -> {
                         if (label.confidence > threshold) {
-                            callback(true, bitmap)
+                            callback(true,label.confidence, bitmap)
                         } else {
-                            callback(false, bitmap)
+                            callback(false,label.confidence, bitmap)
                         }
                     }
                     LABEL_NSFW -> {
                         if (label.confidence < (1 - threshold)) {
-                            callback(true, bitmap)
+                            callback(true,label.confidence, bitmap)
                         } else {
-                            callback(false, bitmap)
+                            callback(false,label.confidence, bitmap)
                         }
                     }
                     else -> {
-                        callback(false, bitmap)
+                        callback(false,0.0F , bitmap)
                     }
                 }
             } catch (e: Exception) {
                 Log.e(TAG, e.localizedMessage ?: "NSFW Scan Error")
-                callback(false, bitmap)
+                callback(false,0.0F , bitmap)
             }
         }.addOnFailureListener { e ->
             Log.e(TAG, e.localizedMessage ?: "NSFW Scan Error")
-            callback(false, bitmap)
+            callback(false,0.0F , bitmap)
         }
     }
 }