# -*- coding: utf-8 -*-
from __future__ import annotations
import json
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

import numpy as np
import onnxruntime as ort
from PIL import Image

import os, numpy as np

DEBUG = os.environ.get("ALAMI_DEBUG") == "1"
def dprint(msg: str = ""):
    if DEBUG:
        print(f"[ALAMI-DEBUG] {msg}", flush=True)


# ------------------------- Small math utils -------------------------

def _sigmoid(x: np.ndarray) -> np.ndarray:
    return 1.0 / (1.0 + np.exp(-x))

def _logit(p: np.ndarray, eps: float = 1e-8) -> np.ndarray:
    p = np.clip(p, eps, 1 - eps)
    return np.log(p) - np.log(1 - p)

def _nms_xyxy_classwise(
    boxes: np.ndarray,  # [N,4]
    scores: np.ndarray, # [N]
    classes: np.ndarray,# [N] int
    iou_thr: float,
    max_det: int
) -> np.ndarray:
    """Gibt Indizes der behaltenen Detections zurück (class-wise Greedy NMS)."""
    keep: List[int] = []
    for c in np.unique(classes):
        idx = np.where(classes == c)[0]
        if idx.size == 0:
            continue
        b = boxes[idx]
        s = scores[idx]
        order = np.argsort(-s)
        idx = idx[order]
        b = b[order]

        while idx.size > 0:
            i = idx[0]
            keep.append(i)
            if len(keep) >= max_det:
                break
            if idx.size == 1:
                break
            iou = _iou_batch_xyxy(b[0], b[1:])
            remain = np.where(iou <= iou_thr)[0] + 1
            idx = idx[remain]
            b = b[remain]
        if len(keep) >= max_det:
            break
    return np.array(keep, dtype=np.int64)

def _iou_batch_xyxy(a: np.ndarray, b: np.ndarray) -> np.ndarray:
    """
    IoU eines einzelnen Kastens a gegen viele b. a: (4,), b: (M,4)
    """
    ax1, ay1, ax2, ay2 = a
    bx1, by1, bx2, by2 = b[:, 0], b[:, 1], b[:, 2], b[:, 3]
    ix1 = np.maximum(ax1, bx1)
    iy1 = np.maximum(ay1, by1)
    ix2 = np.minimum(ax2, bx2)
    iy2 = np.minimum(ay2, by2)
    iw = np.maximum(0.0, ix2 - ix1)
    ih = np.maximum(0.0, iy2 - iy1)
    inter = iw * ih
    area_a = np.maximum(0.0, (ax2 - ax1)) * np.maximum(0.0, (ay2 - ay1))
    area_b = np.maximum(0.0, (bx2 - bx1)) * np.maximum(0.0, (by2 - by1))
    union = area_a + area_b - inter + 1e-9
    return inter / union

# ------------------------- Main class -------------------------

