simplify

detect_demo_pvalue.py

@@ -12,11 +12,12 @@ calibration = hf_hub_download("imatag/stable-signature-bzh-detector-resnet18", f
 with safe_open(calibration, framework="pt") as f:
     calibration_logits = f.get_tensor("logits")
 
-img = Image.open(sys.argv[1]).convert("RGB")
+filename = sys.argv[1]
+img = Image.open(filename).convert("RGB")
 inputs = image_processor(img, return_tensors="pt")
 with torch.no_grad():
-    p = model(**inputs).logits[...,0]
-    p = (1 + torch.searchsorted(calibration_logits, p)) / calibration_logits.shape[0]
+    p = model(**inputs).logits[...,0:1]
+    p = (1 + torch.sum(calibration_logits <= p, dim=-1)) / calibration_logits.shape[0]
     p = p.item()
 
 print(f"approximate p-value: {p}")