Tinh chỉnh SAM 2 trên bộ dữ liệu tùy chỉnh: Hướng dẫn

SAM2 là gì?

SAM2 là mô hình nền tảng do Meta phát triển cho phân đoạn thị giác theo lời nhắc trong ảnh và video. Khác với phiên bản tiền nhiệm SAM chủ yếu tập trung vào ảnh tĩnh, SAM2 được thiết kế để xử lý cả những phức tạp của phân đoạn video.

SAM2 - Nhiệm vụ, Mô hình và Dữ liệu (Nguồn: Ravi et al., 2024)

Mô hình sử dụng kiến trúc transformer với bộ nhớ streaming, cho phép xử lý video theo thời gian thực. Quá trình huấn luyện của SAM 2 bao gồm một tập dữ liệu lớn và đa dạng, nổi bật là SA-V mới, với hơn 600.000 chú giải masklet trên 51.000 video.

Cỗ máy dữ liệu của nó, cho phép thu thập dữ liệu tương tác và cải thiện mô hình, mang lại cho mô hình khả năng phân đoạn mọi thứ có thể. Cơ chế này giúp SAM 2 học liên tục và thích ứng, nhờ đó xử lý dữ liệu mới và thách thức hiệu quả hơn. Tuy nhiên, với các tác vụ chuyên biệt theo miền hoặc đối tượng hiếm, tinh chỉnh là cần thiết để đạt hiệu suất tối ưu.

Vì sao nên tinh chỉnh SAM2?

Trong bối cảnh SAM 2, tinh chỉnh là quá trình tiếp tục huấn luyện mô hình SAM 2 đã được tiền huấn luyện trên một bộ dữ liệu cụ thể để nâng cao hiệu suất cho một tác vụ hay miền nhất định. Dù SAM 2 là công cụ mạnh mẽ được huấn luyện trên tập dữ liệu rộng và đa dạng, tính chất đa dụng của nó không phải lúc nào cũng cho kết quả tối ưu với các tác vụ chuyên sâu hoặc hiếm.

Ví dụ, nếu bạn đang thực hiện một dự án ảnh y khoa yêu cầu nhận diện các loại khối u cụ thể, hiệu suất của mô hình có thể chưa đạt yêu cầu do được huấn luyện theo hướng tổng quát.

Quy trình tinh chỉnh

Tinh chỉnh SAM 2 khắc phục hạn chế này bằng cách cho phép bạn điều chỉnh mô hình theo bộ dữ liệu của riêng mình. Quá trình này cải thiện độ chính xác của mô hình và giúp nó hiệu quả hơn cho trường hợp sử dụng đặc thù của bạn.

Dưới đây là các lợi ích chính của việc tinh chỉnh SAM 2:

1. Cải thiện độ chính xác: Bằng cách tinh chỉnh mô hình trên bộ dữ liệu cụ thể, bạn có thể nâng cao đáng kể độ chính xác, đảm bảo hiệu suất tốt hơn trong ứng dụng mục tiêu.
2. Phân đoạn chuyên biệt: Tinh chỉnh giúp mô hình thành thạo trong việc phân đoạn các loại đối tượng, phong cách thị giác hoặc môi trường cụ thể liên quan đến dự án của bạn, mang lại kết quả phù hợp mà mô hình đa dụng có thể không đạt được.
3. Hiệu quả: Tinh chỉnh thường hiệu quả hơn so với huấn luyện từ đầu. Thông thường cần ít dữ liệu và thời gian hơn, là giải pháp thực tiễn để nhanh chóng thích ứng mô hình cho các tác vụ mới hoặc ngách.

Bắt đầu tinh chỉnh SAM 2: Điều kiện tiên quyết

Để bắt đầu tinh chỉnh SAM 2, bạn cần chuẩn bị các điều kiện sau:

