SAM2 (Segment Anything Model 2) Tutorial

!pip install dtlpy matplotlib numpy Pillow --upgrade --quiet

# Import necessary libraries
import dtlpy as dl
import random
import matplotlib.pyplot as plt
import numpy as np
from PIL import Image
import io

# Login to the Dataloop platform
# This will check if your current token is valid and prompt for login if necessary.
if dl.token_expired():
    dl.login()
print('Successfully logged into Dataloop')

# Get the project
# IMPORTANT: Modify 'your_project_name' to your specific Dataloop project name.
project = dl.projects.get(project_name='your_project_name')
print(f'Connected to project: {project.name}')

# Get the dataset
# IMPORTANT: Modify 'your_dataset_name' to your specific Dataloop dataset name.
dataset = project.datasets.get(dataset_name='your_dataset_name')
print(f'Connected to dataset: {dataset.name}')

# Get a random item from the dataset
items_list = list(dataset.items.list().all)
if items_list:
    item = items_list[random.randint(0, len(items_list) - 1)]
    print(f"Selected random item: {item.name} (ID: {item.id})") 
else:
    print("No items found in the dataset!")
    item = None

def display_item_with_annotations(item):
    """
    Display an image item and its annotations
    
    Args:
        item: A Dataloop item object
    """
    if not item:
        print("No item to display")
        return
        
    # Download the image locally
    buffer = item.download(save_locally=False)
    
    # Convert buffer to image
    image = Image.open(io.BytesIO(buffer))
    plt.figure(figsize=(10, 10))
    plt.imshow(image)
    plt.axis('off')
    plt.title(f"Image: {item.name}")
    plt.show()
    
    # Get annotations for the item
    annotations = item.annotations.list()
    if annotations.items:
        print(f"Found {len(annotations.items)} annotations")
        
        # Display image with annotations
        plt.figure(figsize=(10, 10))
        plt.imshow(image)
        
        # Plot each annotation
        for annotation in annotations.items:
            if hasattr(annotation, 'coordinates'):
                coords = np.array(annotation.coordinates)
                if annotation.type == 'box':
                    x, y, w, h = coords
                    plt.gca().add_patch(plt.Rectangle((x, y), w, h, fill=False, edgecolor='red', linewidth=2))
                    plt.text(x, y, annotation.label, color='white', backgroundcolor='red', fontsize=8)
                elif annotation.type == 'polygon':
                    plt.plot(coords[:, 0], coords[:, 1], 'r-', linewidth=2)
                    plt.text(coords[0, 0], coords[0, 1], annotation.label, color='white', backgroundcolor='red', fontsize=8)
                elif annotation.type == 'point':
                    plt.scatter(coords[0], coords[1], c='red', s=50)
                    plt.text(coords[0], coords[1], annotation.label, color='white', backgroundcolor='red', fontsize=8)
                elif annotation.type == 'segmentation':
                    # Assuming annotation.geo is a Dataloop geo object for segmentation
                    # The geo_to_mask method requires image height and width
                    mask = annotation.geo.geo_to_mask((image.height, image.width)) 
                    plt.imshow(mask, alpha=0.5, cmap='jet')
                    
        plt.axis('off')
        plt.title("Image with Annotations")
        plt.show()
    else:
        print("No annotations found for this item")

# Show the item and item annotations in the notebook
display_item_with_annotations(item)

# Show the item and item annotations in the platform
# This will open the item in your web browser, in the Dataloop annotation studio.
item.open_in_web()

# Get the SAM2 DPK (Dataloop Package)
sam2_dpk = dl.dpks.get(dpk_name='sam2')

# Install the SAM2 app from the DPK if not already installed in the project
try:
    sam2_app = project.apps.install(app_name=sam2_dpk.display_name, dpk=sam2_dpk, custom_installation=sam2_dpk.to_json())
except dl.exceptions.BadRequest as e: # Catch error if app is already installed
    print(f"App {sam2_dpk.display_name} already installed, getting existing app instance.")
    sam2_app = project.apps.get(app_name=sam2_dpk.display_name)

# Get the SAM2 model associated with the installed app
# Models are typically created when an app is installed or can be managed separately.
filters = dl.Filters(resource=dl.FiltersResource.MODEL)
filters.add(field='app.id', values=sam2_app.id)
model = project.models.list(filters=filters)[0][0]

print(f'\nLoaded model: {model.name}')

# Check model status and deploy if not already deployed
if model.status != 'deployed':
    print('Deploying model...')
    model.deploy() # This deploys the model with default settings
print(f'Model status: {model.status}')

# Display model configuration for reference
print('\nModel configuration:')
for key, value in model.configuration.items():
    print(f'{key}: {value}')

# Make prediction
prediction = model.predict(item_ids=[item.id])
prediction.wait() # Wait for the prediction process to complete
print('Prediction completed successfully')

# Get updated item with predictions
# The item in the dataset will now have new annotations generated by the model.
updated_item = dataset.items.get(item_id=item.id)

