automatic_generation_sam2.py

# Source: https://github.com/facebookresearch/sam2/blob/main/notebooks/automatic_mask_generator_example.ipynb

# Install Dependencies
# pip install torch torchvision opencv-python matplotlib Pillow
# pip install git+https://github.com/facebookresearch/sam2.git

# Compatible with Python 3.10+
# Create a virtual environment with a compatible Python version: python3.12 -m venv myenv
# conda create -n myenv python=3.10
# conda activate myenv
# conda install pytorch torchvision torchaudio cudatoolkit=11.3 -c pytorch

'''
With Python 10 on Server Computer:
pip install torch==1.12.1+cu113 torchvision==0.13.1+cu113 torchaudio==0.12.1 --extra-index-url https://download.pytorch.org/whl/cu113
pip install git+https://github.com/facebookresearch/sam2.git
'''

# Download Model Checkpoint
# Linux, MacOS:
# wget -P checkpoints https://dl.fbaipublicfiles.com/segment_anything_2/092824/sam2.1_hiera_large.pt
# wget -P images https://raw.githubusercontent.com/facebookresearch/sam2/main/notebooks/images/cars.jpg
# Windows:
# mkdir checkpoints && curl -o checkpoints/sam2.1_hiera_large.pt https://dl.fbaipublicfiles.com/segment_anything_2/092824/sam2.1_hiera_large.pt
# mkdir images && curl -o images/cars.jpg https://raw.githubusercontent.com/facebookresearch/sam2/main/notebooks/images/cars.jpg

# Directory Structure
# project_directory/
#     images/
#         cars.jpg
#     checkpoints/
#         sam2.1_hiera_large.pt

import os
import numpy as np
import torch
import matplotlib.pyplot as plt
from PIL import Image
from sam2.build_sam import build_sam2
from sam2.automatic_mask_generator import SAM2AutomaticMaskGenerator
import cv2

# Set up device
print(torch.__version__)
print(torch.cuda.is_available())
print(torch.cuda.device_count())
num_gpus = 0

if torch.cuda.is_available():
    device = torch.device("cuda")
    num_gpus = torch.cuda.device_count()
elif torch.backends.mps.is_available():
    device = torch.device("mps")
else:
    device = torch.device("cpu")
print(f"using device: {device}")

# Set random seed for reproducibility
np.random.seed(3)

