Add calibration support to WebUI.

app.py

@@ -1,4 +1,6 @@
-from io import BytesIO
+import csv
+import os
+from io import BytesIO, StringIO
 from threading import Lock
 
 import numpy as np
@@ -31,8 +33,11 @@ model_lock = Lock()
 model, tag_list = load_model("models/jtp-3-hydra.safetensors", device=device)
 model.requires_grad_(False)
 
+def rewrite_tag(tag: str) -> str:
+    return tag.replace("_", " ").replace("vulva", "pussy")
+
 tags = {
-    tag.replace("_", " ").replace("vulva", "pussy"): idx
+    rewrite_tag(tag): idx
     for idx, tag in enumerate(tag_list)
 }
 tag_list = list(tags.keys())
@@ -169,12 +174,19 @@ def cam_composite(image: Image.Image, cam: np.ndarray):
 
     return image
 
-def filter_tags(predictions: dict[str, float], threshold: float):
-    predictions = {
-        key: value
-        for key, value in predictions.items()
-        if value >= threshold
-    }
+def filter_tags(predictions: dict[str, float], threshold: float, calibration: dict[str, float] | None):
+    if calibration is None:
+        predictions = {
+            key: value
+            for key, value in predictions.items()
+            if value >= threshold
+        }
+    else:
+        predictions = {
+            key: value
+            for key, value in predictions.items()
+            if value >= calibration.get(key, float("inf"))
+        }
 
     tag_str = ", ".join(predictions.keys())
     return tag_str, predictions
@@ -224,9 +236,9 @@ def url_submit(url: str):
         processed_image,
     )
 
-def image_changed(image: Image.Image, threshold: float, cam_depth: int):
+def image_changed(image: Image.Image, threshold: float, calibration: dict[str, float] | None, cam_depth: int):
     features, predictions = run_classifier(image, cam_depth)
-    return *filter_tags(predictions, threshold), features, predictions
+    return *filter_tags(predictions, threshold, calibration), features, predictions
 
 def image_clear():
     return (
@@ -235,6 +247,44 @@ def image_clear():
         None, None, {},
     )
 
