152 lines
4.1 KiB
Python
152 lines
4.1 KiB
Python
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") |