【无标题】

from ultralytics import SAM
import cv2
import matplotlib.pyplot as plt

Load the model

model = SAM(‘mobile_sam.pt’)

Predict a segment based on a point prompt

result = model.predict(‘ultralytics/assets/zidane.jpg’, points=[900, 370], labels=[1])

Get the predicted mask

mask = result.pred[0]

Load the input image

input_image = cv2.imread(‘ultralytics/assets/zidane.jpg’)

Overlay the mask on the input image

overlay = cv2.addWeighted(input_image, 0.7, mask, 0.3, 0)

Visualize the result

plt.figure(figsize=(10, 5))
plt.subplot(131)
plt.title(“Input Image”)
plt.imshow(cv2.cvtColor(input_image, cv2.COLOR_BGR2RGB))
plt.subplot(132)
plt.title(“Predicted Mask”)
plt.imshow(mask, cmap=‘gray’)
plt.subplot(133)
plt.title(“Overlay”)
plt.imshow(cv2.cvtColor(overlay, cv2.COLOR_BGR2RGB))
plt.show()

Save the overlay image to a specified folder

output_folder = “output/”
output_path = output_folder + “overlay_image.jpg”
cv2.imwrite(output_path, overlay)

print(f"Overlay image saved to {output_path}")在这里插入代码片