1. Truy cập mô hình và mã nguồn SAM 2: Có quyền truy cập mô hình SAM 2 và mã nguồn của nó. Bạn có thể tải mô hình SAM 2 đã tiền huấn luyện từ kho GitHub của Meta.
2. Bộ dữ liệu phù hợp: Bạn cần một bộ dữ liệu có chứa mask phân đoạn chuẩn. Trong hướng dẫn này, chúng tôi sẽ sử dụng bộ dữ liệu Chest CT Segmentation, bạn có thể tải về và chuẩn bị cho huấn luyện.
3. Tài nguyên tính toán: Tinh chỉnh SAM 2 đòi hỏi phần cứng có năng lực tính toán đủ mạnh. Rất khuyến nghị sử dụng GPU để đảm bảo quá trình hiệu quả và khả thi, đặc biệt khi làm việc với bộ dữ liệu lớn hoặc mô hình phức tạp. Trong ví dụ này, sử dụng GPU A100 trên Google Colab.

Phần mềm và yêu cầu khác:

• Python 3.11 hoặc cao hơn
• PyTorch
• OpenCV: Cài đặt bằng !pip install opencv-python

Chuẩn bị bộ dữ liệu để tinh chỉnh SAM 2

Chất lượng bộ dữ liệu là yếu tố then chốt để tinh chỉnh mô hình SAM 2. Ảnh hoặc video được gán nhãn chất lượng cao với mask phân đoạn chính xác là điều thiết yếu để đạt hiệu suất tối ưu. Chú giải chuẩn xác giúp mô hình học đúng đặc trưng, dẫn đến độ chính xác phân đoạn tốt hơn và khả năng bền vững trong ứng dụng thực tế.

1. Thu thập dữ liệu

Bước đầu tiên là thu thập bộ dữ liệu, nền tảng của quá trình huấn luyện. Chúng tôi lấy dữ liệu từ Kaggle, một nền tảng đáng tin cậy cung cấp nhiều bộ dữ liệu đa dạng. Sử dụng Kaggle API, chúng tôi tải dữ liệu theo đúng định dạng, đảm bảo ảnh và mask phân đoạn tương ứng sẵn sàng cho các bước xử lý tiếp theo.

2. Trích xuất và làm sạch dữ liệu

Sau khi tải về các bộ dữ liệu, chúng tôi thực hiện các bước sau:

• Giải nén và làm sạch: Giải nén dữ liệu từ các tệp zip đã tải và xóa tệp không cần thiết để tiết kiệm dung lượng đĩa.
• Trích xuất ID: Trích xuất định danh duy nhất (ID) cho ảnh và mask để đảm bảo ánh xạ chính xác giữa chúng trong quá trình huấn luyện.
• Loại bỏ tệp không cần thiết: Xóa các tệp nhiễu hoặc không liên quan, chẳng hạn một số ảnh có vấn đề đã biết, để đảm bảo tính toàn vẹn của bộ dữ liệu.

3. Chuyển đổi sang định dạng sử dụng được

Vì mô hình SAM2 yêu cầu đầu vào ở các định dạng cụ thể, chúng tôi chuyển đổi dữ liệu như sau:

• DICOM sang NumPy: Ảnh DICOM được đọc và lưu dưới dạng mảng NumPy, sau đó được đổi kích thước về 512x512 pixel tiêu chuẩn.
• NRRD sang NumPy cho mask: Tương tự, các tệp NRRD chứa mask cho phổi, tim và khí quản được xử lý và lưu thành mảng NumPy. Các mask này sau đó được chỉnh hình dạng cho khớp với ảnh tương ứng.
• Chuyển sang JPG/PNG: Để hiển thị và tương thích tốt hơn, các mảng NumPy được chuyển sang định dạng JPG/PNG. Bước này bao gồm chuẩn hóa giá trị cường độ ảnh và đảm bảo mask có định hướng chính xác.

4. Lưu và tổ chức dữ liệu

Ảnh và mask sau xử lý được sắp xếp vào các thư mục riêng để dễ truy cập trong quá trình tinh chỉnh. Ngoài ra, đường dẫn tới các ảnh và mask này được ghi vào tệp CSV (train.csv) để hỗ trợ tải dữ liệu trong khi huấn luyện.

5. Trực quan hóa và thẩm định

Bước cuối cùng là thẩm định bộ dữ liệu để đảm bảo tính chính xác:

• Trực quan hóa: Chúng tôi hiển thị cặp ảnh–mask bằng cách chồng mask lên ảnh. Điều này giúp kiểm tra căn chỉnh và độ chính xác của mask.
• Kiểm tra: Bằng cách kiểm tra một vài mẫu, chúng tôi xác nhận bộ dữ liệu đã được chuẩn bị đúng và sẵn sàng cho tinh chỉnh.

