process.py

import os
import argparse
import numpy as np
import matplotlib.pyplot as plt
import torch
import cv2

# PyTorch bellek fragmantasyonunu önler
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"

from sam2.build_sam import build_sam2
from sam2.automatic_mask_generator import SAM2AutomaticMaskGenerator


def load_sample_from_npy(npy_path: str) -> dict:
    """Veriyi yükler (.item() kontrolü ile güvenli yükleme)"""
    data_npy = np.load(npy_path, allow_pickle=True, encoding="bytes")
    if isinstance(data_npy, np.ndarray):
        sample = data_npy.item()
    else:
        sample = data_npy

    if "rgb" not in sample:
        raise KeyError(f"'rgb' key not found in {npy_path}")
    return sample


def compute_exclusion_from_background(sample, bg_sample, user_threshold=10.0):
    """
    Background Subtraction. Metre/Pixel ayrımını otomatik yapar.
    """
    depth_key = "depth_raw" if "depth_raw" in sample else "depth"
    has_depth = depth_key in sample and depth_key in bg_sample

    if has_depth:
        print(f"Using Depth ({depth_key}) for Background Subtraction...")
        curr = sample[depth_key]
        bg = bg_sample[depth_key]

        curr = np.nan_to_num(curr, nan=0.0, posinf=0.0, neginf=0.0)
        bg = np.nan_to_num(bg, nan=0.0, posinf=0.0, neginf=0.0)

        diff = cv2.absdiff(curr, bg)

        if diff.dtype == np.float32 or diff.dtype == np.float64:
            if user_threshold > 1.0:
                actual_threshold = 0.02  # 2 cm hassasiyet
                print(f"   [INFO] Veri float (Metre). Threshold: {actual_threshold}m")
            else:
                actual_threshold = user_threshold
            _, fg_mask = cv2.threshold(diff, actual_threshold, 255, cv2.THRESH_BINARY)
        else:
            actual_threshold = int(user_threshold)
            print(f"   [INFO] Veri int (Pixel). Threshold: {actual_threshold}")
            _, fg_mask = cv2.threshold(diff, actual_threshold, 255, cv2.THRESH_BINARY)

        fg_mask = fg_mask.astype(np.uint8)

    else:
        print("Depth not found, falling back to RGB Background Subtraction...")
        curr = sample["rgb"]
        bg = bg_sample["rgb"]
        diff = cv2.absdiff(curr, bg)
        diff_gray = cv2.cvtColor(diff, cv2.COLOR_RGB2GRAY)
        _, fg_mask = cv2.threshold(diff_gray, 30, 255, cv2.THRESH_BINARY)

    kernel = np.ones((3, 3), np.uint8)
    fg_mask = cv2.morphologyEx(fg_mask, cv2.MORPH_OPEN, kernel)
    fg_mask = cv2.morphologyEx(fg_mask, cv2.MORPH_CLOSE, kernel, iterations=2)

    exclusion_mask = fg_mask == 0
    return exclusion_mask, fg_mask


def exclusion_overlap_ratio(seg, excl):
    area = int(seg.sum())
    if area == 0:
        return 0.0
    overlap = int(np.logical_and(seg, excl).sum())
    return overlap / area


def filter_by_exclusion(masks, exclusion_mask, overlap_thresh=0.20):
    kept = []
    for m in masks:
        seg = m["segmentation"].astype(bool)
        r = exclusion_overlap_ratio(seg, exclusion_mask)
        if r < overlap_thresh:
            kept.append(m)
    return kept


def smart_merge_masks(masks, iou_thresh=0.1, ioa_thresh=0.6):
    """Parçalanmış maskeleri birleştirir (Merge)."""
    if not masks:
        return []

    sorted_masks = sorted(masks, key=lambda x: x["area"], reverse=True)
    merged_indices = set()
    final_masks = []

    for i in range(len(sorted_masks)):
        if i in merged_indices:
            continue

        base_mask = sorted_masks[i]["segmentation"].copy()
        base_area = sorted_masks[i]["area"]

        for j in range(i + 1, len(sorted_masks)):
            if j in merged_indices:
                continue

            cand_mask = sorted_masks[j]["segmentation"]
            cand_area = sorted_masks[j]["area"]

            intersection = np.logical_and(base_mask, cand_mask).sum()
            if intersection == 0:
                continue

            union = np.logical_or(base_mask, cand_mask).sum()
            iou = intersection / union
            ioa = intersection / min(base_area, cand_area)

            if iou > iou_thresh or ioa > ioa_thresh:
                base_mask = np.logical_or(base_mask, cand_mask)
                base_area = base_mask.sum()
                merged_indices.add(j)

        new_entry = sorted_masks[i].copy()
        new_entry["segmentation"] = base_mask
        new_entry["area"] = base_area
        final_masks.append(new_entry)

    return final_masks


