first version

sam2/sam_image_mask_extract.py

@@ -0,0 +1,152 @@
+import os
+import cv2
+import json
+import torch
+import numpy as np
+
+from segment_anything import sam_model_registry, SamPredictor
+
+# ---------------- CONFIG ----------------
+SAM_CHECKPOINT = "/media/bs5/New Volume/Brinda/Codes/Florida/SAM/sam_vit_h_4b8939.pth"
+MODEL_TYPE = "vit_h"
+
+IMAGE_DIR = "/media/bs5/New Volume/Brinda/Dataset/Florida/New_Dataset/set2/images_seg"
+LABEL_DIR = "/media/bs5/New Volume/Brinda/Dataset/Florida/New_Dataset/set2/labels_seg/"
+OUTPUT_DIR = "/media/bs5/New Volume/Brinda/Dataset/Florida/classification_data/"
+
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+
+os.makedirs(OUTPUT_DIR, exist_ok=True)
+os.makedirs(os.path.join(OUTPUT_DIR, "masks"), exist_ok=True)
+os.makedirs(os.path.join(OUTPUT_DIR, "overlay"), exist_ok=True)
+os.makedirs(os.path.join(OUTPUT_DIR, "polygons"), exist_ok=True)
+
+# ---------------- LOAD SAM ----------------
+sam = sam_model_registry[MODEL_TYPE](checkpoint=SAM_CHECKPOINT)
+sam.to(DEVICE)
+predictor = SamPredictor(sam)
+
+# ---------------- HELPERS ----------------
+def yolo_to_xyxy(label, img_w, img_h):
+    """
+    YOLO format: class cx cy w h (normalized)
+    Convert → pixel [x1, y1, x2, y2]
+    """
+    cls, cx, cy, w, h = label
+
+    x1 = int((cx - w / 2) * img_w)
+    y1 = int((cy - h / 2) * img_h)
+    x2 = int((cx + w / 2) * img_w)
+    y2 = int((cy + h / 2) * img_h)
+
+    return [x1, y1, x2, y2]
+
+
+def expand_box(x1, y1, x2, y2, pad, w, h):
+    return [
+        max(0, x1 - pad),
+        max(0, y1 - pad),
+        min(w, x2 + pad),
+        min(h, y2 + pad),
+    ]
+
+
+def mask_to_polygon(mask):
+    contours, _ = cv2.findContours(
+        mask.astype(np.uint8),
+        cv2.RETR_EXTERNAL,
+        cv2.CHAIN_APPROX_SIMPLE
+    )
+
+    polygons = []
+    for cnt in contours:
+        cnt = cnt.squeeze()
+        if len(cnt.shape) == 2 and len(cnt) > 2:
+            polygons.append(cnt.tolist())
+
+    return polygons
+
+
+# ---------------- PROCESS ----------------
+for img_file in os.listdir(IMAGE_DIR):
+
+    if not img_file.lower().endswith((".jpg", ".png", ".jpeg")):
+        continue
+
+    image_path = os.path.join(IMAGE_DIR, img_file)
+    label_path = os.path.join(LABEL_DIR, img_file.replace(".jpg", ".txt").replace(".png", ".txt"))
+
+    if not os.path.exists(label_path):
+        continue
+
+    image = cv2.imread(image_path)
+    if image is None:
+        continue
+
+    h, w = image.shape[:2]
+    image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+
+    predictor.set_image(image_rgb)
+
+    overlay = image.copy()
+    output_data = []
+
+    # -------- Read YOLO labels --------
+    with open(label_path, "r") as f:
+        lines = f.readlines()
+
+    for idx, line in enumerate(lines):
+        parts = list(map(float, line.strip().split()))
+        cls_id = int(parts[0])
+        cx, cy, bw, bh = parts[1:]
+
+        # YOLO → XYXY
+        x1, y1, x2, y2 = yolo_to_xyxy([cls_id, cx, cy, bw, bh], w, h)
+
+        # Expand bbox (important for SAM)
+        x1, y1, x2, y2 = expand_box(x1, y1, x2, y2, pad=10, w=w, h=h)
+
+        box = np.array([x1, y1, x2, y2])
+
+        # -------- SAM Prediction --------
+        masks, scores, _ = predictor.predict(
+            box=box,
+            multimask_output=False
+        )
+
+        mask = masks[0]
+
+        # -------- Save mask --------
+        mask_name = f"{img_file.split('.')[0]}_{idx}.png"
+        mask_path = os.path.join(OUTPUT_DIR, "masks", mask_name)
+        cv2.imwrite(mask_path, mask.astype(np.uint8) * 255)
+
+        # -------- Overlay --------
+        color = np.random.randint(0, 255, (3,))
+        overlay[mask] = (overlay[mask] * 0.5 + color * 0.5).astype(np.uint8)
+
+        # -------- Polygon --------
+        polygons = mask_to_polygon(mask)
+
+        output_data.append({
+            "id": idx,
+            "class": cls_id,
+            "bbox": [x1, y1, x2, y2],
+            "polygon": polygons
+        })
+
+    # -------- Save overlay --------
+    cv2.imwrite(
+        os.path.join(OUTPUT_DIR, "overlay", img_file),
+        overlay
+    )
+
+    # -------- Save JSON --------
+    json_path = os.path.join(
+        OUTPUT_DIR, "polygons", img_file.replace(".jpg", ".json").replace(".png", ".json")
+    )
+
+    with open(json_path, "w") as f:
+        json.dump(output_data, f, indent=4)
+
+print(" Processing Complete")