sweep fix + reset

pbc3_sweep.py

@@ -30,68 +30,55 @@ JPEG_QUALITIES = (1, 3, 5, 10, 20, 40, 70, 95)
 AVIF_QUALITIES = (0, 1, 3, 5, 10, 20, 40, 70, 95)
 WEBP_QUALITIES = (0, 1, 3, 5, 10, 20, 40, 70, 95)
 
-# Optimizer search space. (name, kind, *args). Ranges are a starting guess — tweak freely.
 SEARCH_SPACE = [
-    ("patch_count", "int", 1, 250, True), ("search_depth", "int", 20, 1000, True),
-    ("proposal_depth", "int", 5, 200, False), ("exact_depth", "int", 1, 60, False),
-    ("top_k", "cat", [4, 8, 16, 32, 64], False), ("anchor_block_size", "int", 2, 32, False),
-    ("cell_sizes_per_candidate", "int", 1, 6, False),
-    ("search_q_start", "float", 0.0, 1.0), ("search_q_end", "float", 0.0, 1.0),
-    ("q_init", "float", 0.0, 1.0), ("q_start", "float", 0.0, 1.0), ("q_end", "float", 0.0, 1.0),
-    ("patch_palette_bitcount", "int", 1, 5, False),
+    ("patch_count", "cat", [1, 5, 10, 25, 50, 100, 200], True),
+    ("search_depth", "cat", [20, 50, 100, 500, 1000], True),
+    ("proposal_depth", "cat", [1, 5, 10, 20, 50, 100, 200], False),
+    ("exact_depth", "cat", [1, 10, 20, 30, 50], False),
+    ("top_k", "cat", [4, 8, 16, 32, 64], False),
+    ("anchor_block_size", "cat", [2, 4, 8, 16, 32], False),
+    ("cell_sizes_per_candidate", "cat", [1, 2, 3, 4], False),
+    ("search_q_start", "cat", [0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0], False),
+    ("search_q_end", "cat", [0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0], False),
+    ("q_init", "cat", [0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0], False),
+    ("q_start", "cat", [0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0], False),
+    ("q_end", "cat", [0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0], False),
+    ("patch_palette_bitcount", "cat", [1, 2, 3, 4, 5], False),
     ("mask_size", "cat", [4, 6, 8, 16], False),
 ]
 
 _DEFAULTS = PBC3Config().__dict__
-
 PRESETS = ("speed", "balanced", "quality", "compression", "high_quality")
 
-def ensure_pbc_baselines(study, images):
-    have = {t.user_attrs.get("baseline") for t in study.trials}
-    for name in PRESETS:
-        tag = f"preset_{name}"
-        if tag in have:
-            continue
-        _log(f"PBC preset baseline: {name} started")
-        cfg = dict(vars(getattr(PBC3Config, name)()))
-        try:
-            _new_bar()
-            rows = _eval_pbc3(cfg, images, on_image=_bar_tick)
-        except Exception as e:
-            _log(f"preset {name} failed: {e}")
-            continue
-        values = _avg(rows)
-        _add_completed(study, values,
-                       {"kind": "pbc3", "baseline": tag, "preset": name,
-                        "config": cfg, "per_image": rows})
-        _log_result(f"PBC preset baseline: {name} done", values, _is_pareto_candidate(study, values))
+_LOCK = threading.Lock()
+_STOP = threading.Event()
+_THREAD = None
+_STATE = {"running": False, "mode": None, "done": 0, "total": None,
+          "current": None, "started": None, "error": None, "log": []}
 
 
-def suggest_config(trial):
-    cfg = {}
-    for name, kind, *a in SEARCH_SPACE:
-        if kind == "int":
-            cfg[name] = trial.suggest_int(name, a[0], a[1], log=a[2])
-        elif kind == "float":
-            cfg[name] = trial.suggest_float(name, a[0], a[1])
-        else:
-            cfg[name] = trial.suggest_categorical(name, a[0])
-    return cfg
+def _log(msg):
+    with _LOCK:
+        _STATE["log"] = (_STATE["log"] + [f"{time.strftime('%H:%M:%S')} {msg}"])[-60:]
+    print(f"[sweep] {msg}", flush=True)
 
 
-def _coerce(params):
-    out = {}
-    for k, v in params.items():
-        d = _DEFAULTS.get(k)
-        if isinstance(d, bool):
-            out[k] = v if isinstance(v, bool) else str(v).lower() in ("true", "1")
-        elif isinstance(d, int) and not isinstance(d, bool):
-            out[k] = int(v)
-        elif isinstance(d, float):
-            out[k] = float(v)
-        else:
-            out[k] = v
-    return out
+def _new_bar():
+    with _LOCK:
+        _STATE["log"] = (_STATE["log"] + [""])[-60:]
+
+
+def _bar_tick():
+    with _LOCK:
+        if _STATE["log"]:
+            _STATE["log"][-1] += "|"
+
+
+def status():
+    with _LOCK:
+        s = dict(_STATE)
+        s["log"] = list(_STATE["log"])
+        return s
 
 
 def dataset():
