File size: 3,108 Bytes
d575b00 |
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 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 |
import matplotlib
import numpy as np
import torch
from transformers import SamModel, SamProcessor, pipeline
checkpoint = "google/owlvit-base-patch16"
detector = pipeline(model=checkpoint, task="zero-shot-object-detection", device="cuda")
sam_model = SamModel.from_pretrained("facebook/sam-vit-base").cuda()
sam_processor = SamProcessor.from_pretrained("facebook/sam-vit-base")
# image_dims = (256, 256)
image_dims = (224, 224)
def get_bounding_boxes(img, prompt="the black robotic gripper"):
predictions = detector(img, candidate_labels=[prompt], threshold=0.01)
return predictions
def show_box(box, ax, meta, color):
x0, y0 = box["xmin"], box["ymin"]
w, h = box["xmax"] - box["xmin"], box["ymax"] - box["ymin"]
ax.add_patch(
matplotlib.patches.FancyBboxPatch((x0, y0), w, h, edgecolor=color, facecolor=(0, 0, 0, 0), lw=2, label="hehe")
)
ax.text(x0, y0 + 10, "{:.3f}".format(meta["score"]), color="white")
def get_median(mask, p):
row_sum = np.sum(mask, axis=1)
cumulative_sum = np.cumsum(row_sum)
if p >= 1.0:
p = 1
total_sum = np.sum(row_sum)
threshold = p * total_sum
return np.argmax(cumulative_sum >= threshold)
def get_gripper_mask(img, pred):
box = [
round(pred["box"]["xmin"], 2),
round(pred["box"]["ymin"], 2),
round(pred["box"]["xmax"], 2),
round(pred["box"]["ymax"], 2),
]
inputs = sam_processor(img, input_boxes=[[[box]]], return_tensors="pt")
for k in inputs.keys():
inputs[k] = inputs[k].cuda()
with torch.no_grad():
outputs = sam_model(**inputs)
mask = (
sam_processor.image_processor.post_process_masks(
outputs.pred_masks, inputs["original_sizes"], inputs["reshaped_input_sizes"]
)[0][0][0]
.cpu()
.numpy()
)
return mask
def sq(w, h):
return np.concatenate(
[
(np.arange(w * h).reshape(h, w) % w)[:, :, None],
(np.arange(w * h).reshape(h, w) // w)[:, :, None],
],
axis=-1,
)
def mask_to_pos_weighted(mask):
pos = sq(*image_dims)
weight = pos[:, :, 0] + pos[:, :, 1]
weight = weight * weight
x = np.sum(mask * pos[:, :, 0] * weight) / np.sum(mask * weight)
y = get_median(mask * weight, 0.95)
return x, y
def mask_to_pos_naive(mask):
pos = sq(*image_dims)
weight = pos[:, :, 0] + pos[:, :, 1]
min_pos = np.argmax((weight * mask).flatten())
return min_pos % image_dims[0] - (image_dims[0] / 16), min_pos // image_dims[0] - (image_dims[0] / 24)
def get_gripper_pos_raw(img):
# img = Image.fromarray(img.numpy())
predictions = get_bounding_boxes(img)
if len(predictions) > 0:
mask = get_gripper_mask(img, predictions[0])
pos = mask_to_pos_naive(mask)
else:
mask = np.zeros(image_dims)
pos = (-1, -1)
predictions = [None]
# return (int(pos[0]), int(pos[1])), mask, predictions[0]
return (int(pos[0]*224/image_dims[0]), int(pos[1]*224/image_dims[1])), mask, predictions[0]
if __name__ == "__main__":
pass
|