add second class in export to make it compatible with inference api

calibration.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:f0eac4b0be8eb96b2a6fc7596727954ebe600951b3c4719b456948200d75f58e
+oid sha256:ebfb5a2481c315bca61b2348b0268d621a351d48cf5e07785d45d5877f67ebf3
 size 1999934

export_detector.py

@@ -10,9 +10,6 @@ from PIL import Image
 # read logits file
 data=(np.asarray([float(x) for x in open(sys.argv[1]).readlines()]))
 
-# negate for consistency with "1" = "watermarked"
-data = -data
-
 # sort and convert to safetensors format
 data = np.sort(data)
 data_min = data.min()
@@ -47,12 +44,29 @@ image_processor = BlipImageProcessor(do_resize=True,
 
 detector = AutoModelForImageClassification.from_pretrained(detector_path)
 
-# make it output "1" for "watermarked"
+# make it output 2 labels, with label "0" for "watermarked"
 detector.eval()
 with torch.no_grad():
-    detector.classifier[1].weight.copy_(-detector.classifier[1].weight)
+    w0 = detector.classifier[1].weight
+    b0 = detector.classifier[1].bias
+    fdim = w0.shape[1]
+    w = torch.nn.Parameter(torch.zeros((2, fdim), dtype=w0.dtype))
+    w[0, :] = w0
+    w[1, :] = -w0
+    detector.classifier[1].weight = w
+    b = torch.nn.Parameter(torch.zeros((2,), dtype=b0.dtype))
+    b[0] = b0
+    b[1] = -b0
+    detector.classifier[1].bias = b
+labels = ["no watermark detected", "watermarked"]
+label2id, id2label = dict(), dict()
+for i, label in enumerate(labels):
+    label2id[label] = str(i)
+    id2label[str(i)] = label
+detector.config.id2label=id2label
+detector.config.label2id=label2id
 
-#detector.push_to_hub("imatag/stable-signature-bzh-detector-resnet18")
+detector.push_to_hub("imatag/stable-signature-bzh-detector-resnet18")
 #image_processor.push_to_hub("imatag/stable-signature-bzh-detector-resnet18")
 examples = ['examples/not_watermarked.png', 'examples/watermarked.png']