ultralytics / ultralytics

Ultralytics YOLO11 🚀
https://docs.ultralytics.com
GNU Affero General Public License v3.0
36.58k stars 7.05k forks source link

Problem with SAM2 model #19002

Open TasTab007 opened 1 week ago

TasTab007 commented 1 week ago

I am trying to use SAM2 from Ultralytics but it's not working. Is SAM2 is being updated? Why it's not working?

UltralyticsAssistant commented 1 week ago

👋 Hello @TasTab007, thank you for your interest in Ultralytics 🚀! We recommend a visit to the Docs for new users where you can find many Python and CLI usage examples and where many of the most common questions may already be answered.

If this is a 🐛 Bug Report, please provide a minimum reproducible example to help us debug it. This will significantly help us understand the problem and replicate the issue.

If this is a custom training ❓ Question, please provide as much information as possible, including dataset image examples, training logs, and details about your environment. Verify that you are following our Tips for Best Training Results.

Join the Ultralytics community where it suits you best. For real-time chat, head to Discord 🎧. Prefer in-depth discussions? Check out Discourse. Or dive into threads on our Subreddit to share knowledge with the community.

Upgrade

Please ensure that you are using the latest version of the ultralytics package, along with all requirements, in a Python>=3.8 environment with PyTorch>=1.8. Run the following command to update:

pip install -U ultralytics

Environments

YOLO may be run in any of the following up-to-date verified environments (with all dependencies including CUDA/CUDNN, Python and PyTorch preinstalled):

Status

Ultralytics CI

If this badge is green, all Ultralytics CI tests are currently passing. CI tests verify the correct operation of all YOLO Modes and Tasks on macOS, Windows, and Ubuntu every 24 hours and on every commit.

This is an automated response, but please be assured that an Ultralytics engineer will assist you soon. Feel free to provide more details in the meantime, such as any error messages or the steps you followed when encountering the issue. 🚀

TasTab007 commented 1 week ago

I updated the version to newest one and still getting this error:

Image

Y-T-G commented 1 week ago

You need to provide the code and the full error traceback

TasTab007 commented 1 week ago

Here is the code: import os import cv2 import torch import json import logging import time import numpy as np import matplotlib.pyplot as plt from ultralytics import SAM from sklearn.metrics import precision_recall_curve, auc from pycocotools import mask as mask_utils # Import pycocotools for RLE decoding

Set up logging

logging.basicConfig(level=logging.INFO)

Paths for dataset

DATASET_PATH = r"/home/myid/tt59612/Panoptic Segmentation/Dataset/Coral Net.v8-final-dataset.sam2" TRAIN_PATH = os.path.join(DATASET_PATH, "train") VALID_PATH = os.path.join(DATASET_PATH, "valid") TEST_PATH = os.path.join(DATASET_PATH, "test")

Pre-trained SAM2 model path

SAM_MODEL_PATH = r"/home/myid/tt59612/Panoptic Segmentation/sam_l.pt" # Update this path to your SAM2 weights

Set device (Use GPU if available)

device = "cuda" if torch.cuda.is_available() else "cpu" logging.info(f"Using device: {device}")

Output folder for results

RESULTS_DIR = "sam2_fine_tuned_result" SEGMENTED_IMAGES_DIR = os.path.join(RESULTS_DIR, "segmented_images") os.makedirs(SEGMENTED_IMAGES_DIR, exist_ok=True)

Log file for storing evaluation results

EVAL_LOG_FILE = os.path.join(RESULTS_DIR, "evaluation_log.txt")

Function to decode RLE segmentation masks

def decode_rle(segmentation): """Decodes RLE segmentation mask to a binary mask.""" if isinstance(segmentation, dict) and "counts" in segmentation and "size" in segmentation: return mask_utils.decode(segmentation).astype(np.uint8) * 255 # Convert to 0-255 range return None

Function to load images and masks from JSON files

def load_data(image_folder): image_files = [f for f in os.listdir(image_folder) if f.endswith((".jpg", ".png", ".jpeg"))] data = []

