Update all files for SegEarth-OV

pipeline.py

@@ -0,0 +1,476 @@
+"""
+Unified SegEarth pipeline: OV, OV-2 (CLIP-based), OV-3 (SAM3-based).
+Training-free open-vocabulary segmentation for remote sensing.
+"""
+import contextlib
+from pathlib import Path
+from typing import List, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from PIL import Image
+from torchvision import transforms
+
+try:
+    from .upsamplers import get_upsampler, FEATUP_CHECKPOINTS
+except ImportError:
+    from upsamplers import get_upsampler, FEATUP_CHECKPOINTS
+
+try:
+    from .prompts.imagenet_template import openai_imagenet_template, sub_imagenet_template
+except ImportError:
+    openai_imagenet_template = [
+        lambda c: f"a photo of a {c}.",
+        lambda c: f"a bad photo of a {c}.",
+        lambda c: f"a photo of many {c}.",
+        lambda c: f"a photo of the large {c}.",
+        lambda c: f"a photo of the small {c}.",
+    ]
+    sub_imagenet_template = openai_imagenet_template[:7]
+
+
+def get_cls_idx(path: Union[str, Path]) -> Tuple[List[str], List[int]]:
+    """Parse class list file (one line per class, comma-separated synonyms)."""
+    path = Path(path)
+    with open(path) as f:
+        lines = f.readlines()
+    class_names, class_indices = [], []
+    for idx, line in enumerate(lines):
+        names_i = [n.strip() for n in line.strip().split(",")]
+        class_names.extend(names_i)
+        class_indices.extend([idx] * len(names_i))
+    return class_names, class_indices
+
+
+class SegEarthPipelineCLIP:
+    """
+    CLIP-based SegEarth pipeline (OV, OV-2).
+    Uses transformers.CLIPModel + SimFeatUp for dense prediction.
+    """
+
+    def __init__(
+        self,
+        model_id: str = "openai/clip-vit-base-patch16",
+        featup_model: str = "jbu_one",
+        featup_weights_path: Optional[Union[str, Path]] = None,
+        class_names_path: Optional[Union[str, Path]] = None,
+        device: str = "cuda",
+        dtype: torch.dtype = torch.float16,
+        cls_token_lambda: float = -0.3,
+        logit_scale: float = 50.0,
+        prob_thd: float = 0.0,
+        bg_idx: int = 0,
+        slide_crop: int = 0,
+        slide_stride: int = 112,
+        template_set: str = "openai",
+    ):
+        from transformers import CLIPModel, CLIPProcessor
+
+        self.device = device
+        self.dtype = dtype
+        self.cls_token_lambda = cls_token_lambda
+        self.logit_scale = logit_scale
+        self.prob_thd = prob_thd
+        self.bg_idx = bg_idx
+        self.slide_crop = slide_crop
+        self.slide_stride = slide_stride
+        self.output_cls_token = cls_token_lambda != 0
+
+        self.templates = sub_imagenet_template if template_set == "sub" else openai_imagenet_template
+
+        self.clip = CLIPModel.from_pretrained(model_id).to(device).to(dtype).eval()
+        try:
+            self.processor = CLIPProcessor.from_pretrained(model_id)
+        except Exception:
+            # Fallback: use tokenizer only (CLIPProcessor can trigger mistral_common compat in some envs)
+            from transformers import CLIPTokenizer
+            self.processor = None
+            self._tokenizer = CLIPTokenizer.from_pretrained(model_id)
+        self.patch_size = 16
+        self.feat_dim = 512
+
+        # Resolve featup path: self-contained repo only (OV/OV-2/weights/featup)
+        ckpt_name = FEATUP_CHECKPOINTS.get(featup_model, "").split("/")[-1]
+        repo_dir = Path(__file__).parent
+        _candidates = [
+            Path(featup_weights_path) if featup_weights_path else None,
+            repo_dir / "OV" / "weights" / "featup" / ckpt_name,
+            repo_dir / "OV-2" / "weights" / "featup" / ckpt_name,
+            repo_dir / "weights" / "featup" / ckpt_name,
+        ]
+        featup_path = next((p for p in _candidates if p and p.exists()), None)
+
+        self.use_featup = featup_path is not None and featup_path.exists()
+        upsampler_name = "bilinear" if not self.use_featup else featup_model.replace("_maskclip", "")
+        self.upsampler = get_upsampler(upsampler_name, self.feat_dim).to(device).to(dtype).eval()
+
+        if self.use_featup:
+            ckpt = torch.load(featup_path, map_location="cpu")
+            sd = ckpt.get("state_dict", ckpt)
+            weights = {k[10:]: v for k, v in sd.items() if k.startswith("upsampler.")}
+            self.upsampler.load_state_dict(weights, strict=True)
+
+        repo_dir = Path(__file__).parent
+        cls_path = class_names_path or (repo_dir / "configs" / "cls_openearthmap_sar.txt")
+        cls_path = Path(cls_path)
+        if cls_path.exists():
+            self.class_names, self.class_indices = get_cls_idx(cls_path)
+        else:
+            self.class_names = ["building", "road", "water", "vegetation", "bare soil"]
+            self.class_indices = list(range(len(self.class_names)))
+
+        self.num_classes = max(self.class_indices) + 1
+        self.num_queries = len(self.class_indices)
+        self.query_idx = torch.tensor(self.class_indices, dtype=torch.int64, device=device)
+        self._build_query_features()
+
+    def _build_query_features(self):
+        query_features = []
+        with torch.no_grad():
+            tokenizer = getattr(self, "_tokenizer", None) or (self.processor.tokenizer if self.processor else None)
+            for name in self.class_names:
+                texts = [t(name) for t in self.templates]
+                inputs = tokenizer(text=texts, return_tensors="pt", padding=True, truncation=True)
+                inputs = {k: v.to(self.device) for k, v in inputs.items()}
+                out = self.clip.get_text_features(**inputs)
+                if hasattr(out, "shape"):
+                    feat_t = out
+                elif hasattr(out, "pooler_output") and out.pooler_output is not None:
+                    feat_t = out.pooler_output
+                else:
+                    feat_t = out.last_hidden_state.mean(1)
+                feat = feat_t.mean(0) / feat_t.mean(0).norm()
+                query_features.append(feat.unsqueeze(0))
+        self.query_features = torch.cat(query_features, dim=0).to(self.dtype)
+
+    def _encode_image_patches(self, pixel_values: torch.Tensor) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+        out = self.clip.vision_model(pixel_values)
+        hidden = out.last_hidden_state
+        proj = self.clip.visual_projection.weight
+        patch_tokens = hidden[:, 1:, :]
+        patch_feats = patch_tokens @ proj.T
+        cls_token = None
+        if self.output_cls_token:
+            cls_tok = hidden[:, 0:1, :]
+            cls_token = (cls_tok @ proj.T).squeeze(1)
+            cls_token = F.normalize(cls_token, dim=-1)
+        return patch_feats, cls_token
+
+    def _preprocess_image(self, image: Image.Image, size: Optional[int] = 224, keep_size: bool = False) -> torch.Tensor:
+        t = transforms.Compose([
+            transforms.ToTensor(),
+            transforms.Normalize(
+                [0.48145466, 0.4578275, 0.40821073],
+                [0.26862954, 0.26130258, 0.27577711],
+            ),
+        ])
+        x = t(image.convert("RGB"))
+        if not keep_size and size:
+            x = transforms.functional.resize(x, (size, size))
+        return x.unsqueeze(0).to(self.device).to(self.dtype)
+
+    def _compute_padsize(self, H: int, W: int) -> Tuple[int, int, int, int]:
+        l, r, t, b = 0, 0, 0, 0
+        if W % self.patch_size:
+            lr = self.patch_size - (W % self.patch_size)
+            l = lr // 2
+            r = lr - l
+        if H % self.patch_size:
+            tb = self.patch_size - (H % self.patch_size)
+            t = tb // 2
+            b = tb - t
+        return l, r, t, b
+
+    def _forward_single_crop(self, img_tensor: torch.Tensor) -> torch.Tensor:
+        B, C, H, W = img_tensor.shape
+        patch_h, patch_w = H // self.patch_size, W // self.patch_size
+        patch_feats, cls_token = self._encode_image_patches(img_tensor)
+        patch_feats = patch_feats.permute(0, 2, 1).view(B, self.feat_dim, patch_h, patch_w)
+        patch_feats = patch_feats.to(self.dtype)
+        img_tensor = img_tensor.to(self.dtype)
+        patch_feats = self.upsampler(patch_feats, img_tensor)
+        out_h, out_w = H, W
+        patch_feats = patch_feats.view(B, self.feat_dim, -1).permute(0, 2, 1)
+        patch_feats = F.normalize(patch_feats, dim=-1)
+        logits = patch_feats @ self.query_features.T
+        if self.output_cls_token and cls_token is not None:
+            cls_logits = cls_token @ self.query_features.T
+            logits = logits + cls_logits.unsqueeze(1) * self.cls_token_lambda
+        logits = logits.permute(0, 2, 1).view(B, self.num_queries, out_h, out_w)
+        return logits[0]
+
+    def _forward_slide(self, img_tensor: torch.Tensor, ori_shape: Tuple[int, int]) -> torch.Tensor:
+        B, _, h_img, w_img = img_tensor.shape
+        stride = (self.slide_stride, self.slide_stride)
+        crop = (self.slide_crop, self.slide_crop)
+        h_stride, w_stride = stride
+        h_crop, w_crop = crop
+        h_grids = max(h_img - h_crop + h_stride - 1, 0) // h_stride + 1
+        w_grids = max(w_img - w_crop + w_stride - 1, 0) // w_stride + 1
+        preds = img_tensor.new_zeros((B, self.num_queries, h_img, w_img))
+        count_mat = img_tensor.new_zeros((B, 1, h_img, w_img))
+        for h_idx in range(h_grids):
+            for w_idx in range(w_grids):
+                y1 = h_idx * h_stride
+                x1 = w_idx * w_stride
+                y2 = min(y1 + h_crop, h_img)
+                x2 = min(x1 + w_crop, w_img)
+                y1 = max(y2 - h_crop, 0)
+                x1 = max(x2 - w_crop, 0)
+                crop_img = img_tensor[:, :, y1:y2, x1:x2]
+                H, W = crop_img.shape[2:]
+                l, r, t, b = self._compute_padsize(H, W)
+                if any([l, r, t, b]):
+                    crop_img = F.pad(crop_img, (l, r, t, b))
+                crop_logits = self._forward_single_crop(crop_img)
+                if any([l, r, t, b]):
+                    crop_logits = crop_logits[:, t : t + H, l : l + W]
+                pad_crop = F.pad(
+                    crop_logits.unsqueeze(0),
+                    (int(x1), int(preds.shape[3] - x2), int(y1), int(preds.shape[2] - y2)),
+                )
+                preds += pad_crop
+                count_mat[:, :, y1:y2, x1:x2] += 1
+        preds = preds / count_mat.clamp(min=1)
+        logits = F.interpolate(preds, size=ori_shape, mode="bilinear")
+        return logits[0]
+
+    def _postprocess(self, logits: torch.Tensor) -> torch.Tensor:
+        logits = logits * self.logit_scale
+        probs = logits.softmax(0)
+        if self.num_classes != self.num_queries:
+            cls_idx = F.one_hot(self.query_idx, self.num_classes)
+            cls_idx = cls_idx.T.view(self.num_classes, self.num_queries, 1, 1)
+            probs = (probs.unsqueeze(0) * cls_idx).max(1)[0]
+        seg_pred = probs.argmax(0, keepdim=True)
+        if self.prob_thd > 0:
+            max_prob = probs.max(0, keepdim=True)[0]
+            seg_pred[max_prob < self.prob_thd] = self.bg_idx
+        return seg_pred.squeeze(0)
+
+    @torch.no_grad()
+    def __call__(self, image: Union[Image.Image, torch.Tensor], return_logits: bool = False) -> torch.Tensor:
+        if isinstance(image, Image.Image):
+            use_slide = self.slide_crop > 0
+            keep_size = use_slide
+            img_tensor = self._preprocess_image(image, size=224, keep_size=keep_size)
+        else:
+            img_tensor = image.to(self.device).to(self.dtype)
+            if img_tensor.dim() == 3:
+                img_tensor = img_tensor.unsqueeze(0)
+        B, C, H, W = img_tensor.shape
+        ori_shape = (H, W)
+        use_slide = self.slide_crop > 0 and (H > self.slide_crop or W > self.slide_crop)
+        if use_slide:
+            logits = self._forward_slide(img_tensor, ori_shape)
+        else:
+            l, r, t, b = self._compute_padsize(H, W)
+            if any([l, r, t, b]):
+                img_tensor = F.pad(img_tensor, (l, r, t, b))
+                out_h, out_w = img_tensor.shape[2], img_tensor.shape[3]
+            else:
+                out_h, out_w = H, W
+            logits = self._forward_single_crop(img_tensor)
+            if any([l, r, t, b]):
+                logits = logits[:, t : t + H, l : l + W]
+            if (out_h, out_w) != ori_shape:
+                logits = F.interpolate(logits.unsqueeze(0), size=ori_shape, mode="bilinear").squeeze(0)
+        if return_logits:
+            if self.num_classes != self.num_queries:
+                cls_idx = F.one_hot(self.query_idx, self.num_classes)
+                cls_idx = cls_idx.T.view(self.num_classes, self.num_queries, 1, 1)
+                logits = (logits.unsqueeze(0) * cls_idx).max(1)[0]
+            return logits
+        return self._postprocess(logits)
+
+
+class SegEarthPipelineSAM3:
+    """
+    SAM3-based SegEarth pipeline (OV-3).
+    Uses sam3 package for open-vocabulary segmentation.
+    Requires: pip install sam3 (or transformers>=4.45 for Sam3Model)
+    """
+
+    def __init__(
+        self,
+        model_id: str = "facebook/sam3",
+        local_checkpoint: Optional[Union[str, Path]] = None,
+        class_names_path: Optional[Union[str, Path]] = None,
+        device: str = "cuda",
+        prob_thd: float = 0.0,
+        bg_idx: int = 0,
+        slide_crop: int = 0,
+        slide_stride: int = 112,
+        confidence_threshold: float = 0.5,
+        use_sem_seg: bool = True,
+        use_presence_score: bool = True,
+        use_transformer_decoder: bool = True,
+    ):
+        self.device = device
+        self.prob_thd = prob_thd
+        self.bg_idx = bg_idx
+        self.slide_crop = slide_crop
+        self.slide_stride = slide_stride
+        self.confidence_threshold = confidence_threshold
+        self.use_sem_seg = use_sem_seg
+        self.use_presence_score = use_presence_score
+        self.use_transformer_decoder = use_transformer_decoder
+
+        # Workaround for cuDNN "No execution plans support the graph" with SDPA
+        if device == "cuda":
+            if hasattr(torch.backends.cuda, "enable_flash_sdp"):
+                torch.backends.cuda.enable_flash_sdp(False)
+                torch.backends.cuda.enable_mem_efficient_sdp(False)
+            if hasattr(torch.backends.cuda, "enable_math_sdp"):
+                torch.backends.cuda.enable_math_sdp(True)
+
+        try:
+            from sam3 import build_sam3_image_model
+            from sam3.model.sam3_image_processor import Sam3Processor
+        except ImportError:
+            raise ImportError(
+                "SegEarth OV-3 requires the sam3 package. Install from: "
+                "https://github.com/facebookresearch/sam3 or use transformers.Sam3Model.from_pretrained('facebook/sam3')"
+            )
+
+        ckpt_path = Path(local_checkpoint) if local_checkpoint else None
+        if ckpt_path and not ckpt_path.is_absolute():
+            ckpt_path = Path(__file__).parent / "OV-3" / ckpt_path
+        use_safetensors = ckpt_path and str(ckpt_path).endswith(".safetensors") and ckpt_path.exists()
+        use_pt = ckpt_path and (str(ckpt_path).endswith(".pt") or str(ckpt_path).endswith(".bin")) and ckpt_path.exists()
+
+        if use_safetensors:
+            self.model = build_sam3_image_model(checkpoint_path=None, load_from_HF=False, device=device)
+            from safetensors.torch import load_file
+            state_dict = load_file(str(ckpt_path))
+            # HF model.safetensors uses "detector_model." prefix; sam3 expects "detector." -> stripped
+            state_dict = {k.replace("detector_model.", ""): v for k, v in state_dict.items()}
+            self.model.load_state_dict(state_dict, strict=False)
+        elif use_pt:
+            self.model = build_sam3_image_model(checkpoint_path=str(ckpt_path), load_from_HF=False, device=device)
+        else:
+            self.model = build_sam3_image_model(checkpoint_path=None, load_from_HF=True, device=device)
+        self.processor = Sam3Processor(self.model, confidence_threshold=confidence_threshold, device=device)
+
+        repo_dir = Path(__file__).parent
+        cls_path = class_names_path or (repo_dir / "configs" / "cls_openearthmap_sar.txt")
+        cls_path = Path(cls_path)
+        if cls_path.exists():
+            self.class_names, self.class_indices = get_cls_idx(cls_path)
+        else:
+            self.class_names = ["building", "road", "water", "vegetation", "bare soil"]
+            self.class_indices = list(range(len(self.class_names)))
+        self.num_classes = max(self.class_indices) + 1
+        self.num_queries = len(self.class_indices)
+        self.query_idx = torch.tensor(self.class_indices, dtype=torch.int64, device=device)
+
+    def _inference_single_view(self, image: Image.Image) -> torch.Tensor:
+        w, h = image.size
+        seg_logits = torch.zeros((self.num_queries, h, w), device=self.device)
+        sdp_ctx = (
+            torch.backends.cuda.sdp_kernel(enable_flash=False, enable_math=True, enable_mem_efficient=False, enable_cudnn=False)
+            if self.device == "cuda" and hasattr(torch.backends.cuda, "sdp_kernel")
+            else contextlib.nullcontext()
+        )
+        with torch.no_grad(), torch.autocast(device_type="cuda", dtype=torch.bfloat16), sdp_ctx:
+            inference_state = self.processor.set_image(image)
+            for query_idx, query_word in enumerate(self.class_names):
+                self.processor.reset_all_prompts(inference_state)
+                inference_state = self.processor.set_text_prompt(state=inference_state, prompt=query_word)
+                if self.use_transformer_decoder and inference_state.get("masks_logits") is not None:
+                    inst_len = inference_state["masks_logits"].shape[0]
+                    for inst_id in range(inst_len):
+                        instance_logits = inference_state["masks_logits"][inst_id].squeeze()
+                        instance_score = inference_state["object_score"][inst_id]
+                        if instance_logits.shape != (h, w):
+                            instance_logits = F.interpolate(
+                                instance_logits.view(1, 1, *instance_logits.shape),
+                                size=(h, w), mode="bilinear", align_corners=False
+                            ).squeeze()
+                        seg_logits[query_idx] = torch.max(seg_logits[query_idx], instance_logits * instance_score)
+                if self.use_sem_seg and inference_state.get("semantic_mask_logits") is not None:
+                    semantic_logits = inference_state["semantic_mask_logits"]
+                    if semantic_logits.shape != (h, w):
+                        semantic_logits = F.interpolate(
+                            semantic_logits.view(1, 1, *semantic_logits.shape) if semantic_logits.dim() == 2 else semantic_logits.unsqueeze(0),
+                            size=(h, w), mode="bilinear", align_corners=False
+                        ).squeeze()
+                    seg_logits[query_idx] = torch.max(seg_logits[query_idx], semantic_logits)
+                if self.use_presence_score and inference_state.get("presence_score") is not None:
+                    seg_logits[query_idx] = seg_logits[query_idx] * inference_state["presence_score"]
+        return seg_logits
+
+    def slide_inference(self, image: Image.Image) -> torch.Tensor:
+        w_img, h_img = image.size
+        stride = (self.slide_stride, self.slide_stride)
+        crop = (self.slide_crop, self.slide_crop)
+        h_stride, w_stride = stride
+        h_crop, w_crop = crop
+        h_grids = max(h_img - h_crop + h_stride - 1, 0) // h_stride + 1
+        w_grids = max(w_img - w_crop + w_stride - 1, 0) // w_stride + 1
+        preds = torch.zeros((self.num_queries, h_img, w_img), device=self.device)
+        count_mat = torch.zeros((1, h_img, w_img), device=self.device)
+        for h_idx in range(h_grids):
+            for w_idx in range(w_grids):
+                y1 = h_idx * h_stride
+                x1 = w_idx * w_stride
+                y2 = min(y1 + h_crop, h_img)
+                x2 = min(x1 + w_crop, w_img)
+                y1 = max(y2 - h_crop, 0)
+                x1 = max(x2 - w_crop, 0)
+                crop_img = image.crop((x1, y1, x2, y2))
+                crop_seg = self._inference_single_view(crop_img)
+                preds[:, y1:y2, x1:x2] += crop_seg
+                count_mat[:, y1:y2, x1:x2] += 1
+        return preds / count_mat.clamp(min=1)
+
+    @torch.no_grad()
+    def __call__(self, image: Union[Image.Image, torch.Tensor]) -> torch.Tensor:
+        if isinstance(image, torch.Tensor):
+            image = transforms.functional.to_pil_image(image)
+        image = image.convert("RGB")
+        if self.slide_crop > 0 and (image.size[0] > self.slide_crop or image.size[1] > self.slide_crop):
+            seg_logits = self.slide_inference(image)
+        else:
+            seg_logits = self._inference_single_view(image)
+        if self.num_classes != self.num_queries:
+            cls_idx = F.one_hot(self.query_idx, self.num_classes)
+            cls_idx = cls_idx.T.view(self.num_classes, self.num_queries, 1, 1)
+            seg_logits = (seg_logits.unsqueeze(0) * cls_idx).max(1)[0]
+        seg_pred = seg_logits.argmax(0, keepdim=True)
+        if self.prob_thd > 0:
+            max_prob = seg_logits.max(0, keepdim=True)[0]
+            seg_pred[max_prob < self.prob_thd] = self.bg_idx
+        return seg_pred.squeeze(0)
+
+
+def SegEarthPipeline(
+    variant: str = "OV-2",
+    model_id: Optional[str] = None,
+    **kwargs,
+):
+    """
+    Factory for SegEarth pipelines. Load from self-contained subfolders OV/, OV-2/, OV-3/.
+    Args:
+        variant: One of OV, OV-2, OV-3 (or legacy: ov_clip_openai_vitb16, ov2_alignearth_sar, ov3_sam3)
+        model_id: Override HF model ID
+        **kwargs: Passed to pipeline constructor
+    """
+    import json
+    repo_dir = Path(__file__).parent
+    variant_map = {"ov_clip_openai_vitb16": "OV", "ov2_alignearth_sar": "OV-2", "ov3_sam3": "OV-3"}
+    subfolder = variant_map.get(variant, variant)
+    sub_path = repo_dir / subfolder / "pipeline.py"
+    if sub_path.exists():
+        import importlib.util
+        spec = importlib.util.spec_from_file_location(f"segearth_{subfolder}", sub_path)
+        mod = importlib.util.module_from_spec(spec)
+        spec.loader.exec_module(mod)
+        return mod.load(**kwargs) if model_id is None else mod.load(model_id=model_id, **kwargs)
+    # Fallback: legacy flat config
+    if model_id is None:
+        model_id = "BiliSakura/AlignEarth-SAR-ViT-B-16"
+    if variant in ("ov3_sam3", "OV-3"):
+        return SegEarthPipelineSAM3(model_id=model_id or "facebook/sam3", **kwargs)
+    return SegEarthPipelineCLIP(model_id=model_id, **kwargs)