@@ -113,10 +100,16 @@ def _avg(rows):
     return [float(np.mean([r[k] for r in rows])) for k in ("seconds", "bpp", "mse")]
 
 
+def _is_pbc3_trial(t):
+    cfg = t.user_attrs.get("config") or {}
+    kind = t.user_attrs.get("kind") or ("baseline" if cfg.get("codec") else "pbc3")
+    return kind == "pbc3"
+
+
 def _is_pareto_candidate(study, values):
     vals = np.asarray(values, dtype=np.float64)
     for t in study.trials:
-        if t.state != optuna.trial.TrialState.COMPLETE or not t.values or len(t.values) != 3:
+        if not _is_pbc3_trial(t) or t.state != optuna.trial.TrialState.COMPLETE or not t.values or len(t.values) != 3:
             continue
         other = np.asarray(t.values, dtype=np.float64)
         if np.all(other <= vals) and np.any(other < vals):
@@ -124,9 +117,39 @@ def _is_pareto_candidate(study, values):
     return True
 
 
-def _log_result(prefix, values, pareto):
+def _log_result(prefix, values, pareto=None):
     seconds, bpp, mse = values
-    _log(f"{prefix} | MSE {mse:.3f} | bpp {bpp:.5f} | speed {seconds:.3f}s | pareto {'yes' if pareto else 'no'}")
+    msg = f"{prefix} | MSE {mse:.3f} | bpp {bpp:.5f} | speed {seconds:.3f}s"
+    if pareto is not None:
+        msg += f" | pareto {'yes' if pareto else 'no'}"
+    _log(msg)
+
+
+def suggest_config(trial):
+    cfg = {}
+    for name, kind, *a in SEARCH_SPACE:
+        if kind == "int":
+            cfg[name] = trial.suggest_int(name, a[0], a[1], log=a[2])
+        elif kind == "float":
+            cfg[name] = trial.suggest_float(name, a[0], a[1])
+        else:
+            cfg[name] = trial.suggest_categorical(name, a[0])
+    return cfg
+
+
+def _coerce(params):
+    out = {}
+    for k, v in params.items():
+        d = _DEFAULTS.get(k)
+        if isinstance(d, bool):
+            out[k] = v if isinstance(v, bool) else str(v).lower() in ("true", "1")
+        elif isinstance(d, int) and not isinstance(d, bool):
+            out[k] = int(v)
+        elif isinstance(d, float):
+            out[k] = float(v)
+        else:
+            out[k] = v
+    return out
 
 
 def _eval_pbc3(cfg, images, on_image=None):
@@ -146,6 +169,8 @@ def _eval_pbc3(cfg, images, on_image=None):
 def _eval_codec(fmt, q, images, on_image=None):
     rows = []
     for im in images:
+        if _STOP.is_set():
+            raise InterruptedError("sweep stopped")
         arr = im["arr"]
         pixels = arr.shape[0] * arr.shape[1]
         buf = io.BytesIO()
@@ -180,6 +205,8 @@ def ensure_baselines(study, images):
     have = {t.user_attrs.get("baseline") for t in study.trials}
     for fmt, qs in (("JPEG", JPEG_QUALITIES), ("AVIF", AVIF_QUALITIES), ("WEBP", WEBP_QUALITIES)):
         for q in qs:
+            if _STOP.is_set():
+                return
             tag = f"{fmt.lower()}_q{q}"
             if tag in have:
                 continue
@@ -187,6 +214,9 @@ def ensure_baselines(study, images):
             try:
                 _new_bar()
                 rows = _eval_codec(fmt, q, images, on_image=_bar_tick)
+            except InterruptedError:
+                _log(f"baseline {tag} stopped")
+                return
             except Exception as e:
                 _log(f"baseline {tag} failed: {e}")
                 continue
@@ -194,41 +224,30 @@ def ensure_baselines(study, images):
             _add_completed(study, values, {
                 "baseline": tag, "kind": "baseline",
                 "config": {"codec": fmt, "q": q}, "per_image": rows})
-            _log_result(f"baseline {tag} done", values, _is_pareto_candidate(study, values))
-
+            _log_result(f"baseline {tag} done", values)
 
-# ---------------- runner state / control ----------------
-_LOCK = threading.Lock()
-_STOP = threading.Event()
-_THREAD = None
-_STATE = {"running": False, "mode": None, "done": 0, "total": None,
-          "current": None, "started": None, "error": None, "log": []}
 
