models/sam/

[giscus[bot]]

models/sam/

Explore the cutting-edge Segment Anything Model (SAM) from Ultralytics that allows real-time image segmentation. Learn about its promptable segmentation, zero-shot performance, and how to use it.

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

[pabentse]

Would be fantastic with support for the "H" size as well!

[pderrenger]

@pabentse absolutely, the addition of an "H" size model would be a great enhancement to the range of model sizes we offer, catering to different computational needs and use cases. The YOLO community's feedback is invaluable, and suggestions like yours help guide future developments. Keep an eye on our updates! 👀🚀

[kshweikeh]

is there anyway where I can train YOLO SAM on a custom dataset?

[pderrenger]

@kshweikeh hi there!

Currently, the Segment Anything Model (SAM) from Ultralytics is designed primarily for inference and does not support training on custom datasets. However, you can leverage SAM's powerful zero-shot capabilities for various segmentation tasks without the need for additional training.

If you're looking to create a segmentation dataset using SAM, you can use the auto-annotation feature. This allows you to generate segmentation masks using a pre-trained detection model. Here's a quick example:

from ultralytics.data.annotator import auto_annotate

auto_annotate(data="path/to/images", det_model="yolov8x.pt", sam_model="sam_b.pt")

This function will help you annotate your images quickly and efficiently, bypassing the need for manual labeling.

For more detailed information on using SAM, you can visit our Segment Anything Model (SAM) documentation.

If you have any further questions or need additional assistance, feel free to ask! 😊

[Blessing988]

@pderrenger It looks like auto_annotate for SAM only gives out the segmentation masks. What if you want to generate the bounding boxes in additionto the masks to verify if the results are right. Thanks.

[Blessing988]

Also, how will one visualise the segmentation masks after the auto_annotate which have been saved in a directory of txt files.

[pderrenger]

Hi @Blessing988,

Great question! To visualize the segmentation masks generated by the auto_annotate function, you can use the following approach to overlay the masks on the original images and optionally generate bounding boxes for verification.

Here's a Python script to help you visualize the masks and bounding boxes:

import os
import cv2
import numpy as np
from pathlib import Path

def visualize_annotations(image_dir, label_dir, output_dir):
    os.makedirs(output_dir, exist_ok=True)
    
    for image_path in Path(image_dir).glob('*.jpg'):
        image = cv2.imread(str(image_path))
        label_path = Path(label_dir) / (image_path.stem + '.txt')
        
        if not label_path.exists():
            continue
        
        with open(label_path, 'r') as f:
            for line in f:
                parts = line.strip().split()
                cls, x_center, y_center, width, height = map(float, parts)
                
                # Convert YOLO format to bounding box coordinates
                x1 = int((x_center - width / 2) * image.shape[1])
                y1 = int((y_center - height / 2) * image.shape[0])
                x2 = int((x_center + width / 2) * image.shape[1])
                y2 = int((y_center + height / 2) * image.shape[0])
                
                # Draw bounding box
                cv2.rectangle(image, (x1, y1), (x2, y2), (0, 255, 0), 2)
                
                # Optionally, draw the segmentation mask (if available)
                mask_path = Path(label_dir) / (image_path.stem + f'_{int(cls)}.png')
                if mask_path.exists():
                    mask = cv2.imread(str(mask_path), cv2.IMREAD_GRAYSCALE)
                    mask = cv2.resize(mask, (image.shape[1], image.shape[0]))
                    image[mask > 0] = [0, 0, 255]  # Overlay mask in red

        # Save the visualized image
        output_path = Path(output_dir) / image_path.name
        cv2.imwrite(str(output_path), image)

# Example usage
visualize_annotations('path/to/images', 'path/to/labels', 'path/to/output')

This script reads the images and corresponding label files, draws the bounding boxes, and overlays the segmentation masks if available. The visualized images are saved in the specified output directory.

For more detailed information on using SAM and its features, you can refer to the Segment Anything Model (SAM) documentation.

Feel free to reach out if you have any further questions or need additional assistance! 😊

[dreambit]

I have images where a single object (such as a car) is partially covered by small trees, resulting in the object being divided into multiple visible fragments. In coco it could look like
{
"id": 1,
"image_id": 1,
"category_id": 1,
"bbox": [x, y, width, height],
"segmentation": [
[polygon1_x1, polygon1_y1, polygon1_x2, polygon1_y2, ...],
[polygon2_x1, polygon2_y1, polygon2_x2, polygon2_y2, ...]
],
"area": 500,
"iscrowd": 0
}
Each polygon is an array of coordinates defining a closed shape

Is it possible with auto_annotate and yolo format?

[glenn-jocher]

@dreambit hi there!