class ModelBundle:
    def __init__(self, bundle_dir: Path):
        self.bundle_dir = bundle_dir
        self.onnx = bundle_dir / "model.onnx"
        if not self.onnx.exists():
            raise FileNotFoundError(f"ONNX not found: {self.onnx}")

        # names.json is authoritative; fall back to model_card.json dataset classes
        # so a bundle without names.json still serves instead of crashing at boot.
        self.names = self._read_json("names.json", required=False)
        if self.names is None:
            card = self._read_json("model_card.json", required=False) or {}
            self.names = (card.get("dataset") or {}).get("classes")
        if not self.names:
            raise FileNotFoundError(
                f"Neither names.json nor model_card.json dataset.classes found in {bundle_dir}"
            )
    # Normalize names -> list (if dict)
        if isinstance(self.names, dict):
            try:
                keys = [int(k) for k in self.names.keys()] if self.names else []
                arr = [None] * (max(keys) + 1 if keys else 0)
                for k, v in self.names.items():
                    arr[int(k)] = v
                self.names = arr
            except Exception:
                self.names = list(self.names.values())
        self.names = [("" if n is None else str(n)) for n in self.names]

        self.post_cfg = self._read_json("postprocess_config.json", required=True)
        self.calibration = self.post_cfg.get("calibration") or None
        self.conf_thr = float(self.post_cfg.get("confidence_threshold", 0.25))
        self.iou_thr = float(self.post_cfg.get("iou_threshold", 0.50))
        self.max_det = int(self.post_cfg.get("max_detections", 300))

        # preprocessing mode (from postprocess_config.json), default 'letterbox'
        self.preprocess_mode = str(self.post_cfg.get("preprocess", "letterbox")).lower()
        if self.preprocess_mode not in ("letterbox", "resize"):
            self.preprocess_mode = "letterbox"

        # holds last affine used during load_image(); consumed by back-projection
        self._affine: Optional[Dict[str, float]] = None

        providers = ort.get_available_providers()
        # prefer CUDA if available, fallback CPU
        if "CUDAExecutionProvider" in providers:
            self.session = ort.InferenceSession(
                str(self.onnx),
                providers=["CUDAExecutionProvider", "CPUExecutionProvider"]
            )
        else:
            self.session = ort.InferenceSession(
                str(self.onnx),
                providers=["CPUExecutionProvider"]
            )

        # input / output info
        io = self.session.get_inputs()[0]
        self.input_name = io.name
        self.input_shape = tuple(io.shape)  # [batch, ch, h, w] (may be dynamic)
        self.imgsz = self._infer_imgsz(self.input_shape)

    # Try to identify the two main outputs:
    # - preds: [1, N, 4 + nc + nm]  (N varies)
    # - proto: [1, mask_dim(=32), H/4, W/4]
        self.pred_out_name, self.proto_out_name = self._resolve_output_names()

    # ---------- IO ----------

    def _read_json(self, name: str, required: bool = False) -> Any:
        p = self.bundle_dir / name
        if not p.exists():
            if required:
                raise FileNotFoundError(p)
            return None
        return json.loads(p.read_text(encoding="utf-8"))

    @staticmethod
    def _infer_imgsz(shape: Tuple[Any, ...]) -> int:
        # YOLOv8 standard: [batch, 3, H, W]
        try:
            h = int(shape[2]) if shape[2] is not None else 640
            w = int(shape[3]) if shape[3] is not None else 640
            assert h == w
            return h
        except Exception:
            return 640

    # ---------- preprocessing ----------

    def load_image(self, img_path: Path) -> Tuple[np.ndarray, Tuple[int, int, float, int, int]]:
        """
                Load image and perform letterbox resize to self.imgsz.
                Return:
                    - Tensor [1,3,H,W] float32
                    - Meta: (w0, h0, r, pad_w, pad_h) for back-projection
        """
        dprint("load_image() called")
        img = Image.open(img_path).convert("RGB")
        w0, h0 = img.size
        dprint(f"orig_size=(w0={w0}, h0={h0}), preprocess_mode={self.preprocess_mode}")

        if self.preprocess_mode == "letterbox":
            # letterbox to square imgsz
            r = min(self.imgsz / h0, self.imgsz / w0)
            nw, nh = int(round(w0 * r)), int(round(h0 * r))
            img_resized = img.resize((nw, nh), Image.BILINEAR)
            canvas = Image.new("RGB", (self.imgsz, self.imgsz), (114, 114, 114))
            pad_w, pad_h = (self.imgsz - nw) // 2, (self.imgsz - nh) // 2
            dprint(f"letterbox: r={r:.6f}, nw={nw}, nh={nh}, pad_w={pad_w}, pad_h={pad_h}")
            canvas.paste(img_resized, (pad_w, pad_h))
            arr = np.asarray(canvas).astype(np.float32)
            # store affine for back-projection
            self._affine = {"mode": "letterbox", "w0": w0, "h0": h0, "r": r, "pad_w": pad_w, "pad_h": pad_h}
        else:
            # plain resize (no padding) to (imgsz, imgsz)
            img_resized = img.resize((self.imgsz, self.imgsz), Image.BILINEAR)
            arr = np.asarray(img_resized).astype(np.float32)
            sx = w0 / float(self.imgsz)
            sy = h0 / float(self.imgsz)
            dprint(f"resize: sx={w0/float(self.imgsz):.6f}, sy={h0/float(self.imgsz):.6f}")
            # store affine for back-projection
            self._affine = {"mode": "resize", "w0": w0, "h0": h0, "sx": sx, "sy": sy}

        arr = arr.transpose(2, 0, 1) / 255.0  # [3,H,W], 0..1
        arr = np.expand_dims(arr, 0)          # [1,3,H,W]
        # keep meta tuple for backward-compat (letterbox values; unused for 'resize')
        meta = (w0, h0, self._affine.get("r", 1.0), self._affine.get("pad_w", 0), self._affine.get("pad_h", 0))
        dprint(f"tensor_shape={arr.shape}, meta={meta}")
        return arr, meta

    # ---------- inference (raw) ----------

    def infer(self, img_tensor: np.ndarray) -> Dict[str, np.ndarray]:
        """
    Raw ONNX outputs. Kept return contract for backward compatibility.
        """
        outputs = self.session.run(None, {self.input_name: img_tensor})
        out = {}
        for i, o in enumerate(outputs):
            out[f"out{i}"] = o
        return out

    # ---------- high-level prediction ----------

    def predict(
        self,
        image_path: Path,
        return_masks: bool = True,
        mask_threshold: float = 0.5
    ) -> Dict[str, Any]:
        """
    Run end-to-end inference incl. postprocessing.
    Return compatible to val_predictions.json:
        {
          "path": <str>,
          "orig_shape": [h, w],
          "boxes": [{"xyxy":[x1,y1,x2,y2], "cls":int, "conf":float}, ...],
          "masks": Optional[List[np.ndarray or None]]  # binary HxW (when return_masks=True)
        }
        """
        tensor, meta = self.load_image(image_path)
        raw = self.session.run(None, {self.input_name: tensor})

        outs = self.session.get_outputs()
        for i, arr in enumerate(raw):
            dprint(f"onnx_out{i}: name={outs[i].name}, shape={arr.shape}, ndim={arr.ndim}, dtype={arr.dtype}")
            # nur sehr kleine Kostprobe loggen (erste 2x5 Werte flach)
            flat = arr.ravel()
            sample = np.array2string(flat[:10], precision=4, suppress_small=True)
            dprint(f"onnx_out{i}_sample={sample}")

        preds, proto = self._pick_preds_and_proto(raw)

        boxes, scores, clses, mask_coef = self._decode_preds(preds)  # imgsz-Koords
        if boxes.size == 0:
            return {
                "path": str(image_path),
                "orig_shape": [meta[1], meta[0]],
                "boxes": [],
                "masks": None
            }

        # Temperature scaling (optional)
        if self.calibration and "temperature" in self.calibration:
            T = float(self.calibration.get("temperature", 1.0))
            # numerically stable: sigmoid(logit(p)/T)
            scores = _sigmoid(_logit(scores) / max(T, 1e-9))

        # Threshold
        # mask_coef ist None bei reinen DETEKTIONS-Bundles (kein Seg-Kopf) — z. B.
        # dem Scene-Gate-Bundle. Nur indizieren, wenn es existiert.
        th_mask = scores >= self.conf_thr
        boxes, scores, clses = boxes[th_mask], scores[th_mask], clses[th_mask]
        mask_coef = mask_coef[th_mask] if mask_coef is not None else None
        if boxes.size == 0:
            return {
                "path": str(image_path),
                "orig_shape": [meta[1], meta[0]],
                "boxes": [],
                "masks": None
            }
        dprint(f"after conf_thr({self.conf_thr}): kept={boxes.shape[0]}")

      # NMS (class-wise)
        keep = _nms_xyxy_classwise(boxes, scores, clses, self.iou_thr, self.max_det)
        boxes, scores, clses = boxes[keep], scores[keep], clses[keep]
        mask_coef = mask_coef[keep] if mask_coef is not None else None

        dprint(f"after NMS(iou={self.iou_thr}): kept={boxes.shape[0]}")

      # Back-project to original image
      # boxes = self._unletterbox_boxes(boxes, meta)
        boxes = self._backproject_boxes(boxes, meta)

      # Reconstruct masks (optional)
        masks_out = None
        if return_masks and proto is not None and mask_coef is not None and mask_coef.size > 0:
            masks_out = self._reconstruct_masks(proto, mask_coef, boxes, meta, mask_threshold)

        result_boxes = [
            {"xyxy": boxes[i].tolist(), "cls": int(clses[i]), "conf": float(scores[i])}
            for i in range(boxes.shape[0])
        ]

        return {
            "path": str(image_path),
            "orig_shape": [meta[1], meta[0]],  # [h,w]
            "boxes": result_boxes,
            "masks": masks_out  # list of binary HxW arrays (or None)
        }

    # ---------- internes Postprocessing ----------

    def _resolve_output_names(self) -> Tuple[Optional[str], Optional[str]]:
        """
    Try to resolve prediction and proto outputs based on shapes.
        """
        outs = self.session.get_outputs()
        pred_name = None
        proto_name = None
        for o in outs:
            shape = tuple(o.shape)
            # Proto-Kandidaten: 4D, häufig [1, 32, H/4, W/4]
            if len(shape) == 4 and shape[0] in (1, None) and shape[1] and shape[1] >= 16:
                proto_name = o.name if proto_name is None else proto_name
            # Pred-Kandidaten: 3D [1, N, 4+nc+nm]
            if len(shape) == 3 and shape[0] in (1, None) and (shape[2] is None or shape[2] >= 20):
                pred_name = o.name if pred_name is None else pred_name
        return pred_name, proto_name

    def _pick_preds_and_proto(self, outputs: List[np.ndarray]) -> Tuple[np.ndarray, Optional[np.ndarray]]:
        """
    Pick relevant tensors based on resolved names.
    Fallback: heuristic by rank.
        """
        outs = self.session.get_outputs()
        name_to_arr = {outs[i].name: outputs[i] for i in range(len(outs))}
        preds = None
        proto = None
        if self.pred_out_name in name_to_arr:
            preds = name_to_arr[self.pred_out_name]
        if self.proto_out_name in name_to_arr:
            proto = name_to_arr[self.proto_out_name]

        # Heuristik-Fallback
        if preds is None or preds.ndim != 3:
            for arr in outputs:
                if arr.ndim == 3:
                    preds = arr
                    break
        if proto is None:
            for arr in outputs:
                if arr.ndim == 4:
                    proto = arr
                    break
        if preds is None:
            # Last resort: take the first tensor
            preds = outputs[0]

        dprint(f"pick_preds_and_proto: preds_shape={None if preds is None else preds.shape}, "
            f"proto_shape={None if proto is None else proto.shape}")
        if preds is not None and preds.ndim == 3:
            b, a, c = preds.shape
            dprint(f"preds dims: b={b}, a={a}, c={c} (expect [1,N,D])")

        return preds, proto

    def _decode_preds(self, preds: np.ndarray) -> Tuple[np.ndarray, np.ndarray, np.ndarray, Optional[np.ndarray]]:
        """
        Decode YOLOv8-style ONNX head to (boxes_xyxy_imgsz, scores, clses, mask_coef).

        Supports:
        - Seg head without obj: 4 (xywh) + nc + nm
        - Seg head with obj:    4 (xywh) + 1 (obj) + nc + nm
        """
        # Ensure 3D: [1, N, D] or [1, D, N]
        if preds.ndim != 3:
            preds = preds.reshape(1, preds.shape[0], preds.shape[1])

        P = preds[0]  # [N, D] or [D, N]
        dprint(f"_decode_preds: raw P shape={P.shape}")

        # Detect and fix [D, N] -> [N, D]
        if P.shape[0] <= 256 and P.shape[1] >= 1000:
            dprint("P appears to be [D,N]; transposing to [N,D].")
            P = P.T

        dprint(f"_decode_preds: normalized P shape={P.shape}")

        if P.size == 0 or P.shape[1] < 4:
            dprint("No valid prediction channels; returning empty.")
            return (np.zeros((0, 4), np.float32),
                    np.zeros((0,), np.float32),
                    np.zeros((0,), np.int64),
                    None)

        N, D = P.shape
        nc = len(self.names)

        # Boxes are xywh in imgsz space (Ultralytics export)
        xywh = P[:, 0:4].astype(np.float32)
        xywh_max = float(np.nanmax(xywh)) if xywh.size else 0.0
        dprint(f"xywh_max={xywh_max:.4f}, nc={nc}, D={D}")

        if xywh_max <= 1.5:
            xywh *= float(self.imgsz)
            dprint("xywh interpreted as normalized; scaled by imgsz.")

        # Infer layout using known seg pattern (proto channels)
        nm_candidate = D - 4 - nc  # remaining dims after xywh + cls
        obj = None
        mask_coef: Optional[np.ndarray] = None

        if nm_candidate == 32:
            # Typical YOLOv8-seg: 4 + nc + 32 (no obj)
            cls_start = 4
            cls_end = 4 + nc
            mask_start = cls_end
            cls_scores = P[:, cls_start:cls_end].astype(np.float32)
            mask_coef = P[:, mask_start:mask_start + nm_candidate].astype(np.float32)
            obj = np.ones((N, 1), dtype=np.float32)
            dprint(f"layout=xywh+cls+mask (no obj), nm={nm_candidate}")
        elif nm_candidate > 32:
            # Likely: 4 + 1 + nc + nm (with obj)
            obj_idx = 4
            cls_start = 5
            cls_end = 5 + nc
            nm = D - (5 + nc)
            if nm <= 0:
                dprint(f"Inconsistent head layout (nm={nm}); returning empty.")
                return (np.zeros((0, 4), np.float32),
                        np.zeros((0,), np.float32),
                        np.zeros((0,), np.int64),
                        None)
            obj = P[:, obj_idx:obj_idx + 1].astype(np.float32)
            cls_scores = P[:, cls_start:cls_end].astype(np.float32)
            mask_coef = P[:, cls_end:cls_end + nm].astype(np.float32)
            dprint(f"layout=xywh+obj+cls+mask, nm={nm}")
        else:
            # Fallback: assume 4 + nc (+ optional mask), no obj
            if D >= 4 + nc:
                cls_start = 4
                cls_end = 4 + nc
                nm = max(0, D - (4 + nc))
                cls_scores = P[:, cls_start:cls_end].astype(np.float32)
                mask_coef = P[:, cls_end:cls_end + nm].astype(np.float32) if nm > 0 else None
                obj = np.ones((N, 1), dtype=np.float32)
                dprint(f"layout=fallback xywh+cls(+mask), nm={nm}")
            else:
                dprint(f"Unexpected head layout: D={D}, nc={nc}, nm_candidate={nm_candidate}; returning empty.")
                return (np.zeros((0, 4), np.float32),
                        np.zeros((0,), np.float32),
                        np.zeros((0,), np.int64),
                        None)

        # If logits detected, apply sigmoid
        if obj is not None and (obj.max() > 1.0 or obj.min() < 0.0):
            dprint("objectness appears to be logits; applying sigmoid.")
            obj = _sigmoid(obj)

        if cls_scores.max() > 1.0 or cls_scores.min() < 0.0:
            dprint("class scores appear to be logits; applying sigmoid.")
            cls_scores = _sigmoid(cls_scores)

        # Clip to [0,1] after sigmoid
        obj = np.clip(obj, 0.0, 1.0)
        cls_scores = np.clip(cls_scores, 0.0, 1.0)

        # xywh -> xyxy in imgsz space
        x, y, w, h = xywh.T
        x1 = x - w / 2.0
        y1 = y - h / 2.0
        x2 = x + w / 2.0
        y2 = y + h / 2.0
        boxes_xyxy = np.stack([x1, y1, x2, y2], axis=1).astype(np.float32)

        # Clip to imgsz
        boxes_xyxy[:, 0] = np.clip(boxes_xyxy[:, 0], 0, self.imgsz)
        boxes_xyxy[:, 1] = np.clip(boxes_xyxy[:, 1], 0, self.imgsz)
        boxes_xyxy[:, 2] = np.clip(boxes_xyxy[:, 2], 0, self.imgsz)
        boxes_xyxy[:, 3] = np.clip(boxes_xyxy[:, 3], 0, self.imgsz)

        # Final scores
        clses = np.argmax(cls_scores, axis=1).astype(np.int64)
        max_cls = cls_scores[np.arange(N), clses]
        scores = (obj.flatten() * max_cls).astype(np.float32)

        if boxes_xyxy.size:
            w_box = boxes_xyxy[:, 2] - boxes_xyxy[:, 0]
            h_box = boxes_xyxy[:, 3] - boxes_xyxy[:, 1]
            dprint(
                f"boxes_xyxy stats: median_w={float(np.median(w_box)):.2f}, "
                f"median_h={float(np.median(h_box)):.2f}, "
                f"zeros_w={int(np.sum(w_box <= 1e-3))}"
            )
            dprint(
                f"scores stats: min={float(scores.min()):.4f}, "
                f"max={float(scores.max()):.4f}, "
                f"mean={float(scores.mean()):.4f}"
            )

        return boxes_xyxy, scores, clses, mask_coef

    def _unletterbox_boxes(self, boxes_imgsz: np.ndarray, meta: Tuple[int, int, float, int, int]) -> np.ndarray:
        """
    Transform boxes from imgsz-space (letterbox) back to original image (w0,h0).
        """
        w0, h0, r, pad_w, pad_h = meta
    # Remove padding and scale back
        boxes = boxes_imgsz.copy()
        boxes[:, [0, 2]] -= pad_w
        boxes[:, [1, 3]] -= pad_h
        boxes /= max(r, 1e-9)
        # clamp
        boxes[:, 0] = np.clip(boxes[:, 0], 0, w0)
        boxes[:, 2] = np.clip(boxes[:, 2], 0, w0)
        boxes[:, 1] = np.clip(boxes[:, 1], 0, h0)
        boxes[:, 3] = np.clip(boxes[:, 3], 0, h0)
        return boxes
    
    def _backproject_boxes(self, boxes_imgsz: np.ndarray, meta: Tuple[int, int, float, int, int]) -> np.ndarray:
        """
        Project boxes from imgsz-space back to original image using last used affine.
        Supports modes: 'letterbox' and 'resize'.
        """
        if self._affine is None:
            # fallback: behave like old letterbox
            return self._unletterbox_boxes(boxes_imgsz, meta)

        mode = self._affine.get("mode", "letterbox")
        w0 = float(self._affine.get("w0", meta[0]))
        h0 = float(self._affine.get("h0", meta[1]))

        boxes = boxes_imgsz.copy().astype(np.float32)

        if mode == "letterbox":
            r = float(self._affine.get("r", meta[2]))
            pad_w = float(self._affine.get("pad_w", meta[3]))
            pad_h = float(self._affine.get("pad_h", meta[4]))
            boxes[:, [0, 2]] -= pad_w
            boxes[:, [1, 3]] -= pad_h
            boxes /= max(r, 1e-9)
        else:
            # plain resize back-projection
            sx = float(self._affine.get("sx", w0 / float(self.imgsz)))
            sy = float(self._affine.get("sy", h0 / float(self.imgsz)))
            boxes[:, [0, 2]] *= sx
            boxes[:, [1, 3]] *= sy

        dprint(f"backproject: mode={self._affine.get('mode','?')}, affine={self._affine}")
        if boxes_imgsz.size:
            dprint(f"pre-backproj sample[0]={np.array2string(boxes_imgsz[0], precision=2)}")

        # clamp
        boxes[:, 0] = np.clip(boxes[:, 0], 0, w0)
        boxes[:, 2] = np.clip(boxes[:, 2], 0, w0)
        boxes[:, 1] = np.clip(boxes[:, 1], 0, h0)
        boxes[:, 3] = np.clip(boxes[:, 3], 0, h0)
 
        if boxes.size:
            w = boxes[:, 2] - boxes[:, 0]
            h = boxes[:, 3] - boxes[:, 1]
            dprint(f"post-backproj sample[0]={np.array2string(boxes[0], precision=2)}, "
                f"median_w={np.median(w):.2f}, median_h={np.median(h):.2f}, zeros_w={int(np.sum(w<=1e-3))}")
    
        return boxes

    def _reconstruct_masks(
        self,
        proto: np.ndarray,          # [1, c, mh, mw]
        mask_coef: np.ndarray,      # [K, c]
        boxes_xyxy: np.ndarray,     # [K, 4] in Originalbild-Koords
        meta: Tuple[int, int, float, int, int],
        thr: float
    ) -> List[Optional[np.ndarray]]:
        """
    Reconstruct binary masks in original image space (H=h0, W=w0).
    Simplified implementation (simpler than Ultralytics' ROI rasterization).
        """
        w0, h0, r, pad_w, pad_h = meta
        # proto -> (c, mh, mw)
        p = proto[0]
        # (c, mh, mw) -> (mh, mw, c)
        p = np.transpose(p, (1, 2, 0))  # [mh, mw, c]
        mh, mw, cdim = p.shape
        if mask_coef.shape[1] != cdim:
            # incompatible dimension, skip masks
            return [None] * boxes_xyxy.shape[0]

        # lineare Kombi
        # logits: [mh, mw, K] = p @ mask_coef^T
        logits = np.tensordot(p, mask_coef.T, axes=([2], [0]))  # [mh, mw, K]
        probs = _sigmoid(logits)

    # upscale to imgsz
    # (mh,mw) ~ imgsz/4; scale to imgsz, then remove letterbox, then to (h0,w0)
        probs = np.transpose(probs, (2, 0, 1))  # [K, mh, mw]
        masks_imgsz = []
        for k in range(probs.shape[0]):
            mask_k = Image.fromarray((probs[k] * 255).astype(np.uint8), mode="L")
            mask_k = mask_k.resize((self.imgsz, self.imgsz), Image.BILINEAR)
            # Remove letterbox
            canvas = np.array(mask_k, dtype=np.float32) / 255.0  # imgsz x imgsz
            # remove padding
            # Note: padding in load_image is evenly distributed due to integer division
            y1, y2 = pad_h, self.imgsz - pad_h
            x1, x2 = pad_w, self.imgsz - pad_w
            canvas = canvas[y1:y2, x1:x2]
            # scale back
            if canvas.size == 0:
                masks_imgsz.append(None)
                continue
            mask_full = Image.fromarray((canvas * 255).astype(np.uint8), mode="L")
            mask_full = mask_full.resize((w0, h0), Image.BILINEAR)
            bin_mask = (np.array(mask_full, dtype=np.float32) / 255.0) >= float(thr)
            masks_imgsz.append(bin_mask.astype(np.uint8))
        return masks_imgsz