Tại đây là một notebook nhanh giúp bạn đi qua phần mã tạo bộ dữ liệu. Bạn có thể đi theo lộ trình tạo dữ liệu này hoặc dùng trực tiếp bất kỳ bộ dữ liệu nào có sẵn trực tuyến với cùng định dạng như đã nêu trong phần điều kiện tiên quyết.

Tinh chỉnh SAM2

Segment Anything Model 2 gồm nhiều thành phần, nhưng mấu chốt để tinh chỉnh nhanh là chỉ huấn luyện các thành phần nhẹ như bộ giải mã mask và bộ mã hóa lời nhắc, thay vì toàn bộ mô hình. Các bước tinh chỉnh mô hình như sau:

Bước 1: Cài đặt SAM-2

Để bắt đầu tinh chỉnh, chúng ta cần cài thư viện SAM-2, rất quan trọng cho Segment Anything Model (SAM2). Mô hình này được thiết kế để xử lý hiệu quả nhiều tác vụ phân đoạn. Cài đặt gồm việc clone kho SAM-2 từ GitHub và cài đặt các phụ thuộc cần thiết.

!git clone https://github.com/facebookresearch/segment-anything-2
%cd /content/segment-anything-2
!pip install -q -e .Được hỗ trợ bởi

Đoạn mã này đảm bảo thư viện SAM2 được cài đúng và sẵn sàng dùng trong quy trình tinh chỉnh.

Bước 2: Tải bộ dữ liệu

Sau khi cài đặt thư viện SAM-2, bước tiếp theo là lấy bộ dữ liệu dùng cho tinh chỉnh. Chúng tôi dùng một bộ dữ liệu trên Kaggle, cụ thể là bộ phân đoạn CT ngực chứa ảnh và mask của phổi, tim và khí quản.

Bộ dữ liệu bao gồm:

• images.zip: Ảnh định dạng RGB
• masks.zip: Mask phân đoạn định dạng RGB
• train.csv: Tệp CSV với tên ảnh

Ảnh từ bộ dữ liệu chụp CT

Trong bài viết này, chúng ta chỉ dùng ảnh và mask của phổi để phân đoạn. Kaggle API cho phép tải bộ dữ liệu trực tiếp về môi trường làm việc. Chúng ta bắt đầu bằng cách tải lên tệp kaggle.json từ Kaggle để dễ dàng truy cập bất kỳ bộ dữ liệu nào.

Để lấy kaggle.json, vào tab Settings trong hồ sơ người dùng của bạn và chọn Create New Token. Thao tác này sẽ kích hoạt tải xuống Kaggle.json, một tệp chứa thông tin xác thực API của bạn.

# get dataset from Kaggle
from google.colab import files
files.upload()  # This will prompt you to upload the kaggle.json file

!mkdir -p ~/.kaggle
!mv kaggle.json ~/.kaggle/
!chmod 600 ~/.kaggle/kaggle.json
!kaggle datasets download -d polomarco/chest-ct-segmentationĐược hỗ trợ bởi

Giải nén dữ liệu:

!unzip chest-ct-segmentation.zip -d chest-ct-segmentationĐược hỗ trợ bởi

Khi bộ dữ liệu đã sẵn sàng, hãy bắt đầu quá trình tinh chỉnh. Như tôi đã đề cập, chìa khóa ở đây là chỉ tinh chỉnh các thành phần nhẹ của SAM2, như bộ giải mã mask và bộ mã hóa lời nhắc, thay vì toàn bộ mô hình. Cách tiếp cận này hiệu quả hơn và cần ít tài nguyên hơn.

Bước 3: Tải checkpoint của SAM-2

Để tinh chỉnh, chúng ta cần bắt đầu với trọng số mô hình SAM2 đã tiền huấn luyện. Những trọng số này, gọi là "checkpoint", là điểm xuất phát cho huấn luyện tiếp. Các checkpoint đã được huấn luyện trên phạm vi hình ảnh rộng, và bằng cách tinh chỉnh trên bộ dữ liệu cụ thể của chúng ta, có thể đạt hiệu suất tốt hơn cho tác vụ mục tiêu.