-
-def _log(msg):
-    with _LOCK:
-        _STATE["log"] = (_STATE["log"] + [f"{time.strftime('%H:%M:%S')} {msg}"])[-60:]
-    print(f"[sweep] {msg}", flush=True)
-
-
-def _new_bar():
-    """Append a fresh, empty progress line that _bar_tick grows in place."""
-    with _LOCK:
-        _STATE["log"] = (_STATE["log"] + [""])[-60:]
-
-
-def _bar_tick():
-    """Add one '|' to the current progress line (called once per image)."""
-    with _LOCK:
-        if _STATE["log"]:
-            _STATE["log"][-1] += "|"
-
-
-def status():
-    with _LOCK:
-        s = dict(_STATE)
-        s["log"] = list(_STATE["log"])
-        return s
+def ensure_pbc_baselines(study, images):
+    have = {t.user_attrs.get("baseline") for t in study.trials}
+    for name in PRESETS:
+        if _STOP.is_set():
+            return
+        tag = f"preset_{name}"
+        if tag in have:
+            continue
+        _log(f"PBC preset baseline: {name} started")
+        cfg = dict(vars(getattr(PBC3Config, name)()))
+        try:
+            _new_bar()
+            rows = _eval_pbc3(cfg, images, on_image=_bar_tick)
+        except Exception as e:
+            _log(f"preset {name} failed: {e}")
+            continue
+        values = _avg(rows)
+        _add_completed(study, values,
+                       {"kind": "pbc3", "baseline": tag, "preset": name,
+                        "config": cfg, "per_image": rows})
+        _log_result(f"PBC preset baseline: {name} done", values, _is_pareto_candidate(study, values))
 
 
 def _grid_trials(spec):
@@ -244,7 +263,7 @@ def start(mode, spec=None):
         if _STATE["running"]:
             return {"error": "A sweep is already running."}
     if not os.path.isdir(DATA_DIR):
-        return {"error": f"Dataset folder missing: hpt_data/"}
+        return {"error": "Dataset folder missing: hpt_data/"}
     _STOP.clear()
     _THREAD = threading.Thread(target=_run, args=(mode, spec or {}), daemon=True)
     _THREAD.start()
@@ -267,6 +286,8 @@ def _run(mode, spec):
         _log(f"loaded {len(images)} images")
         study = get_study()
         ensure_baselines(study, images)
+        if _STOP.is_set():
+            return
         ensure_pbc_baselines(study, images)
         if _STOP.is_set():
             return

server.py

@@ -516,6 +516,9 @@ def _pbc3_trials(study, mp_min=None, mp_max=None):
                  if t.state == optuna.trial.TrialState.COMPLETE and t.values and len(t.values) == 3]
     rows = []
     for t in completed:
+        cfg = t.user_attrs.get("config") or {}
+        kind = t.user_attrs.get("kind") or ("baseline" if cfg.get("codec") else "pbc3")
+
         if mp_min is not None or mp_max is not None:
             sub = [r for r in (t.user_attrs.get("per_image") or []) if _in_range(r["mp"], mp_min, mp_max)]
             if not sub:
@@ -525,22 +528,24 @@ def _pbc3_trials(study, mp_min=None, mp_max=None):
             mse = float(np.mean([r["mse"] for r in sub]))
         else:
             seconds, bpp, mse = t.values
-        rows.append((t, seconds, bpp, mse))
 
-    if rows:
-        orient = np.array([[-s, -b, -m] for _, s, b, m in rows], dtype=np.float64)  # minimize all three
+        rows.append((t, seconds, bpp, mse, kind, cfg))
+
+    pareto = [False] * len(rows)
+    pbc_rows = [(i, s, b, m) for i, (_, s, b, m, kind, _) in enumerate(rows) if kind == "pbc3"]
+    if pbc_rows:
+        orient = np.array([[-s, -b, -m] for _, s, b, m in pbc_rows], dtype=np.float64)
         mask = _pareto_mask(orient)
-    else:
-        mask = []
+        for (i, _, _, _), is_pareto in zip(pbc_rows, mask):
+            pareto[i] = bool(is_pareto)
 
     out = []
-    for i, (t, s, b, m) in enumerate(rows):
-        cfg = t.user_attrs.get("config") or {}
+    for i, (t, s, b, m, kind, cfg) in enumerate(rows):
         out.append({
-            "number": t.number, "speed": s, "bpp": b, "mse": m,   # "speed" key carries seconds
-            "pareto": bool(mask[i]), "baseline": t.user_attrs.get("baseline"),
+            "number": t.number, "speed": s, "bpp": b, "mse": m,
+            "pareto": pareto[i], "baseline": t.user_attrs.get("baseline"),
             "preset": t.user_attrs.get("preset"),
-            "kind": t.user_attrs.get("kind") or ("baseline" if cfg.get("codec") else "pbc3"),
+            "kind": kind,
             "codec": cfg.get("codec"), "q": cfg.get("q"), "params": cfg,
         })
     return out