def masks_to_label_map(masks, shape):
    H, W = shape[:2]
    label_map = np.zeros((H, W), dtype=np.uint8)
    for i, mask_data in enumerate(masks):
        obj_id = i + 1
        mask = mask_data["segmentation"].astype(bool)
        label_map[mask] = obj_id
    return label_map


def main():
    parser = argparse.ArgumentParser()

    # Bash script uyumluluğu için -e eklendi (kullanmasak bile hata vermesin)
    parser.add_argument(
        "-e",
        "--environment",
        type=str,
        required=False,
        help="Environment name (Passed by bash script)",
    )

    # Input path
    parser.add_argument(
        "-i",
        "--input_path",
        type=str,
        required=True,
        help="Input directory or raw_capture.npy path",
    )

    # Opsiyonel: Eğer verilmezse input klasöründe aranır
    parser.add_argument(
        "-bg",
        "--bg_path",
        type=str,
        default="/home/furkand/dev/DiffusionProject/background.npy",
        help="Background npy path (Optional)",
    )
    parser.add_argument(
        "-o", "--output_path", type=str, required=False, help="Output directory"
    )

    # SAM2 Configs
    parser.add_argument(
        "--checkpoint",
        default="/home/furkand/dev/sam2/checkpoints/sam2.1_hiera_base_plus.pt",
    )
    parser.add_argument("--cfg", default="configs/sam2.1/sam2.1_hiera_b+.yaml")

    # Algoritma Ayarları
    parser.add_argument("--merge-iou", type=float, default=0.1)
    parser.add_argument("--merge-ioa", type=float, default=0.5)
    parser.add_argument("--excl-overlap", type=float, default=0.50)
    parser.add_argument("--depth-thresh", type=float, default=10.0)
    parser.add_argument("--min-area-frac", type=float, default=0.001)

    args = parser.parse_args()

    # --- 1. Path Ayarlama ---
    print(f"Environment: {args.environment}")

    # Eğer input bir klasör ise (bash script klasör yolluyor: $target_dir/temp)
    if os.path.isdir(args.input_path):
        target_file = os.path.join(args.input_path, "raw_capture.npy")
        input_dir = args.input_path  # Background'u burada arayacağız
    else:
        target_file = args.input_path
        input_dir = os.path.dirname(args.input_path)

    if not os.path.exists(target_file):
        raise FileNotFoundError(f"Input file not found: {target_file}")

    # Background Dosyasını Bulma Mantığı
    # Bash script -bg yollamıyor, bu yüzden input klasöründe 'background.npy' arıyoruz.
    if args.bg_path is None:
        possible_bg = os.path.join(input_dir, "background.npy")
        if os.path.exists(possible_bg):
            args.bg_path = possible_bg
            print(f"Auto-detected background: {args.bg_path}")
        else:
            # Belki bir üst klasördedir? (storage/test_sam2/background.npy gibi)
            parent_bg = os.path.join(os.path.dirname(input_dir), "background.npy")
            if os.path.exists(parent_bg):
                args.bg_path = parent_bg
                print(f"Auto-detected background (parent dir): {args.bg_path}")
            else:
                # Background yoksa hata vermek yerine boş bir background yaratabiliriz ama şimdilik hata verelim.
                # Ya da capture.py'nin background'u nereye kaydettiğine emin olmamız lazım.
                # Şimdilik aynı klasörde olduğunu varsayıyoruz.
                print(
                    f"[WARNING] background.npy not found in {input_dir}. Assuming static scene logic might fail."
                )
                # raise FileNotFoundError(f"Background file not found in {input_dir}")
                # Hata vermeyip devam edelim mi? Hayır, compute_exclusion patlar.
                # Geçici çözüm: Eğer background yoksa kodun patlamaması için depth farkını iptal edebiliriz ama
                # senin senaryonda background.npy kesin vardır diye umuyorum.
                pass

    # Output path
    if args.output_path is None:
        output_dir = os.path.dirname(target_file)
    else:
        output_dir = args.output_path
        os.makedirs(output_dir, exist_ok=True)

    # --- 2. Verileri Yükle ---
    sample = load_sample_from_npy(target_file)

    # Background yükleme (Eğer bulunduysa)
    if args.bg_path and os.path.exists(args.bg_path):
        bg_sample = load_sample_from_npy(args.bg_path)
        has_bg = True
    else:
        print(
            "!!! Background data missing. Running SAM2 without background subtraction !!!"
        )
        has_bg = False
        bg_sample = None

    img = sample["rgb"]
    if img.dtype != np.uint8:
        img = np.clip(img, 0, 255).astype(np.uint8)

    # --- 3. Exclusion Mask (Background Varsa) ---
    if has_bg:
        print("Computing exclusion mask...")
        exclusion, fg_vis = compute_exclusion_from_background(
            sample, bg_sample, user_threshold=args.depth_thresh
        )
        if np.count_nonzero(fg_vis) == 0:
            print("[UYARI] Foreground simsiyah! Threshold ayarını kontrol et.")
    else:
        # Background yoksa exclusion maskesi boş (False) olsun
        exclusion = np.zeros(img.shape[:2], dtype=bool)

    # --- 4. SAM2 Modeli ---
    device = "cuda" if torch.cuda.is_available() else "cpu"
    torch.cuda.empty_cache()
    print(f"Building SAM2 on {device}...")

    model = build_sam2(args.cfg, args.checkpoint, device=device)

    mask_generator = SAM2AutomaticMaskGenerator(
        model,
        points_per_side=32,
        points_per_batch=32,
        pred_iou_thresh=0.86,
        stability_score_thresh=0.92,
    )

    # --- 5. Generate ---
    print("Running SAM2 Generator...")
    with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16):
        masks = mask_generator.generate(img)
    print(f"SAM2 generated {len(masks)} raw masks.")

    # --- 6. Filtrele ---
    if has_bg:
        masks2 = filter_by_exclusion(masks, exclusion, overlap_thresh=args.excl_overlap)
        print(f"After background filter: {len(masks2)} masks.")
    else:
        masks2 = masks

    # --- 7. Merge ---
    H, W = img.shape[:2]
    min_area = int(H * W * args.min_area_frac)
    masks_clean = [m for m in masks2 if m["area"] > min_area]

    print("Merging fragmented parts...")
    final_masks = smart_merge_masks(
        masks_clean, iou_thresh=args.merge_iou, ioa_thresh=args.merge_ioa
    )
    print(f"Final Merged Objects: {len(final_masks)}")

    # --- 8. Kaydetme (seg_data.npy) ---
    seg_map = masks_to_label_map(final_masks, img.shape)

    # Orijinal verileri al
    depth_key = "depth_raw" if "depth_raw" in sample else "depth"

    save_data = {
        "rgb": sample["rgb"],
        "depth": sample[depth_key],
        "seg": seg_map,
        "K": sample.get("K", np.eye(3)),
    }

    save_filename = "seg_data.npy"
    save_full_path = os.path.join(output_dir, save_filename)

    np.save(save_full_path, save_data)
    print(f"\n>>> Results saved to: {save_full_path}")

    # Görsel Kaydet (Debugging için, pipeline içinde engel olmasın diye show() kapalı)
    vis_path = os.path.join(output_dir, "segmentation_vis.png")
    plt.figure(figsize=(10, 5))
    plt.subplot(1, 2, 1)
    plt.title("RGB")
    plt.imshow(img)
    plt.axis("off")
    plt.subplot(1, 2, 2)
    plt.title(f"Segmentation (N={len(final_masks)})")
    plt.imshow(seg_map, cmap="jet", interpolation="nearest")
    plt.axis("off")
    plt.savefig(vis_path)
    print(f">>> Visual saved to: {vis_path}")
    plt.close()


if __name__ == "__main__":
    main()