!wget -O sam2_hiera_tiny.pt "https://dl.fbaipublicfiles.com/segment_anything_2/072824/sam2_hiera_tiny.pt"
!wget -O sam2_hiera_small.pt "https://dl.fbaipublicfiles.com/segment_anything_2/072824/sam2_hiera_small.pt"
!wget -O sam2_hiera_base_plus.pt "https://dl.fbaipublicfiles.com/segment_anything_2/072824/sam2_hiera_base_plus.pt"
!wget -O sam2_hiera_large.pt "https://dl.fbaipublicfiles.com/segment_anything_2/072824/sam2_hiera_large.pt"Được hỗ trợ bởi

Ở bước này, chúng ta tải về nhiều checkpoint SAM-2 tương ứng với các kích thước mô hình khác nhau (ví dụ: tiny, small, base_plus, large). Việc chọn checkpoint có thể điều chỉnh dựa trên tài nguyên tính toán sẵn có và tác vụ cụ thể.

Bước 4: Chuẩn bị dữ liệu

Sau khi tải bộ dữ liệu, bước tiếp theo là chuẩn bị nó cho huấn luyện. Điều này bao gồm chia bộ dữ liệu thành tập huấn luyện và kiểm thử, đồng thời tạo các cấu trúc dữ liệu có thể đưa vào mô hình SAM 2 trong quá trình tinh chỉnh.

%cd /content/segment-anything-2

import os
import pandas as pd
import cv2
import torch
import torch.nn.utils
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.colors as mcolors
from sklearn.model_selection import train_test_split
from sam2.build_sam import build_sam2
from sam2.sam2_image_predictor import SAM2ImagePredictor

# Path to the chest-ct-segmentation dataset folder
data_dir = "/content/segment-anything-2/chest-ct-segmentation"
images_dir = os.path.join(data_dir, "images/images")
masks_dir = os.path.join(data_dir, "masks/masks")

# Load the train.csv file
train_df = pd.read_csv(os.path.join(data_dir, "train.csv"))

# Split the data into two halves: one for training and one for testing
train_df, test_df = train_test_split(train_df, test_size=0.2, random_state=42)

# Prepare the training data list
train_data = []
for index, row in train_df.iterrows():
   image_name = row['ImageId']
   mask_name = row['MaskId']

   # Append image and corresponding mask paths
   train_data.append({
       "image": os.path.join(images_dir, image_name),
       "annotation": os.path.join(masks_dir, mask_name)
   })

# Prepare the testing data list (if needed for inference or evaluation later)
test_data = []
for index, row in test_df.iterrows():
   image_name = row['ImageId']
   mask_name = row['MaskId']

   # Append image and corresponding mask paths
   test_data.append({
       "image": os.path.join(images_dir, image_name),
       "annotation": os.path.join(masks_dir, mask_name)
   })
Được hỗ trợ bởi

Chúng tôi chia bộ dữ liệu thành tập huấn luyện (80%) và tập kiểm thử (20%) để có thể đánh giá hiệu suất mô hình sau huấn luyện. Dữ liệu huấn luyện sẽ dùng để tinh chỉnh SAM 2, còn dữ liệu kiểm thử dùng cho suy luận và đánh giá.

Sau khi chia bộ dữ liệu thành tập huấn luyện và kiểm thử, bước tiếp theo là tạo mask nhị phân, chọn các điểm khóa trong các mask này, và trực quan hóa các thành phần để đảm bảo dữ liệu đã được xử lý đúng.

1. Đọc và đổi kích thước ảnh: Quy trình bắt đầu bằng việc chọn ngẫu nhiên một ảnh và mask tương ứng từ bộ dữ liệu. Ảnh được chuyển từ BGR sang RGB, là định dạng màu tiêu chuẩn cho hầu hết các mô hình học sâu. Chú giải tương ứng (mask) được đọc ở chế độ thang xám. Sau đó, cả ảnh và mask được đổi kích thước về tối đa 1024 pixel, giữ nguyên tỷ lệ để phù hợp yêu cầu đầu vào của mô hình và giảm tải tính toán.