+def threshold_input(predictions: dict[str, float], threshold: float):
+    return (
+        *filter_tags(predictions, threshold, None), None,
+        gr.Slider(label="Tag Threshold", elem_classes=[]),
+        gr.Textbox(label="Upload Calibration")
+    )
+
+def parse_calibration(data) -> dict[str, float]:
+    return {
+        rewrite_tag(row["tag"]): float(row["threshold"])
+        for row in csv.DictReader(data)
+    }
+
+def calibration_load(predictions: dict[str, float]):
+    try:
+        with open("calibration.csv", "r", encoding="utf-8", newline="") as csv:
+            calibration = parse_calibration(csv)
+    except Exception:
+        return gr.skip(), gr.skip(), gr.skip(), gr.skip(), gr.Textbox(label="Invalid Calibration File")
+
+    return (
+        *filter_tags(predictions, 0.0, calibration), calibration,
+        gr.Slider(label="Using Default Calibration", elem_classes=["inactive-slider"]),
+        gr.Textbox(label="Change Calibration")
+    )
+
+def calibration_changed(predictions: dict[str, float], calibration_data: bytes):
+    try:
+        calibration = parse_calibration(StringIO(calibration_data.decode("utf-8")))
+    except Exception:
+        return gr.skip(), gr.skip(), gr.skip(), gr.skip(), gr.Textbox(label="Invalid Calibration File")
+
+    return (
+        *filter_tags(predictions, 0.0, calibration), calibration,
+        gr.Slider(label="Using Uploaded Calibration", elem_classes=["inactive-slider"]),
+        gr.Textbox(label="Change Calibration")
+    )
+
 def cam_changed(
     display_image: Image.Image,
     image: Image.Image, features: dict[str, Tensor],
@@ -258,6 +308,11 @@ custom_css = """
     ) !important;
     background-size: 100% 100% !important;
 }
+
+.inactive-slider input[type=range] {
+    --slider-color: grey !important;
+}
+
 #image_container-image {
     width: 100%;
     aspect-ratio: 1 / 1;
@@ -280,6 +335,7 @@ with gr.Blocks(
     image_state = gr.State()
     features_state = gr.State()
     predictions_state = gr.State(value={})
+    calibration_state = gr.State()
 
     gr.HTML(
         "<h1 style='display:flex; flex-flow: row nowrap; align-items: center;'>"
@@ -321,8 +377,24 @@ with gr.Blocks(
                 )
 
         with gr.Column():
-            threshold_slider = gr.Slider(minimum=0.00, maximum=1.00, step=0.01, value=0.30, label="Tag Threshold")
-            tag_string = gr.Textbox(lines=3, label="Tags", show_label=True, show_copy_button=True)
+            with gr.Row(variant="panel"):
+                threshold_slider = gr.Slider(
+                    minimum=0.00, maximum=1.00, step=0.01, value=0.30,
+                    label="Tag Threshold", scale=4
+                )
+
+                with gr.Column(), gr.Group():
+                    calibration_default = gr.Button(
+                        interactive=os.path.exists("calibration.csv"),
+                        value="Default Calibration", size="lg",
+                    )
+
+                    calibration_upload = gr.UploadButton(
+                        file_count="single", file_types=["text"], type="binary",
+                        label="Upload Calibration", size="md", variant="secondary",
+                    )
+
+            tag_string = gr.Textbox(lines=3, label="Tags", show_copy_button=True)
             tag_box = gr.Label(num_top_classes=250, show_label=False, show_heading=False)
 
     image.upload(
@@ -337,7 +409,7 @@ with gr.Blocks(
         show_progress_on=[image]
     ).then(
         fn=image_changed,
-        inputs=[image_state, threshold_slider, cam_depth],
+        inputs=[image_state, threshold_slider, calibration_state, cam_depth],
         outputs=[
             tag_string, tag_box,
             features_state, predictions_state,
@@ -358,7 +430,7 @@ with gr.Blocks(
         show_progress_on=[url]
     ).then(
         fn=image_changed,
-        inputs=[image_state, threshold_slider, cam_depth],
+        inputs=[image_state, threshold_slider, calibration_state, cam_depth],
         outputs=[
             tag_string, tag_box,
             features_state, predictions_state,
@@ -379,13 +451,29 @@ with gr.Blocks(
     )
 
     threshold_slider.input(
-        fn=filter_tags,
+        fn=threshold_input,
         inputs=[predictions_state, threshold_slider],
-        outputs=[tag_string, tag_box],
+        outputs=[tag_string, tag_box, calibration_state, threshold_slider, calibration_upload],
         trigger_mode='always_last',
         show_progress='hidden'
     )
 
+    calibration_default.click(
+        fn=calibration_load,
+        inputs=[predictions_state],
+        outputs=[tag_string, tag_box, calibration_state, threshold_slider, calibration_upload],
+        show_progress='hidden'
+    )
+
+    calibration_upload.upload(
+        fn=calibration_changed,
+        inputs=[predictions_state, calibration_upload],
+        outputs=[tag_string, tag_box, calibration_state, threshold_slider, calibration_upload],
+        trigger_mode='always_last',
+        show_progress='minimal',
+        show_progress_on=[calibration_upload],
+    )
+
     cam_tag.input(
         fn=cam_changed,
         inputs=[
@@ -430,5 +518,12 @@ with gr.Blocks(
         show_progress_on=[cam_tag]
     )
 
+    scan_timer = gr.Timer()
+    scan_timer.tick(
+        fn=lambda: gr.Button(interactive=os.path.exists("calibration.csv")),
+        outputs=[calibration_default],
+        show_progress='hidden'
+    )
+
 if __name__ == "__main__":
     demo.launch()