From e49cee0e77abdfe6e9783e4b8c8a7612d4d5ea83 Mon Sep 17 00:00:00 2001
From: nipunru <nipuna.tc09@gmail.com>
Date: Sun, 29 Mar 2020 18:54:39 +0530
Subject: [PATCH] Add confidence level

---
 .../com/nipunru/nsfwdetector/NSFWDetector.kt   | 18 +++++++++---------
 1 file changed, 9 insertions(+), 9 deletions(-)

diff --git a/nsfwdetector/src/main/java/com/nipunru/nsfwdetector/NSFWDetector.kt b/nsfwdetector/src/main/java/com/nipunru/nsfwdetector/NSFWDetector.kt
index 781bd3d..81ad9c4 100644
--- 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)
         }
     }
 }
\ No newline at end of file