Upload finetune-sam2.py

finetune-sam2.py

@@ -0,0 +1,426 @@
+import os
+import random
+import pandas as pd
+import cv2
+import torch
+import torch.nn.utils
+import torch.nn.functional as F
+import numpy as np
+import matplotlib.pyplot as plt
+import matplotlib.colors as mcolors
+from sklearn.model_selection import train_test_split
+ 
+from sam2.build_sam import build_sam2
+from sam2.sam2_image_predictor import SAM2ImagePredictor
+
+def set_seeds():
+    SEED_VALUE = 42
+    random.seed(SEED_VALUE)
+    np.random.seed(SEED_VALUE)
+    torch.manual_seed(SEED_VALUE)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed(SEED_VALUE)
+        torch.cuda.manual_seed_all(SEED_VALUE)
+        torch.backends.cudnn.deterministic = True
+        torch.backends.cudnn.benchmark = True
+ 
+set_seeds()
+
+data_dir = "./sam2-data"
+images_dir = os.path.join(data_dir, "images")
+masks_dir = os.path.join(data_dir, "masks")
+ 
+train_df = pd.read_csv(os.path.join(data_dir, "train.csv"))
+ 
+train_df, test_df = train_test_split(train_df, test_size=0.1, random_state=42)
+ 
+train_data = []
+for index, row in train_df.iterrows():
+   image_name = row['imageid']
+   mask_name = row['maskid']
+   train_data.append({
+       "image": os.path.join(images_dir, image_name),
+       "annotation": os.path.join(masks_dir, mask_name)
+   })
+ 
+test_data = []
+
+for index, row in test_df.iterrows():
+   image_name = row['imageid']
+   mask_name = row['maskid']
+   test_data.append({
+       "image": os.path.join(images_dir, image_name),
+       "annotation": os.path.join(masks_dir, mask_name)
+   })
+
+def read_batch(data, visualize_data=True):
+   ent = data[np.random.randint(len(data))]
+   Img = cv2.imread(ent["image"])[..., ::-1]  
+   ann_map = cv2.imread(ent["annotation"], cv2.IMREAD_GRAYSCALE)
+ 
+   if Img is None or ann_map is None:
+       print(f"Error: Could not read image or mask from path {ent['image']} or {ent['annotation']}")
+       return None, None, None, 0
+ 
+   r = np.min([1024 / Img.shape[1], 1024 / Img.shape[0]])
+   Img = cv2.resize(Img, (int(Img.shape[1] * r), int(Img.shape[0] * r)))
+   ann_map = cv2.resize(ann_map, (int(ann_map.shape[1] * r), int(ann_map.shape[0] * r)),
+                        interpolation=cv2.INTER_NEAREST)
+ 
+   binary_mask = np.zeros_like(ann_map, dtype=np.uint8)
+   points = []
+   inds = np.unique(ann_map)[1:]
+   for ind in inds:
+       mask = (ann_map == ind).astype(np.uint8)
+       binary_mask = np.maximum(binary_mask, mask)
+ 
+   eroded_mask = cv2.erode(binary_mask, np.ones((5, 5), np.uint8), iterations=1)
+   coords = np.argwhere(eroded_mask > 0)
+   if len(coords) > 0:
+       for _ in inds:
+           yx = np.array(coords[np.random.randint(len(coords))])
+           points.append([yx[1], yx[0]])
+   points = np.array(points)
+ 
+   if visualize_data:
+       plt.figure(figsize=(15, 5))
+       plt.subplot(1, 3, 1)
+       plt.title('Original Image')
+       plt.imshow(Img)
+       plt.axis('off')
+ 
+       plt.subplot(1, 3, 2)
+       plt.title('Binarized Mask')
+       plt.imshow(binary_mask, cmap='gray')
+       plt.axis('off')
+ 
+       plt.subplot(1, 3, 3)
+       plt.title('Binarized Mask with Points')
+       plt.imshow(binary_mask, cmap='gray')
+       colors = list(mcolors.TABLEAU_COLORS.values())
+       for i, point in enumerate(points):
+           plt.scatter(point[0], point[1], c=colors[i % len(colors)], s=100)
+       plt.axis('off')
+ 
+       plt.tight_layout()
+       plt.show()
+ 
+   binary_mask = np.expand_dims(binary_mask, axis=-1)
+   binary_mask = binary_mask.transpose((2, 0, 1))
+   points = np.expand_dims(points, axis=1)
+   return Img, binary_mask, points, len(inds)
+ 
+# Img1, masks1, points1, num_masks = read_batch(train_data, visualize_data=True)
+def _to_hydra_name(x):
+    if not x:
+        return None
+    s = str(x).replace("\\", "/")
+    if s.endswith(".yaml"):
+        s = s[:-5]
+    # Normalize absolute/relative repo paths to hydra names:
+    # /.../sam2/sam2/configs/sam2.1/sam2.1_hiera_s  -> configs/sam2.1/sam2.1_hiera_s
+    # ./sam2/configs/sam2.1/sam2.1_hiera_s          -> configs/sam2.1/sam2.1_hiera_s
+    if "/sam2/configs/" in s:
+        return s.split("/sam2/")[1]       # keep from 'configs/...'
+    if s.startswith("sam2/configs/"):
+        return s[len("sam2/"):]           # strip leading 'sam2/'
+    return s
+
+sam2_checkpoint = "./checkpoints/sam2.1_hiera_large.pt"
+model_cfg = "./sam2/configs/sam2.1/sam2.1_hiera_l.yaml"
+
+model_cfg = _to_hydra_name(model_cfg) 
+sam2_model = build_sam2(model_cfg, sam2_checkpoint, device="cuda")
+predictor = SAM2ImagePredictor(sam2_model)
+ 
+predictor.model.sam_mask_decoder.train(True)
+predictor.model.sam_prompt_encoder.train(True)
+
+scaler = torch.amp.GradScaler()
+NO_OF_STEPS = 1200
+FINE_TUNED_MODEL_NAME = "fine_tuned_sam2"
+ 
+optimizer = torch.optim.AdamW(params=predictor.model.parameters(),
+                              lr=0.00005,
+                              weight_decay=1e-4)
+ 
+scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=1000, gamma=0.6)
+accumulation_steps = 8
+
+def train(predictor, train_data, step, mean_iou):
+    # Ensure rolling mean is numeric
+    if mean_iou is None or (isinstance(mean_iou, float) and (mean_iou != mean_iou)):  # NaN
+        mean_iou = 0.0
+
+    eps = 1e-6
+
+    predictor.model.train()
+    with torch.amp.autocast(device_type='cuda'):
+        image, mask, input_point, num_masks = read_batch(train_data, visualize_data=False)
+
+        # If this batch is unusable, keep the rolling mean unchanged
+        if image is None or mask is None or num_masks == 0:
+            return mean_iou
+
+        input_label = np.ones((num_masks, 1), dtype=np.int64)
+
+        if not isinstance(input_point, np.ndarray) or not isinstance(input_label, np.ndarray):
+            return mean_iou
+        if input_point.size == 0 or input_label.size == 0:
+            return mean_iou
+
+        predictor.set_image(image)
+        mask_input, unnorm_coords, labels, unnorm_box = predictor._prep_prompts(
+            input_point, input_label, box=None, mask_logits=None, normalize_coords=True
+        )
+        if (
+            unnorm_coords is None or labels is None or
+            unnorm_coords.shape[0] == 0 or labels.shape[0] == 0
+        ):
+            return mean_iou
+
+        sparse_embeddings, dense_embeddings = predictor.model.sam_prompt_encoder(
+            points=(unnorm_coords, labels), boxes=None, masks=None
+        )
+
+        batched_mode = unnorm_coords.shape[0] > 1
+        high_res_features = [lvl[-1].unsqueeze(0) for lvl in predictor._features["high_res_feats"]]
+
+        low_res_masks, prd_scores, _, _ = predictor.model.sam_mask_decoder(
+            image_embeddings=predictor._features["image_embed"][-1].unsqueeze(0),
+            image_pe=predictor.model.sam_prompt_encoder.get_dense_pe(),
+            sparse_prompt_embeddings=sparse_embeddings,
+            dense_prompt_embeddings=dense_embeddings,
+            multimask_output=True,
+            repeat_image=batched_mode,
+            high_res_features=high_res_features,
+        )
+
+        prd_masks = predictor._transforms.postprocess_masks(
+            low_res_masks, predictor._orig_hw[-1]
+        )
+
+        gt_mask = torch.tensor(mask.astype(np.float32), device='cuda')
+        prd_mask = torch.sigmoid(prd_masks[:, 0])
+
+        # BCE-style seg loss (numerically stable enough with eps)
+        seg_loss = (-gt_mask * torch.log(prd_mask + eps)
+                    - (1 - gt_mask) * torch.log((1 - prd_mask) + eps)).mean()
+
+        # IoU with safeties
+        pred_bin = (prd_mask > 0.5).float()
+        inter = (gt_mask * pred_bin).sum(dim=(1, 2))
+        denom = gt_mask.sum(dim=(1, 2)) + pred_bin.sum(dim=(1, 2)) - inter
+        iou = inter / (denom + eps)
+
+        score_loss = torch.abs(prd_scores[:, 0] - iou).mean()
+        loss = seg_loss + 0.05 * score_loss
+
+        # grad accumulation
+        loss = loss / accumulation_steps
+        scaler.scale(loss).backward()
+
+        torch.nn.utils.clip_grad_norm_(predictor.model.parameters(), max_norm=1.0)
+
+        did_optimizer_step = False
+        if step % accumulation_steps == 0:
+            # Optimizer step first, then scheduler.step()  (fixes the warning)
+            scaler.step(optimizer)
+            scaler.update()
+            optimizer.zero_grad(set_to_none=True)
+            did_optimizer_step = True
+
+        # Step the LR scheduler only when we actually step the optimizer
+        if did_optimizer_step:
+            scheduler.step()
+
+        # Update rolling mean IoU (robust to NaN/inf)
+        iou_np = iou.detach().float().cpu().numpy()
+        iou_np = np.nan_to_num(iou_np, nan=0.0, posinf=1.0, neginf=0.0)
+        mean_iou = float(mean_iou * 0.99 + 0.01 * float(np.mean(iou_np)))
+
+        if step % 100 == 0:
+            current_lr = optimizer.param_groups[0]["lr"]
+            print(f"Step {step}: LR={current_lr:.6f}  IoU={mean_iou:.6f}  SegLoss={seg_loss.item():.6f}")
+
+    return mean_iou
+
+def validate(predictor, test_data, step, mean_iou):
+    # Always have a numeric baseline
+    if mean_iou is None or (isinstance(mean_iou, float) and (mean_iou != mean_iou)):  # NaN check
+        mean_iou = 0.0
+
+    predictor.model.eval()
+    with torch.amp.autocast(device_type='cuda'):
+        with torch.no_grad():
+            image, mask, input_point, num_masks = read_batch(test_data, visualize_data=False)
+
+            # If this batch is unusable, keep the rolling mean unchanged
+            if image is None or mask is None or num_masks == 0:
+                return mean_iou
+
+            input_label = np.ones((num_masks, 1), dtype=np.int64)
+
+            if not isinstance(input_point, np.ndarray) or not isinstance(input_label, np.ndarray):
+                return mean_iou
+            if input_point.size == 0 or input_label.size == 0:
+                return mean_iou
+
+            predictor.set_image(image)
+            mask_input, unnorm_coords, labels, unnorm_box = predictor._prep_prompts(
+                input_point, input_label, box=None, mask_logits=None, normalize_coords=True
+            )
+
+            if (
+                unnorm_coords is None or labels is None or
+                unnorm_coords.shape[0] == 0 or labels.shape[0] == 0
+            ):
+                return mean_iou
+
+            sparse_embeddings, dense_embeddings = predictor.model.sam_prompt_encoder(
+                points=(unnorm_coords, labels), boxes=None, masks=None
+            )
+
+            batched_mode = unnorm_coords.shape[0] > 1
+            high_res_features = [lvl[-1].unsqueeze(0) for lvl in predictor._features["high_res_feats"]]
+            low_res_masks, prd_scores, _, _ = predictor.model.sam_mask_decoder(
+                image_embeddings=predictor._features["image_embed"][-1].unsqueeze(0),
+                image_pe=predictor.model.sam_prompt_encoder.get_dense_pe(),
+                sparse_prompt_embeddings=sparse_embeddings,
+                dense_prompt_embeddings=dense_embeddings,
+                multimask_output=True,
+                repeat_image=batched_mode,
+                high_res_features=high_res_features,
+            )
+
+            prd_masks = predictor._transforms.postprocess_masks(
+                low_res_masks, predictor._orig_hw[-1]
+            )
+
+            gt_mask = torch.tensor(mask.astype(np.float32), device='cuda')
+            prd_mask = torch.sigmoid(prd_masks[:, 0])
+
+            # BCE-style seg loss
+            eps = 1e-6
+            seg_loss = (-gt_mask * torch.log(prd_mask + eps)
+                        - (1 - gt_mask) * torch.log((1 - prd_mask) + eps)).mean()
+
+            # IoU with numerical safety
+            pred_bin = (prd_mask > 0.5).float()
+            inter = (gt_mask * pred_bin).sum(dim=(1, 2))
+            denom = gt_mask.sum(dim=(1, 2)) + pred_bin.sum(dim=(1, 2)) - inter
+            iou = inter / (denom + eps)  # avoid 0/0
+
+            # Score loss
+            score_loss = torch.abs(prd_scores[:, 0] - iou).mean()
+            loss = seg_loss + 0.05 * score_loss
+            loss = loss / accumulation_steps  # assumes defined elsewhere
+
+            if step % 100 == 0:
+                torch.save(predictor.model.state_dict(), f"./checkpoints-ft/{FINE_TUNED_MODEL_NAME}_{step}.pt")
+
+            iou_np = iou.detach().float().cpu().numpy()
+            iou_np = np.nan_to_num(iou_np, nan=0.0, posinf=1.0, neginf=0.0)
+            mean_iou = float(mean_iou * 0.99 + 0.01 * float(np.mean(iou_np)))
+
+            if step % 100 == 0:
+                current_lr = optimizer.param_groups[0]["lr"]
+                print(f"Step {step}: LR={current_lr:.6f}  Valid_IoU={mean_iou:.6f}  SegLoss={seg_loss.item():.6f}")
+
+    return mean_iou
+
+train_mean_iou = 0
+valid_mean_iou = 0
+ 
+# for step in range(1, NO_OF_STEPS + 1):
+#     train_mean_iou = train(predictor, train_data, step, train_mean_iou)
+#     valid_mean_iou = validate(predictor, test_data, step, valid_mean_iou)
+
+def read_image(image_path, mask_path):  # read and resize image and mask
+   img = cv2.imread(image_path)[..., ::-1]  # Convert BGR to RGB
+   mask = cv2.imread(mask_path, 0)
+   r = np.min([1024 / img.shape[1], 1024 / img.shape[0]])
+   img = cv2.resize(img, (int(img.shape[1] * r), int(img.shape[0] * r)))
+   mask = cv2.resize(mask, (int(mask.shape[1] * r), int(mask.shape[0] * r)), interpolation=cv2.INTER_NEAREST)
+   return img, mask
+ 
+def get_points(mask, num_points):  # Sample points inside the input mask
+   points = []
+   coords = np.argwhere(mask > 0)
+   for i in range(num_points):
+       yx = np.array(coords[np.random.randint(len(coords))])
+       points.append([[yx[1], yx[0]]])
+   return np.array(points)
+
+for n in range(3):
+    selected_entry = random.choice(test_data)
+    print(selected_entry)
+image_path = selected_entry['image']
+mask_path = selected_entry['annotation']
+print(mask_path,'mask path')
+ 
+# Load the selected image and mask
+image, target_mask = read_image(image_path, mask_path)
+ 
+# Generate random points for the input
+num_samples = 30  # Number of points per segment to sample
+input_points = get_points(target_mask, num_samples)
+ 
+# Load the fine-tuned model
+FINE_TUNED_MODEL_WEIGHTS = "./checkpoints-ft/fine_tuned_sam2_1200.pt"
+sam2_model = build_sam2(model_cfg, sam2_checkpoint, device="cuda")
+ 
+# Build net and load weights
+predictor = SAM2ImagePredictor(sam2_model)
+predictor.model.load_state_dict(torch.load(FINE_TUNED_MODEL_WEIGHTS))
+ 
+ 
+ 
+# Perform inference and predict masks
+with torch.no_grad():
+   predictor.set_image(image)
+   masks, scores, logits = predictor.predict(
+       point_coords=input_points,
+       point_labels=np.ones([input_points.shape[0], 1])
+   )
+ 
+# Process the predicted masks and sort by scores
+np_masks = np.array(masks[:, 0])
+np_scores = scores[:, 0]
+sorted_masks = np_masks[np.argsort(np_scores)][::-1]
+ 
+# Initialize segmentation map and occupancy mask
+seg_map = np.zeros_like(sorted_masks[0], dtype=np.uint8)
+occupancy_mask = np.zeros_like(sorted_masks[0], dtype=bool)
+ 
+# Combine masks to create the final segmentation map
+for i in range(sorted_masks.shape[0]):
+   mask = sorted_masks[i]
+   if (mask * occupancy_mask).sum() / mask.sum() > 0.15:
+       continue
+ 
+   mask_bool = mask.astype(bool)
+   mask_bool[occupancy_mask] = False  # Set overlapping areas to False in the mask
+   seg_map[mask_bool] = i + 1  # Use boolean mask to index seg_map
+   occupancy_mask[mask_bool] = True  # Update occupancy_mask
+ 
+# Visualization: Show the original image, mask, and final segmentation side by side
+plt.figure(figsize=(18, 6))
+ 
+plt.subplot(1, 3, 1)
+plt.title('Test Image')
+plt.imshow(image)
+plt.axis('off')
+ 
+plt.subplot(1, 3, 2)
+plt.title('Original Mask')
+plt.imshow(target_mask, cmap='gray')
+plt.axis('off')
+ 
+plt.subplot(1, 3, 3)
+plt.title('Final Segmentation')
+plt.imshow(seg_map, cmap='jet')
+plt.axis('off')
+ 
+plt.tight_layout()
+plt.show()