Более четкая сегментация SAM (Segment Anything) ⇐ Python

[Anonymous]

Мне нужно сегментировать объекты, присутствующие на изображении, следующим образом:
Изображение 1
Изображение 2
Я решил использовать AI SAM (Segment Anything) от Meta, чтобы обрезать эти объекты.
Мой код выглядит следующим образом:

Код: Выделить всё

import cv2
import numpy as np
import sys, os
from pathlib import Path
import torch
import supervision as sv
from segment_anything import sam_model_registry, SamAutomaticMaskGenerator, SamPredictor

# Get command line arguments
num = sys.argv[2]
img_path = sys.argv[1]
folder_path = sys.argv[3]

# Load the image from the given path
img = cv2.imread(img_path)

# Function to preprocess the image for cutting
def preprocess_image_cut(img):
gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
return gray

# Function to obtain a clean image
def get_clean_image(img, filter_level=1):
gray = preprocess_image_cut(img)
blurred_image = cv2.GaussianBlur(gray, (3, 3), 0)
_, binary_image = cv2.threshold(blurred_image, 210, 255, cv2.THRESH_BINARY)

kernel = np.ones((1, 1), np.uint8)
result_image = cv2.morphologyEx(binary_image, cv2.MORPH_OPEN, kernel)
result_image = cv2.medianBlur(result_image, filter_level)

output_image = cv2.bitwise_or(gray, result_image)
return output_image

# Function to set up the pipeline
def setup_pipeline():
HOME = 'C:/'
CHECKPOINT_PATH = os.path.join(HOME, 'weights', 'sam_vit_h_4b8939.pth')
DEVICE = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
MODEL_TYPE = "vit_h"

sam = sam_model_registry[MODEL_TYPE](checkpoint=CHECKPOINT_PATH).to(device=DEVICE)
mask_generator = SamAutomaticMaskGenerator(sam)
return mask_generator

# Function to run the SAM model
def run_sam(mask_generator, clean_image):
image_rgb = cv2.cvtColor(clean_image, cv2.COLOR_BGR2RGB)
image_bgr = cv2.cvtColor(image_rgb, cv2.COLOR_RGB2BGR)
sam_result = mask_generator.generate(image_rgb)
mask_annotator = sv.MaskAnnotator(color_lookup=sv.ColorLookup.INDEX)
detections = sv.Detections.from_sam(sam_result=sam_result)

annotated_image = mask_annotator.annotate(scene=image_bgr.copy(), detections=detections)

masks = [mask['segmentation'] for mask in sorted(sam_result, key=lambda x: x['area', reverse=True])]
return masks

# Set up the mask generator
mask_generator = setup_pipeline()

if __name__ == "__main__":
if img_path.endswith(".png"):
save_path = folder_path
os.makedirs(save_path, exist_ok=True)
if os.path.isdir(f'{folder_path}/PROCESSED'):
os.makedirs(f'{folder_path}/PROCESSED', exist_ok=True)
img = cv2.imread(img_path)
gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
clean_image = get_clean_image(img, 3)

print("Starting mask generation")
masks = run_sam(mask_generator, clean_image)

for i, mask in enumerate(masks):
image_rgb = cv2.cvtColor(clean_image, cv2.COLOR_BGR2RGB)
image_bgra = cv2.cvtColor(image_rgb, cv2.COLOR_RGB2BGRA)
extracted_region = np.zeros_like(image_bgra)
extracted_region[mask] = image_bgra[mask]
x, y, w, h = cv2.boundingRect(mask.astype(np.uint8))
extracted_region = extracted_region[y:y+h, x:x+w]
extracted_region[:, :, 3] = (mask[y:y+h, x:x+w] > 0) * 255
save_path2 = f'{folder_path}/PROCESSED/img_{x}_{y}_{w}_{h}.png'
cv2.imwrite(save_path2, extracted_region)

print("Finished")

Но у меня возникают трудности с некоторыми моментами, например с тем, что я не хочу, чтобы Сэм сегментировал числа или линии, связывающие объекты с числами.
Может ли кто-нибудь предложить какие-нибудь идеи по этой теме? Также принимаются другие методы сегментации изображения.

Подробнее здесь: https://stackoverflow.com/questions/788 ... t-anything