Spaces:
Runtime error
Runtime error
File size: 1,126 Bytes
19d592d 146111c 19d592d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 |
import argparse
import cv2
import numpy as np
from segment_anything import SamPredictor, sam_model_registry
# Argument parser
parser = argparse.ArgumentParser()
parser.add_argument("-i", "--image", required=True, help="Path to the image")
args = parser.parse_args()
# Set hyperparameters
sam_checkpoint = "./models/sam_vit_h_4b8939.pth"
model_type = "vit_h"
device = "cpu"
# Load model
sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)
sam.to(device=device)
predictor = SamPredictor(sam)
# Preprocessing the image
image = cv2.imread(args.image)
predictor.set_image(image)
# SAM Encoder for embedding
embedding = predictor.get_image_embedding()
np.save("models/embedding.npy", embedding)
# SAM Decoder for segmentation
input_point = np.array([[1300, 950]])
input_label = np.array([1])
mask, score, logit = predictor.predict(
point_coords=input_point,
point_labels=input_label,
multimask_output=False,
)
# Save output
h, w = mask.shape[-2:]
mask = mask.reshape(h, w, 1)
## Mask has a 255 or 0 value
mask = (mask * 255).astype(np.uint8)
## Save mask image
cv2.imwrite("mask.png", mask[:, :])
|