Upload app.py

2.CNN/app.py

@@ -225,7 +225,7 @@ def compose_dual_canvas(left_input, right_input):
     return combined
 
 
-def preprocess_image(img_input, stroke_scale: float = 1.0, *, auto_balance: bool = True):
+def preprocess_image(img_input, stroke_scale: float = 1.0):
     ensure_model_loaded()
     img = extract_canvas_array(img_input)
     if img is None:
@@ -301,14 +301,14 @@ def preprocess_image(img_input, stroke_scale: float = 1.0, *, auto_balance: bool
     arr_resized = np.clip(arr_resized * stroke_scale, 0.0, 1.0)
 
     auto_balance_scale = 1.0
-    balanced_mass_fraction = float(arr_resized.sum() / (TARGET_HEIGHT * TARGET_WIDTH))
-    if auto_balance:
-        target_mass = sum(METRIC_TARGETS["mass_fraction"]) / 2.0
-        arr_resized, auto_balance_scale, balanced_mass_fraction = _auto_balance_stroke(
-            arr_resized,
-            target_mass_fraction=target_mass,
-            clamp=(0.7, 1.4),
-        )
+    # Match the dataset's global mean intensity (more faithful than a fixed midpoint)
+    pre_balance_mass_fraction = float(arr_resized.mean())
+    target_mass = float(mean.mean())
+    arr_resized, auto_balance_scale, balanced_mass_fraction = _auto_balance_stroke(
+        arr_resized,
+        target_mass_fraction=target_mass,
+        clamp=(0.6, 1.6),
+    )
 
     # Light recentering by center-of-mass to reduce sensitivity to placement
     mass = arr_resized
@@ -345,9 +345,11 @@ def preprocess_image(img_input, stroke_scale: float = 1.0, *, auto_balance: bool
         std_safe,
     )
     diagnostics["applied_auto_balance"] = {
-        "enabled": bool(auto_balance),
+        "enabled": True,
         "scale": float(auto_balance_scale),
         "mass_fraction_after": float(balanced_mass_fraction),
+        "mass_fraction_before": float(pre_balance_mass_fraction),
+        "target_mass_fraction": float(target_mass),
     }
 
     return augmented_standardized, arr_centered, mean_diff_uint8, diagnostics
@@ -521,7 +523,7 @@ def enrich_diagnostics(stats, probs):
     return stats
 
 
-def predict_number(left_canvas, right_canvas, stroke_scale, auto_balance):
+def predict_number(left_canvas, right_canvas, stroke_scale):
     ensure_model_loaded()
     combined_canvas = compose_dual_canvas(left_canvas, right_canvas)
     if combined_canvas is None:
@@ -534,7 +536,6 @@ def predict_number(left_canvas, right_canvas, stroke_scale, auto_balance):
     result = preprocess_image(
         combined_canvas,
         stroke_scale=stroke_scale,
-        auto_balance=bool(auto_balance),
     )
     if result is None:
         blank_probs = {f"{i:02d}": 0.0 for i in range(OUTPUT_CLASSES)}
@@ -582,11 +583,6 @@ with gr.Blocks() as demo:
                 step=0.05,
                 label="Stroke Intensity (scale)",
             )
-            auto_balance = gr.Checkbox(
-                value=True,
-                label="Auto Balance Stroke Thickness",
-                info="Automatically rescales the digit to match training mass and brightness.",
-            )
 
         with gr.Column(scale=1):
             pred_box = gr.Number(label="Predicted Number", precision=0, value=None)
@@ -605,7 +601,6 @@ with gr.Blocks() as demo:
                     left_canvas,
                     right_canvas,
                     stroke_slider,
-                    auto_balance,
                     pred_box,
                     prob_table,
                     preview,
@@ -616,19 +611,19 @@ with gr.Blocks() as demo:
 
     predict_btn.click(
         fn=predict_number,
-        inputs=[left_canvas, right_canvas, stroke_slider, auto_balance],
+        inputs=[left_canvas, right_canvas, stroke_slider],
         outputs=[pred_box, prob_table, preview, mean_diff_view, diagnostics_box],
     )
     # On Spaces, avoid per-stroke inference to prevent event floods
     if not IS_SPACE:
         left_canvas.change(
             fn=predict_number,
-            inputs=[left_canvas, right_canvas, stroke_slider, auto_balance],
+            inputs=[left_canvas, right_canvas, stroke_slider],
             outputs=[pred_box, prob_table, preview, mean_diff_view, diagnostics_box],
         )
         right_canvas.change(
             fn=predict_number,
-            inputs=[left_canvas, right_canvas, stroke_slider, auto_balance],
+            inputs=[left_canvas, right_canvas, stroke_slider],
             outputs=[pred_box, prob_table, preview, mean_diff_view, diagnostics_box],
         )