Thank you for your detailed question. Yes, you can use the auto_annotate function in the Ultralytics framework to handle images where objects are partially occluded and appear as multiple fragments. This function can generate segmentation masks for such complex scenarios.

Here's how you can achieve this:

1. Prepare Your Dataset: Ensure your images are organized in a directory.
2. Use the auto_annotate Function: This function will use a pre-trained YOLO detection model and the SAM segmentation model to annotate your images.

Here's a code example to get you started:

from ultralytics.data.annotator import auto_annotate

# Path to your images
data_path = "path/to/your/images"

# Auto-annotate using YOLOv8 and SAM models
auto_annotate(data=data_path, det_model="yolov8x.pt", sam_model="sam_b.pt")

Explanation:

• data: Path to the folder containing your images.
• det_model: The pre-trained YOLO detection model, which defaults to yolov8x.pt.
• sam_model: The pre-trained SAM segmentation model, which defaults to sam_b.pt.

This setup will generate segmentation masks for your images, even if the objects are partially occluded.

For more detailed information on using SAM and its capabilities, you can refer to the Segment Anything Model (SAM) documentation.

If you encounter any issues or need further assistance, please provide a reproducible example so we can better understand and address your specific scenario. You can find more information on creating a minimum reproducible example here.

Feel free to update to the latest versions of the packages to ensure you have the latest features and fixes.

Best of luck with your project! 🚀

[dreambit]

Glenn thanks, but as I see in docs https://docs.ultralytics.com/datasets/segment/#ultralytics-yolo-format
...
In this format, is the index of the class for the object, and ... are the bounding coordinates of the object's segmentation mask. The coordinates are separated by spaces.

So it means that I could have only one segmentation mask?

[dreambit]

0 0.681 0.485 0.670 0.487 0.676 0.487
1 0.504 0.000 0.501 0.004 0.498 0.004 0.493 0.010 0.492 0.0104

[dreambit]

Glen, i ran the command you specified, and used an image from coco data set to test, as you can see segmentations are not correct, only half of the objects are segmented:

https://ibb.co/P4Pt2NT

Here is original image from coco with car mask
https://ibb.co/RjYyRdk

[glenn-jocher]

Hi @dreambit,

Thank you for sharing the images and details. I understand the challenge you're facing with partial segmentations. The auto_annotate function should ideally handle such cases, but it seems like the segmentation results are not as expected.

Steps to Troubleshoot and Improve Segmentation:

1. Verify Latest Versions: Ensure you are using the latest versions of the Ultralytics and SAM models. This can help avoid any bugs that might have been fixed in recent updates.

2. Reproducible Example: If the issue persists, providing a minimum reproducible example would be very helpful. This allows us to better understand and diagnose the problem. You can find guidance on creating a reproducible example here.

3. Adjusting Prompts: Sometimes, adjusting the prompts (bounding boxes or points) can improve segmentation results. Here’s an example of how you can use different prompts:

   from ultralytics import SAM
   
   # Load a model
   model = SAM("sam_b.pt")
   
   # Run inference with bounding boxes prompt
   results = model("path/to/image.jpg", bboxes=[[x1, y1, x2, y2], [x3, y3, x4, y4]])
   
   # Run inference with points prompt
   results = model("path/to/image.jpg", points=[[x1, y1], [x2, y2]], labels=[1, 1])

4. Segment Everything: If specific prompts are not yielding the desired results, you can try segmenting the entire image without prompts:

   from ultralytics import SAM
   
   # Load a model
   model = SAM("sam_b.pt")
   
   # Segment the whole image
   results = model("path/to/image.jpg")

5. Multiple Masks Handling: If your object is divided into multiple fragments, ensure that the segmentation model is capable of handling multiple masks per object. The YOLO format typically supports one segmentation mask per object, but you can adapt the output to handle multiple masks by combining them post-processing.

Example of Handling Multiple Masks:

If you need to combine multiple masks into a single annotation, you can do this post-processing:

import numpy as np

# Assuming `results` contains multiple masks for the same object
combined_mask = np.zeros_like(results[0].mask)
for result in results:
    combined_mask = np.maximum(combined_mask, result.mask)

# Save or use the combined mask

Documentation and Resources:

For more detailed information on using SAM and its capabilities, please refer to the Segment Anything Model (SAM) documentation.

If you continue to face issues, please share a reproducible example, and we’ll be happy to assist further. Your feedback is valuable, and we aim to provide the best support possible.

[kh-diab]

can we use SAM to process labeled images with bounding box to segment the object inside the box and generate new labels files like yolo or coco format
if yes how please

[pderrenger]