# Show the item and item annotations (including predictions) in the notebook
display_item_with_annotations(updated_item)

# Show the item and item annotations (including predictions) in the platform
updated_item.open_in_web()

# Run this cell to open the item in the Dataloop UI for manual annotation
# After adding and saving a bounding box in the UI, come back to this notebook.
item.open_in_web()

# Define and upload a bounding box annotation
image_annotation = dl.AnnotationCollection()
image_annotation.add(annotation_definition=dl.Box(left=float(100),  # left coordinate in pixels from the left edge of the image
                                                top=float(100),    # top coordinate in pixels from the top edge of the image
                                                right=float(200),  # right coordinate in pixels from the left edge of the image
                                                bottom=float(200), # bottom coordinate in pixels from the top edge of the image
                                                label='your_label' # Assign a meaningful label to your prompt box
                                                ))
item.annotations.upload(image_annotation)

# Fetch the item again to get any annotations added via UI or code
annotated_item = dataset.items.get(item_id=item.id)

# Make prediction using the item with the prompt
prediction = model.predict(item_ids=[annotated_item.id])
prediction.wait()
print('Prediction completed successfully')

# Get updated item with predictions based on the prompt
updated_item = dataset.items.get(item_id=item.id)

# Show the item and its annotations (including new predictions) in the notebook
display_item_with_annotations(updated_item)

# Show the item and its annotations in the platform
updated_item.open_in_web()

# Clone the base SAM2 model to create a trainable version
# It's good practice to give the cloned model a descriptive name.
clone_model = model.clone(model_name=f"{model.name}-finetuned", # Or provide a new name like f"{model.name}-finetune"
                          dataset=dataset, # Associate the dataset for training context
                          description='Cloned model for training')

print('Model cloned successfully')

# Display the configuration of the cloned model
print('\nCloned model configuration:')
for key, value in clone_model.configuration.items():
    print(f'{key}: {value}')

# Option 1: Open the cloned model in the Dataloop UI to modify configuration
clone_model.open_in_web()

# Option 2: Modify configuration programmatically (e.g., number of epochs)
# Make sure the key 'num_epochs' (or similar, like 'epochs') exists in your model's configuration.
clone_model.configuration['num_epochs'] = 10 # Set desired number of epochs
clone_model.update() # Save the changes

clone_model = dl.models.get(model_id=clone_model.id)

# Display model configuration again to confirm changes
print('\nUpdated model configuration:')
for key, value in clone_model.configuration.items():
    print(f'{key}: {value}')

# Add subsets to the dataset
# This operation tags items with 'train', 'validation', or 'test'.
filters = dl.Filters(field='type', values='file') # Filter to select only 'file' type items (e.g., images)
dataset.split_ml_subsets(
    # Define the query to filter items (only files in this case)
    items_query=filters,
    # Split the dataset into train (60%), validation (20%), and test (20%) subsets
    # These will be used for training, validating, and testing the model respectively
    # Adjust percentages as needed for your dataset size and requirements.
    percentages={'train': 60, 'validation': 20, 'test': 20}
)

# Create filters to select items tagged as 'train' and 'validation'
# These tags are automatically added by the split_ml_subsets method.
train_filters = dl.Filters(field="metadata.system.tags.train", values=True)
validation_filters = dl.Filters(field="metadata.system.tags.validation", values=True)

# Add the train subset to the model, specifying the dataset and the filter
# This tells the model which items to use for training
clone_model.add_subset(subset_name="train", subset_filter=train_filters) # dataset_id is taken from clone_model.dataset_id if not specified

# Add the validation subset to the model
# This tells the model which items to use for validation during training
clone_model.add_subset(subset_name="validation", subset_filter=validation_filters)

# Update the model to save the subset configurations
# The system_metadata=True parameter ensures system metadata (like subset definitions) is updated.
clone_model.update(system_metadata=True)

# Start the model training process
# This will use the items defined in the 'train' and 'validation' subsets.
# This returns an execution object that can be used to track training progress.
execution = clone_model.train()

# Print confirmation message
print('Model training started successfully')
print(f"Training initiated. Execution ID: {execution.id}")

# Wait for the training execution to complete
execution.wait() 
print('Model training finished successfully')

# Open the model page in the Dataloop platform to monitor training
clone_model.open_in_web()

# Create filters to select items tagged as 'test'
# This filter will identify all items in the dataset that have been tagged for testing.
test_filters = dl.Filters(field='metadata.system.tags.test', values=True)

# Start the model evaluation process on the test subset
# - dataset_id: Specifies which dataset to use for evaluation (taken from clone_model.dataset_id by default if associated).
# - filters: Applies the test filters to only evaluate on test items.
execution = clone_model.evaluate(dataset_id=dataset.id, # Explicitly providing dataset.id
                                 filters=test_filters)

# Print confirmation message
print('Model evaluation started successfully')

# Wait for the evaluation execution to complete
execution.wait()
print('Model evaluation finished successfully')

clone_model.open_in_web()