Spaces:
Sleeping
Sleeping
File size: 2,873 Bytes
dccc960 |
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 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 |
import numpy as np
import cv2
import matplotlib.pyplot as plt
from regionspot.modeling.regionspot import build_regionspot_model
from regionspot import RegionSpot_Predictor
# Function to show masks on an image
def show_mask(mask, ax, random_color=False):
if random_color:
color = np.concatenate([np.random.random(3), np.array([0.6])], axis=0)
else:
color = np.array([30/255, 144/255, 255/255, 0.6])
h, w = mask.shape[-2:]
mask_image = mask.reshape(h, w, 1) * color.reshape(1, 1, -1)
ax.imshow(mask_image)
# Function to show points on an image
def show_points(coords, labels, ax, marker_size=375):
pos_points = coords[labels == 1]
neg_points = coords[labels == 0]
ax.scatter(pos_points[:, 0], pos_points[:, 1], color='green', marker='*', s=marker_size, edgecolor='white', linewidth=1.25)
ax.scatter(neg_points[:, 0], neg_points[:, 1], color='red', marker='*', s=marker_size, edgecolor='white', linewidth=1.25)
# Function to show bounding boxes on an image
def show_box(box, ax):
x0, y0 = box[0], box[1]
w, h = box[2] - x0, box[3] - y0
ax.add_patch(plt.Rectangle((x0, y0), w, h, edgecolor='green', facecolor='none', lw=2))
# Read image and set up model
image = cv2.imread('assets/image.jpg')
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) # Convert image to RGB format
# Multiple boxes
box_prompt = np.array([[64, 926, 804, 1978], [1237, 490, 1615, 771.], [1510, 64, 1670, 167]])
ckpt_path = '/path/to/model_weights.pth'
clip_type = 'CLIP_400M_Large_336'
clip_input_size = 336
custom_vocabulary = ["Smoothie bowl", "Banana", "Strawberry", "Chia seeds", "Shredded coconut", "Wooden spoons", "Grapefruit", "Goji berries", "Flaxseeds seeds"]
# Build and initialize the model
model, msg = build_regionspot_model(is_training=False, image_size=clip_input_size, clip_type=clip_type, pretrain_ckpt=ckpt_path, custom_vocabulary=custom_vocabulary)
# Create predictor and set image
predictor = RegionSpot_Predictor(model.cuda())
predictor.set_image(image, clip_input_size=clip_input_size)
# Prediction based on box prompt
masks, mask_iou_score, class_score, class_index = predictor.predict(
point_coords=None,
point_labels=None,
box=box_prompt,
multimask_output=False,
)
# Extract class name and display image with masks and box
fig, ax = plt.subplots(figsize=(10, 10))
ax.imshow(image)
for idx in range(len(box_prompt)):
show_mask(masks[idx], ax)
show_box(box_prompt[idx], ax) # Assuming box_prompt contains all your boxes
# You might want to modify the text display for multiple classes as well
class_name = custom_vocabulary[int(class_index[idx])]
ax.text(box_prompt[idx][0], box_prompt[idx][1] - 10, class_name, color='white', fontsize=14, bbox=dict(facecolor='green', edgecolor='green', alpha=0.6))
ax.axis('off')
plt.show()
fig.savefig('result.png')
plt.close(fig)
|