def read_batch(data, visualize_data=False):
   # Select a random entry
   ent = data[np.random.randint(len(data))]

   # Get full paths
   Img = cv2.imread(ent["image"])[..., ::-1]  # Convert BGR to RGB
   ann_map = cv2.imread(ent["annotation"], cv2.IMREAD_GRAYSCALE)  # Read annotation as grayscale

   if Img is None or ann_map is None:
       print(f"Error: Could not read image or mask from path {ent['image']} or {ent['annotation']}")
       return None, None, None, 0

   # Resize image and mask
   r = np.min([1024 / Img.shape[1], 1024 / Img.shape[0]])  # Scaling factor
   Img = cv2.resize(Img, (int(Img.shape[1] * r), int(Img.shape[0] * r)))
   ann_map = cv2.resize(ann_map, (int(ann_map.shape[1] * r), int(ann_map.shape[0] * r)), interpolation=cv2.INTER_NEAREST)
Được hỗ trợ bởi

2. Nhị phân hóa mask phân đoạn: Mask chú giải đa lớp (có thể có nhiều lớp đối tượng gán bằng các giá trị điểm ảnh khác nhau) được chuyển thành mask nhị phân. Mask này làm nổi bật tất cả vùng quan tâm trong ảnh, đơn giản hóa tác vụ phân đoạn thành bài toán phân loại nhị phân: đối tượng vs. nền. Mask nhị phân sau đó được xói mòn (erode) bằng kernel 5x5.

Xói mòn giúp giảm nhẹ kích thước mask, tránh hiệu ứng biên khi chọn điểm. Nhờ đó, các điểm được chọn nằm sâu trong phần nội bộ đối tượng thay vì gần rìa, nơi có thể nhiễu hoặc mơ hồ.

Các điểm khóa được chọn từ bên trong mask đã xói mòn. Các điểm này đóng vai trò là lời nhắc trong quá trình tinh chỉnh, định hướng mô hình tập trung vào đâu. Điểm được chọn ngẫu nhiên từ phần nội bộ của đối tượng để đảm bảo tính đại diện và không bị ảnh hưởng bởi vùng biên nhiễu.

   ### Continuation of read_batch() ###

   # Initialize a single binary mask
   binary_mask = np.zeros_like(ann_map, dtype=np.uint8)
   points = []

   # Get binary masks and combine them into a single mask
   inds = np.unique(ann_map)[1:]  # Skip the background (index 0)
   for ind in inds:
       mask = (ann_map == ind).astype(np.uint8)  # Create binary mask for each unique index
       binary_mask = np.maximum(binary_mask, mask)  # Combine with the existing binary mask

   # Erode the combined binary mask to avoid boundary points
   eroded_mask = cv2.erode(binary_mask, np.ones((5, 5), np.uint8), iterations=1)

   # Get all coordinates inside the eroded mask and choose a random point
   coords = np.argwhere(eroded_mask > 0)
   if len(coords) > 0:
       for _ in inds:  # Select as many points as there are unique labels
           yx = np.array(coords[np.random.randint(len(coords))])
           points.append([yx[1], yx[0]])

   points = np.array(points)
Được hỗ trợ bởi

3. Trực quan hóa: Bước này rất quan trọng để xác minh các bước tiền xử lý dữ liệu đã được thực hiện đúng. Bằng cách quan sát trực quan các điểm trên mask đã nhị phân hóa, bạn có thể đảm bảo mô hình sẽ nhận được đầu vào phù hợp trong quá trình huấn luyện. Cuối cùng, mask nhị phân được thay đổi hình dạng và định dạng đúng (kích thước phù hợp với đầu vào mô hình), và các điểm cũng được đổi hình dạng để sử dụng tiếp trong huấn luyện. Hàm trả về ảnh đã xử lý, mask nhị phân, các điểm đã chọn và số lượng mask tìm thấy.

    ### Continuation of read_batch() ###

    if visualize_data:
        # Plotting the images and points
        plt.figure(figsize=(15, 5))

        # Original Image
        plt.subplot(1, 3, 1)
        plt.title('Original Image')
        plt.imshow(img)
        plt.axis('off')

        # Segmentation Mask (binary_mask)
        plt.subplot(1, 3, 2)
        plt.title('Binarized Mask')
        plt.imshow(binary_mask, cmap='gray')
        plt.axis('off')

        # Mask with Points in Different Colors
        plt.subplot(1, 3, 3)
        plt.title('Binarized Mask with Points')
        plt.imshow(binary_mask, cmap='gray')

        # Plot points in different colors
        colors = list(mcolors.TABLEAU_COLORS.values())
        for i, point in enumerate(points):
            plt.scatter(point[0], point[1], c=colors[i % len(colors)], s=100, label=f'Point {i+1}')  # Corrected to plot y, x order

        # plt.legend()
        plt.axis('off')

        plt.tight_layout()
        plt.show()

    binary_mask = np.expand_dims(binary_mask, axis=-1)  # Now shape is (1024, 1024, 1)
    binary_mask = binary_mask.transpose((2, 0, 1))
    points = np.expand_dims(points, axis=1)

    # Return the image, binarized mask, points, and number of masks
    return img, binary_mask, points, len(inds)

