Upload pipeline.py

pipeline.py

@@ -922,6 +922,16 @@ class Q2KQuantizer:
                     stop_event=stop_event,
                     device=dev_str,
                 )
+                if seg_perm is not None:
+                    from .segmentation import band_labels
+
+                    Xh = np.stack([b.dequantize() for b in res.blocks]).astype(np.float64)
+                    labels_slot = np.take_along_axis(seg_labels, seg_perm, axis=1)
+                    labels_rec = band_labels(Xh, max(2, int(cfg.segment_k)))
+                    res.stats.segment_rounds = seg_rounds
+                    res.stats.segment_flips = int(np.count_nonzero(labels_rec != labels_slot))
+                    res.segment_perm = seg_perm
+                    res.segment_labels = seg_labels
                 if cfg.verbose:
                     _report(res.stats)
                 return res