Yes, you can use SAM to process labeled images with bounding boxes to segment the objects inside and generate new label files in formats like YOLO or COCO. You can achieve this by using the auto_annotate function in the Ultralytics framework. This function allows you to auto-annotate your dataset using a pre-trained detection model and SAM for segmentation. For more details, you can refer to the Segment Anything Model documentation. If you need further assistance, feel free to ask!

[kh-diab]

Hello @pderrenger

Thanks
I tried to run the below but I face error all the time

import os
import cv2
import numpy as np
from ultralytics import SAM

Load the SAM model

model = SAM("sam_b.pt") # Update the model path if needed

Function to load YOLO annotations

def load_yolo_annotations(annotation_path, image_width, image_height):
bboxes = []
try:
with open(annotation_path, 'r') as file:
lines = file.readlines()
for line in lines:
parts = line.strip().split()
class_id = int(parts[0])
x_center = float(parts[1]) * image_width
y_center = float(parts[2]) * image_height
width = float(parts[3]) * image_width
height = float(parts[4]) * image_height
x = x_center - width / 2
y = y_center - height / 2
bboxes.append((x, y, x + width, y + height)) # Update to [x1, y1, x2, y2]
except Exception as e:
print(f"Error reading annotation file {annotation_path}: {e}")
return bboxes

Paths

images_folder = r'E:\MSc\tesi\20210901_Herrenpuent_Taenikon_Matrice_300_15m_1.6m_s\500 D'
annotations_folder = r'C:\Users\Khaled\OneDrive - OST\Study Material\Msc Thesis\dataset\Group D yolo format\obj_train_data'
output_folder = r'C:\Users\Khaled\OneDrive - OST\Study Material\Msc Thesis\dataset\test output sam'

if not os.path.exists(output_folder):
os.makedirs(output_folder)

Process each image and its corresponding annotation

total_images = len(os.listdir(images_folder))
processed_images = 0