# Visualize the data
Img1, masks1, points1, num_masks = read_batch(train_data, visualize_data=True)
Được hỗ trợ bởi

Đoạn mã trên trả về hình ảnh sau, gồm ảnh gốc từ bộ dữ liệu cùng mask nhị phân và mask nhị phân với các điểm.

Ảnh gốc, mask nhị phân và mask nhị phân với điểm cho bộ dữ liệu.

Bước 5: Tinh chỉnh mô hình SAM2

Tinh chỉnh mô hình SAM2 bao gồm nhiều bước như tải mô hình, thiết lập bộ tối ưu và lịch trình, và cập nhật trọng số mô hình lặp đi lặp lại dựa trên dữ liệu huấn luyện.

Tải checkpoint mô hình:

sam2_checkpoint = "sam2_hiera_small.pt"  # @param ["sam2_hiera_tiny.pt", "sam2_hiera_small.pt", "sam2_hiera_base_plus.pt", "sam2_hiera_large.pt"]
model_cfg = "sam2_hiera_s.yaml" # @param ["sam2_hiera_t.yaml", "sam2_hiera_s.yaml", "sam2_hiera_b+.yaml", "sam2_hiera_l.yaml"]

sam2_model = build_sam2(model_cfg, sam2_checkpoint, device="cuda")
predictor = SAM2ImagePredictor(sam2_model)
Được hỗ trợ bởi

Chúng ta bắt đầu bằng việc xây dựng mô hình SAM2 từ các checkpoint đã tiền huấn luyện. Mô hình sau đó được bọc trong lớp predictor, giúp đơn giản hóa quá trình thiết lập ảnh, mã hóa lời nhắc và giải mã mask.

Cấu hình siêu tham số

Chúng tôi cấu hình một số siêu tham số để đảm bảo mô hình học hiệu quả, như tốc độ học, weight decay và số bước tích lũy gradient. Các siêu tham số này điều khiển quá trình học, bao gồm tốc độ cập nhật trọng số và cách tránh overfitting. Bạn có thể thử nghiệm với chúng.

# Train mask decoder.
predictor.model.sam_mask_decoder.train(True)

# Train prompt encoder.
predictor.model.sam_prompt_encoder.train(True)

# Configure optimizer.
optimizer=torch.optim.AdamW(params=predictor.model.parameters(),lr=0.0001,weight_decay=1e-4) #1e-5, weight_decay = 4e-5

# Mix precision.
scaler = torch.cuda.amp.GradScaler()

# No. of steps to train the model.
NO_OF_STEPS = 3000 # @param 

# Fine-tuned model name.
FINE_TUNED_MODEL_NAME = "fine_tuned_sam2"
Được hỗ trợ bởi

Bộ tối ưu chịu trách nhiệm cập nhật trọng số mô hình, trong khi scheduler điều chỉnh tốc độ học trong quá trình huấn luyện để cải thiện hội tụ. Bằng cách tinh chỉnh các tham số này, chúng ta có thể đạt độ chính xác phân đoạn tốt hơn.

Bắt đầu huấn luyện

Quá trình tinh chỉnh thực tế là lặp, trong mỗi bước, một lô ảnh và mask chỉ cho phổi được đưa qua mô hình, và hàm mất mát được tính toán để cập nhật trọng số mô hình.

# Initialize scheduler
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=500, gamma=0.2) # 500 , 250, gamma = 0.1
accumulation_steps = 4  # Number of steps to accumulate gradients before updating