def show_anns(anns, borders=True):
    if len(anns) == 0:
        return
    sorted_anns = sorted(anns, key=(lambda x: x['area']), reverse=True)
    ax = plt.gca()
    ax.set_autoscale_on(False)

    img = np.ones((sorted_anns[0]['segmentation'].shape[0], sorted_anns[0]['segmentation'].shape[1], 4))
    img[:, :, 3] = 0
    for ann in sorted_anns:
        m = ann['segmentation']
        color_mask = np.concatenate([np.random.random(3), [0.5]])
        img[m] = color_mask
        if borders:
            contours, _ = cv2.findContours(m.astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
            contours = [cv2.approxPolyDP(contour, epsilon=0.01, closed=True) for contour in contours]
            cv2.drawContours(img, contours, -1, (0, 0, 1, 0.4), thickness=1)

    ax.imshow(img)

# Load the model
sam2_checkpoint = 'checkpoints/sam2.1_hiera_large.pt'
model_cfg = "configs/sam2.1/sam2.1_hiera_l.yaml"

sam2 = build_sam2(model_cfg, sam2_checkpoint, device=device, apply_postprocessing=False)
#mask_generator = SAM2AutomaticMaskGenerator(sam2)

mask_generator = SAM2AutomaticMaskGenerator(
    model=sam2,
    points_per_side=64,
    points_per_batch=128,
    pred_iou_thresh=0.7,
    stability_score_thresh=0.92,
    stability_score_offset=0.7,
    crop_n_layers=1,
    box_nms_thresh=0.7,
    crop_n_points_downscale_factor=2,
    min_mask_region_area=25.0,
    use_m2m=True,
)


def get_image_dimensions(image_path):
    image = cv2.imread(image_path)
    if image is None:
        raise FileNotFoundError(f"Image file not found: {image_path}")
    return image.shape[1], image.shape[0]

def extract_bounding_boxes(yolo_label_path, image_width, image_height):
    bounding_boxes = []
    
    try:
        with open(yolo_label_path, 'r') as file:
            lines = file.readlines()
            for line in lines:
                values = line.strip().split()
                class_id = int(values[0])
                x_center = float(values[1])
                y_center = float(values[2])
                width = float(values[3])
                height = float(values[4])
                
                # Convert normalized values to absolute pixel values
                x_min = (x_center - width / 2) * image_width
                y_min = (y_center - height / 2) * image_height
                x_max = (x_center + width / 2) * image_width
                y_max = (y_center + height / 2) * image_height
                
                bounding_boxes.append((x_min, y_min, x_max, y_max, class_id))
    
    except Exception as e:
        print(f"Error reading {yolo_label_path}: {e}")
    
    return bounding_boxes

# Segment fish in one given image
def segmentImage(filename, label_path, name):
    image = cv2.imread(filename)
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    image_width, image_height = get_image_dimensions(filename)
    masks = mask_generator.generate(image)

    plt.figure(figsize=(15, 15))
    plt.imshow(image)
    show_anns(masks)
    plt.axis('off')
    #plt.show()
    
    # Calculate area and create a list of tuples (area, mask)
    masks_with_area = [(mask['area'], mask) for mask in masks]

    # Sort masks by area (largest to smallest)
    sorted_masks = sorted(masks_with_area, key=lambda x: x[0], reverse=True)
    
    bounding_boxes = extract_bounding_boxes(label_path, image_width, image_height)
    counter = 0 # counter for naming

    # Loop through all bounding boxes in the image
    for bounding_box in bounding_boxes:
        masks_within_bbox = []
        x_min, y_min, x_max, y_max, class_id = bounding_box

        # Check if the bounding box is touching the edges of the image
        if (x_min <= 5 or y_min <= 5 or x_max >= image_width-5 or y_max >= image_height-5):
            print(f"Skipping bounding box {bounding_box} as it touches the image edge.")
            continue

        # Loop through each mask to check if it lies within the bounding box
        for idx, mask in enumerate(sorted_masks):  # Use sorted_masks from previous code
            segmentation = mask[1]['segmentation']  # Access the mask from the tuple
            coordinates = np.column_stack(np.where(segmentation > 0))  # Get the coordinates

            # Check if all coordinates are within the bounding box
            if np.all((coordinates[:, 0] >= bounding_box[1]) & (coordinates[:, 0] <= bounding_box[3]) &
                    (coordinates[:, 1] >= bounding_box[0]) & (coordinates[:, 1] <= bounding_box[2])):
                masks_within_bbox.append(mask)  # Add mask to the list if it is within the bounding box

        # Print the masks that are within the bounding box
        print(f'Within Bounding Box {counter}:')
        for idx, (area, mask) in enumerate(masks_within_bbox):
            print(f"Mask {idx} - Area: {area}")
            
        # Find the largest mask within the bounding box
        if masks_within_bbox:
            largest_mask = max(masks_within_bbox, key=lambda x: x[0])  # Get the mask with the largest area
            largest_segmentation = largest_mask[1]['segmentation']
            
            # Create an image to visualize the largest mask with transparency
            mask_image = np.zeros((image.shape[0], image.shape[1], 4), dtype=np.uint8)  # Create a blank RGBA image
            mask_image[..., 3] = 0  # Set alpha channel to 0 (fully transparent)

            # Copy the mask pixels and set alpha to 255 (fully opaque)
            mask_image[largest_segmentation > 0, :3] = image[largest_segmentation > 0]  # RGB channels
            mask_image[largest_segmentation > 0, 3] = 255  # Alpha channel
            
            # Display the original image and the largest mask
            showing_box = image.copy()
            cv2.rectangle(showing_box, (int(x_min), int(y_min)), (int(x_max), int(y_max)), (255, 0, 0), 2)
            plt.figure(figsize=(15, 15))
            plt.subplot(1, 2, 1)
            plt.imshow(showing_box)
            plt.title('Original Image')
            plt.axis('off')

            plt.subplot(1, 2, 2)
            plt.imshow(mask_image)
            plt.title('Largest Mask Within Bounding Box')
            plt.axis('off')
            #plt.show()

            # Save largest mask image
            output_path = os.path.join('segmented', 'images', f'{name}seg{counter}.png')
            os.makedirs(os.path.dirname(output_path), exist_ok=True)
            save_image = Image.fromarray(mask_image).convert('RGBA')
            save_image.save(output_path)

            # Save largest mask label in YOLO format
            h, w = image.shape[:2]

            # Calculate YOLO format values
            x_center = (x_min + x_max) / 2 / w
            y_center = (y_min + y_max) / 2 / h
            bbox_width = (x_max - x_min) / w
            bbox_height = (y_max - y_min) / h

            # Save mask label
            output_label_path = os.path.join('segmented', 'labels', f'{name}seg{counter}.txt')
            os.makedirs(os.path.dirname(output_label_path), exist_ok=True)
            print(f'Saving image at {output_path} with label path at {output_label_path}')
            with open(output_label_path, 'w') as file:
                file.write(f"{class_id} {x_center} {y_center} {bbox_width} {bbox_height}\n")
            counter += 1
        else:
            print("No masks found within the bounding box.")
            continue

def process_all_files(images_folder, labels_folder):
    if not os.path.isdir(images_folder):
        raise FileNotFoundError(f"Images directory does not exist: {images_folder}")
    if not os.path.isdir(labels_folder):
        raise FileNotFoundError(f"Labels directory does not exist: {labels_folder}")
    counter = 0 # counter for naming
    for file_name in os.listdir(images_folder):
        print(f"Processing Image {counter}:")
        if file_name.lower().endswith(('.jpg', '.png')):
            image_path = os.path.join(images_folder, file_name)
            label_path = os.path.join(labels_folder, file_name.rsplit('.', 1)[0] + '.txt')

            if not os.path.isfile(label_path):
                print(f"Label file not found for {file_name}")
                continue
                
            segmentImage(image_path, label_path, f'image{counter}')
            counter += 1

process_all_files(os.path.join('processed_data', 'images'), os.path.join('processed_data', 'labels'))