for image_filename in os.listdir(images_folder):
if image_filename.endswith(('.JPG', '.jpg', '.png')):
image_path = os.path.join(images_folder, image_filename)
annotation_path = os.path.join(annotations_folder, image_filename.replace('.JPG', '.txt').replace('.jpg', '.txt').replace('.png', '.txt'))

    print(f"Processing {image_filename}...")

    # Check if annotation file exists
    if not os.path.exists(annotation_path):
        print(f"Annotation file not found for {image_filename}")
        continue

    # Load the image
    image = cv2.imread(image_path)
    if image is None:
        print(f"Failed to load image {image_path}")
        continue

    image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    height, width, _ = image.shape

    # Load YOLO annotations
    bboxes = load_yolo_annotations(annotation_path, width, height)
    if not bboxes:
        print(f"No bounding boxes found in {annotation_path}")
        continue

    # Perform segmentation using SAM model
    results = model(image_rgb, bboxes=bboxes)

    for idx, result in enumerate(results):
        if hasattr(result, 'masks') and result.masks is not None:
            masks = result.masks.cpu().numpy()  # Convert to numpy array
            for m_idx, mask in enumerate(masks):
                chw_mask = mask.astype(np.uint8)  # Ensure mask is in uint8 format
                hwc_mask = np.transpose(chw_mask, (1, 2, 0))  # Convert CHW to HWC
                contours, _ = cv2.findContours(hwc_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
                
                segmented_image = cv2.bitwise_and(image, image, mask=hwc_mask)

                # Save the segmented image
                output_path = os.path.join(output_folder, image_filename.replace('.JPG', f'_segmented_{idx}_{m_idx}.jpg').replace('.jpg', f'_segmented_{idx}_{m_idx}.jpg').replace('.png', f'_segmented_{idx}_{m_idx}.png'))
                cv2.imwrite(output_path, segmented_image)
                print(f"Saved segmented image to {output_path}")

    processed_images += 1
    print(f"Processed {processed_images}/{total_images} images.")

print("Processing completed.")

error is

AttributeError: 'Masks' object has no attribute 'astype'. See valid attributes below.

A class for storing and manipulating detection masks.

    xy (list): A list of segments in pixel coordinates.
    xyn (list): A list of normalized segments.

Methods:
    cpu(): Returns the masks tensor on CPU memory.
    numpy(): Returns the masks tensor as a numpy array.
    cuda(): Returns the masks tensor on GPU memory.
    to(device, dtype): Returns the masks tensor with the specified device and dtype.

PS C:\Users\Khaled> & C:/MSc_Workspace/anaconda3/envs/GenEnv/python.exe "c:/Users/Khaled/OneDrive - OST/Study Material/Msc Thesis/coding/sam segment.py"
Processing DJI_20210901122818_0695.JPG...

0: 1024x1024 67330.4ms
Speed: 48.0ms preprocess, 67330.4ms inference, 8095.5ms postprocess per image at shape (1, 3, 1024, 1024)
Traceback (most recent call last):
File "c:\Users\Khaled\OneDrive - OST\Study Material\Msc Thesis\coding\sam segment.py", line 75, in
chw_mask = mask.astype(np.uint8) # Ensure mask is in uint8 format
File "C:\MSc_Workspace\anaconda3\envs\GenEnv\lib\site-packages\ultralytics\utils_init_.py", line 162, in getattr
raise AttributeError(f"'{name}' object has no attribute '{attr}'. See valid attributes below.\n{self.doc}")
AttributeError: 'Masks' object has no attribute 'astype'. See valid attributes below.

A class for storing and manipulating detection masks.

Attributes:
    xy (list): A list of segments in pixel coordinates.
    xyn (list): A list of normalized segments.

Methods:
    cpu(): Returns the masks tensor on CPU memory.
    numpy(): Returns the masks tensor as a numpy array.
    cuda(): Returns the masks tensor on GPU memory.
    to(device, dtype): Returns the masks tensor with the specified device and dtype.

do you have idea how to solve it thanks

[NasserAlqeifi]

Correct me if I'm wrong, when I use the auto_annotate function I will only get labels for my dataset with the classes in Yolo pre-trained? What if I need classes that are not available within Yolo? For example, I want to get a label for the road, and this is not available in any version of Yolo, but it is available in SAM?

[glenn-jocher]

@NasserAlqeifi you are correct. The auto_annotate function uses YOLO's pre-trained classes for labeling. If you need classes not available in YOLO, like "road," you can use SAM directly for segmentation. For more details, please refer to the Segment Anything Model documentation.

[glenn-jocher]

@NasserAlqeifi but also remember SAM identifies segments but not classes, so it may find a sign or a road, but it won't be able to assign a class to them by itself.

[NasserAlqeifi]

Thank you for your quick response. Indeed, I tried to use SAM directly and found that it does not save the class names in the name, but the class number.

from ultralytics import SAM

Load a model

model = SAM("sam2_l.pt")

Predict results

results = model.predict(source=r"/testSAM/image.png", stream=True, bboxes=None, points=None, labels=None, show=False)

Iterate over results, print names, and save each result to a file

for i, result in enumerate(results):
# Print the names detected in the result
print(f"Result {i}: Detected Names - {result.names}")

# Save the result as an image file
result.save(f"output_{i}.jpg")

------------- output ---------

image 1/1 /testSAM/image.png: 1024x1024 1 0, 1 1, 1 2, 1 3, 1 4, 1 5, 1 6, 1 7, 1 8, 1 9, 1 10, 1 11, 1 12, 1 13, 1 14, 1 15, 1 16, 1 17, 1 18, 1 19, 1 20, 1 21, 1 22, 1 23, 1 24, 1 25, 1 26, 1 27, 1 28, 1 29, 1 30, 1 31, 1 32, 23758.2ms
Result 0: Detected Names - {0: '0', 1: '1', 2: '2', 3: '3', 4: '4', 5: '5', 6: '6', 7: '7', 8: '8', 9: '9', 10: '10', 11: '11', 12: '12', 13: '13', 14: '14', 15: '15', 16: '16', 17: '17', 18: '18', 19: '19', 20: '20', 21: '21', 22: '22', 23: '23', 24: '24', 25: '25', 26: '26', 27: '27', 28: '28', 29: '29', 30: '30', 31: '31', 32: '32'}
Speed: 20.3ms preprocess, 23758.2ms inference, 7.1ms postprocess per image at shape (1, 3, 1024, 1024)

[glenn-jocher]

Thank you for your feedback. Currently, SAM does not assign class names, only class numbers. If you need specific class names, you may need to implement a custom mapping or post-processing step to assign the desired class names based on your requirements.

[MeenakshiKK236]

Why are my segmentation results being automatically saved into this directory while using the SAM model to perform segmentation on my custom dataset? "Results saved to runs/segment/predict4"

[glenn-jocher]

@MeenakshiKK236 the results are saved to the runs/segment/predict4 directory by default. You can change the save directory by specifying the project and name arguments in your command. For more details, please refer to the SAM documentation.

[meenakshi71]

I've been trying to export the SAM model (which i had trained on my custom dataset) into ONNX but the output is always wrong. It is just generating a grid instead of a proper segmentation mask every time even after providing multiple kinds of prompts.

[glenn-jocher]

@meenakshi71 it seems the issue may stem from the export process or configuration. Ensure you are using the correct export parameters and supported formats. SAM models are primarily designed for real-time inference and not all functionalities may transfer seamlessly to ONNX. Refer to the SAM export documentation for guidance. If the problem persists, consider testing the ONNX model with simpler prompts to isolate the issue.