for step in range(1, NO_OF_STEPS + 1):
   with torch.cuda.amp.autocast():
       image, mask, input_point, num_masks = read_batch(train_data, visualize_data=False)
       if image is None or mask is None or num_masks == 0:
           continue

       input_label = np.ones((num_masks, 1))
       if not isinstance(input_point, np.ndarray) or not isinstance(input_label, np.ndarray):
           continue

       if input_point.size == 0 or input_label.size == 0:
           continue

       predictor.set_image(image)
       mask_input, unnorm_coords, labels, unnorm_box = predictor._prep_prompts(input_point, input_label, box=None, mask_logits=None, normalize_coords=True)
       if unnorm_coords is None or labels is None or unnorm_coords.shape[0] == 0 or labels.shape[0] == 0:
           continue

       sparse_embeddings, dense_embeddings = predictor.model.sam_prompt_encoder(
           points=(unnorm_coords, labels), boxes=None, masks=None,
       )

       batched_mode = unnorm_coords.shape[0] > 1
       high_res_features = [feat_level[-1].unsqueeze(0) for feat_level in predictor._features["high_res_feats"]]
       low_res_masks, prd_scores, _, _ = predictor.model.sam_mask_decoder(
           image_embeddings=predictor._features["image_embed"][-1].unsqueeze(0),
           image_pe=predictor.model.sam_prompt_encoder.get_dense_pe(),
           sparse_prompt_embeddings=sparse_embeddings,
           dense_prompt_embeddings=dense_embeddings,
           multimask_output=True,
           repeat_image=batched_mode,
           high_res_features=high_res_features,
       )
       prd_masks = predictor._transforms.postprocess_masks(low_res_masks, predictor._orig_hw[-1])

       gt_mask = torch.tensor(mask.astype(np.float32)).cuda()
       prd_mask = torch.sigmoid(prd_masks[:, 0])
       seg_loss = (-gt_mask * torch.log(prd_mask + 0.000001) - (1 - gt_mask) * torch.log((1 - prd_mask) + 0.00001)).mean()

       inter = (gt_mask * (prd_mask > 0.5)).sum(1).sum(1)
       iou = inter / (gt_mask.sum(1).sum(1) + (prd_mask > 0.5).sum(1).sum(1) - inter)
       score_loss = torch.abs(prd_scores[:, 0] - iou).mean()
       loss = seg_loss + score_loss * 0.05

       # Apply gradient accumulation
       loss = loss / accumulation_steps
       scaler.scale(loss).backward()

       # Clip gradients
       torch.nn.utils.clip_grad_norm_(predictor.model.parameters(), max_norm=1.0)

       if step % accumulation_steps == 0:
           scaler.step(optimizer)
           scaler.update()
           predictor.model.zero_grad()

       # Update scheduler
       scheduler.step()

       if step % 500 == 0:
           FINE_TUNED_MODEL = FINE_TUNED_MODEL_NAME + "_" + str(step) + ".torch"
           torch.save(predictor.model.state_dict(), FINE_TUNED_MODEL)

       if step == 1:
           mean_iou = 0

       mean_iou = mean_iou * 0.99 + 0.01 * np.mean(iou.cpu().detach().numpy())

       if step % 100 == 0:
           print("Step " + str(step) + ":\t", "Accuracy (IoU) = ", mean_iou)
Được hỗ trợ bởi

Trong mỗi vòng lặp, mô hình xử lý một lô ảnh, tính toán mask phân đoạn và so sánh với ground truth để tính loss. Loss này được dùng để điều chỉnh trọng số mô hình, dần cải thiện hiệu suất. Sau khoảng 3000 epoch huấn luyện, chúng tôi đạt độ chính xác (IoU - Intersection over Union) khoảng 72%.

Bước 6: Suy luận với mô hình đã tinh chỉnh

Mô hình sau đó có thể dùng cho suy luận, dự đoán mask phân đoạn trên ảnh mới chưa thấy. Bắt đầu với các hàm trợ giúp read_images và get_points để lấy ảnh suy luận và mask của nó cùng các điểm khóa.

def read_image(image_path, mask_path):  # read and resize image and mask
   img = cv2.imread(image_path)[..., ::-1]  # Convert BGR to RGB
   mask = cv2.imread(mask_path, 0)
   r = np.min([1024 / img.shape[1], 1024 / img.shape[0]])
   img = cv2.resize(img, (int(img.shape[1] * r), int(img.shape[0] * r)))
   mask = cv2.resize(mask, (int(mask.shape[1] * r), int(mask.shape[0] * r)), interpolation=cv2.INTER_NEAREST)
   return img, mask

