output with softmax

app.py

@@ -12,13 +12,16 @@ model, transforms = get_watermarks_detection_model(
 predictor = WatermarksPredictor(model, transforms, 'cpu')
 
 
-def predict(image):
-    result = predictor.predict_image(image)
-    return 'watermarked' if result else 'clean' # prints "watermarked"
+def predict(image, threshold=0.5):
+    result = predictor.predict_image_confidence(image)
+    values = result.tolist()
+    wm_flag = 1 if values[1] >= threshold else 0
+    return 'watermarked' if wm_flag else 'clean', "%.4f" % values[1] # prints "watermarked"
 
 
 examples = glob.glob(os.path.join('images', 'clean', '*'))
 examples.extend(glob.glob(os.path.join('images', 'watermark', '*')))
-iface = gr.Interface(fn=predict, inputs=[gr.inputs.Image(type="pil")], 
-             examples=examples, outputs="text")
+examples = [[e, 0.5] for e in examples]
+iface = gr.Interface(fn=predict, inputs=[gr.inputs.Image(type="pil"), gr.inputs.Number(label="threshold", default=0.5), ], 
+             examples=examples, outputs=[gr.outputs.Textbox(label="class"), gr.outputs.Textbox(label="wm_confidence")])
 iface.launch()

wmdetection/pipelines/predictor.py

@@ -51,6 +51,12 @@ class WatermarksPredictor:
         outputs = self.wm_model(input_img.to(self.device))
         result = torch.max(outputs, 1)[1].cpu().reshape(-1).tolist()[0]
         return result
+    
+    def predict_image_confidence(self, pil_image):
+        pil_image = pil_image.convert("RGB")
+        input_img = self.classifier_transforms(pil_image).float().unsqueeze(0)
+        outputs = self.wm_model(input_img.to(self.device))
+        return torch.nn.functional.softmax(outputs, dim=1).cpu().reshape(-1)
         
     def run(self, files, num_workers=8, bs=8, pbar=True):
         eval_dataset = ImageDataset(files, self.classifier_transforms)