for image_file in image_files:
    image_path = os.path.join(image_folder, image_file)
    json_path = os.path.splitext(image_path)[0] + ".json"

    if not os.path.exists(json_path):
        logging.warning(f"Skipping {image_file}: No corresponding JSON file found.")
        continue

    # Load image
    img = cv2.imread(image_path)
    if img is None:
        logging.warning(f"Skipping {image_file}: Unable to read image.")
        continue

    # Load JSON annotation
    with open(json_path, "r") as f:
        annotations = json.load(f)

    # Extract instance segmentation masks
    masks = []
    for ann in annotations.get("annotations", []):
        segmentation = ann.get("segmentation", None)

        if segmentation is None:
            logging.warning(f"Skipping {image_file}: No segmentation data found.")
            continue

        # Handle RLE-encoded masks
        if isinstance(segmentation, dict) and "counts" in segmentation:
            mask = decode_rle(segmentation)
            if mask is not None:
                masks.append(mask)
            else:
                logging.warning(f"Skipping {image_file}: Failed to decode RLE mask.")

        # Handle polygon segmentation
        elif isinstance(segmentation, list):  
            mask = np.zeros(img.shape[:2], dtype=np.uint8)
            for poly in segmentation:
                pts = np.array(poly, np.int32).reshape((-1, 2))
                cv2.fillPoly(mask, [pts], 255)
            masks.append(mask)

        else:
            logging.warning(f"Skipping {image_file}: Unknown segmentation format.")

    if len(masks) > 0:
        data.append((img, masks))

return data

try:

Load SAM2 model

sam_model = SAM(SAM_MODEL_PATH).to(device)

# Load training and validation data
logging.info("Loading training data...")
train_data = load_data(TRAIN_PATH)
logging.info(f"Loaded {len(train_data)} training samples.")

logging.info("Loading validation data...")
valid_data = load_data(VALID_PATH)
logging.info(f"Loaded {len(valid_data)} validation samples.")

# Fine-tune SAM2 on the dataset
logging.info("Starting SAM2 fine-tuning...")
sam_model.train(data=train_data, val=valid_data, epochs=30, imgsz=650)

# Save the fine-tuned model
fine_tuned_sam_model = os.path.join(RESULTS_DIR, "fine_tuned_sam2.pt")
sam_model.save(fine_tuned_sam_model)
logging.info(f"✅ SAM2 fine-tuning completed. Model saved at {fine_tuned_sam_model}")

# Apply SAM2 on test images
logging.info("Applying SAM2 instance segmentation on test images...")

# Metrics storage
iou_scores = []
dice_scores = []
processing_times = []

for image_name in os.listdir(TEST_PATH):
    image_path = os.path.join(TEST_PATH, image_name)

    if not image_path.endswith((".jpg", ".png", ".jpeg")):
        continue

    img = cv2.imread(image_path)
    if img is None:
        logging.warning(f"Skipping {image_name}: Unable to read image.")
        continue

    start_time = time.time()
    results = sam_model(img)
    end_time = time.time()

    if isinstance(results, list) and len(results) > 0:
        segmented_image = results[0].masks
    else:
        segmented_image = None

    # Save segmented images
    output_path = os.path.join(SEGMENTED_IMAGES_DIR, image_name)
    if segmented_image is not None:
        cv2.imwrite(output_path, segmented_image)
    else:
        cv2.imwrite(output_path, img)

    # Record processing time
    processing_times.append(end_time - start_time)

# Compute FPS and Memory Usage
avg_fps = len(processing_times) / sum(processing_times)
gpu_memory = torch.cuda.memory_allocated() / (1024 ** 3)

# Save evaluation log
with open(EVAL_LOG_FILE, "w") as f:
    f.write("🟢 Evaluation Results:\n")
    f.write(f"📌 FPS (Inference Speed): {avg_fps:.2f} frames/sec\n")
    f.write(f"📌 GPU Memory Usage: {gpu_memory:.2f} GB\n")

logging.info(f"✅ Segmentation and evaluation completed. Results saved in: {RESULTS_DIR}")

except Exception as e: logging.error(f"❌ An error occurred: {str(e)}")

Y-T-G commented 1 week ago

You can't train SAM in Ultralytics.

Also follow Ultralytics guide instead of AI-generated code.

https://docs.ultralytics.com/models/sam-2/