def get_points(mask, num_points):  # Sample points inside the input mask
   points = []
   coords = np.argwhere(mask > 0)
   for i in range(num_points):
       yx = np.array(coords[np.random.randint(len(coords))])
       points.append([[yx[1], yx[0]]])
   return np.array(points)Được hỗ trợ bởi

Sau đó tải các ảnh mẫu bạn muốn suy luận, cùng với trọng số vừa tinh chỉnh, và thực hiện suy luận trong ngữ cảnh torch.no_grad().

# Randomly select a test image from the test_data
selected_entry = random.choice(test_data)
image_path = selected_entry['image']
mask_path = selected_entry['annotation']

# Load the selected image and mask
image, mask = read_image(image_path, mask_path)

# Generate random points for the input
num_samples = 30  # Number of points per segment to sample
input_points = get_points(mask, num_samples)

# Load the fine-tuned model
FINE_TUNED_MODEL_WEIGHTS = "fine_tuned_sam2_1000.torch"
sam2_model = build_sam2(model_cfg, sam2_checkpoint, device="cuda")

# Build net and load weights
predictor = SAM2ImagePredictor(sam2_model)
predictor.model.load_state_dict(torch.load(FINE_TUNED_MODEL_WEIGHTS))

# Perform inference and predict masks
with torch.no_grad():
   predictor.set_image(image)
   masks, scores, logits = predictor.predict(
       point_coords=input_points,
       point_labels=np.ones([input_points.shape[0], 1])
   )

# Process the predicted masks and sort by scores
np_masks = np.array(masks[:, 0])
np_scores = scores[:, 0]
sorted_masks = np_masks[np.argsort(np_scores)][::-1]

# Initialize segmentation map and occupancy mask
seg_map = np.zeros_like(sorted_masks[0], dtype=np.uint8)
occupancy_mask = np.zeros_like(sorted_masks[0], dtype=bool)

# Combine masks to create the final segmentation map
for i in range(sorted_masks.shape[0]):
   mask = sorted_masks[i]
   if (mask * occupancy_mask).sum() / mask.sum() > 0.15:
       continue

   mask_bool = mask.astype(bool)
   mask_bool[occupancy_mask] = False  # Set overlapping areas to False in the mask
   seg_map[mask_bool] = i + 1  # Use boolean mask to index seg_map
   occupancy_mask[mask_bool] = True  # Update occupancy_mask

# Visualization: Show the original image, mask, and final segmentation side by side
plt.figure(figsize=(18, 6))

plt.subplot(1, 3, 1)
plt.title('Test Image')
plt.imshow(image)
plt.axis('off')

plt.subplot(1, 3, 2)
plt.title('Original Mask')
plt.imshow(mask, cmap='gray')
plt.axis('off')

plt.subplot(1, 3, 3)
plt.title('Final Segmentation')
plt.imshow(seg_map, cmap='jet')
plt.axis('off')

plt.tight_layout()
plt.show()Được hỗ trợ bởi

Ở bước này, chúng ta dùng mô hình đã tinh chỉnh để tạo mask phân đoạn cho ảnh kiểm thử. Các mask dự đoán sau đó được hiển thị cùng ảnh gốc và mask ground truth để đánh giá hiệu suất mô hình.

Ảnh phân đoạn cuối cùng trên dữ liệu kiểm thử

Kết luận

Tinh chỉnh SAM2 là cách thực tế để tăng cường khả năng cho các tác vụ cụ thể. Dù bạn làm về ảnh y khoa, xe tự hành hay chỉnh sửa video, tinh chỉnh cho phép bạn dùng SAM2 theo nhu cầu riêng. Bằng cách làm theo hướng dẫn này, bạn có thể điều chỉnh SAM2 cho dự án của mình và đạt kết quả phân đoạn tiên tiến.

Với các trường hợp nâng cao hơn, hãy cân nhắc tinh chỉnh thêm các thành phần khác của SAM2 như bộ mã hóa ảnh. Dù cần nhiều tài nguyên hơn, cách này mang lại độ linh hoạt và cải thiện hiệu năng lớn hơn.