update sam
Browse files- sam_utils.py +0 -13
sam_utils.py
CHANGED
@@ -23,28 +23,15 @@ def sam_out_nosave(predictor, input_image, *bbox_sliders):
|
|
23 |
start_time = time.time()
|
24 |
predictor.set_image(image)
|
25 |
|
26 |
-
h, w, _ = image.shape
|
27 |
-
input_point = np.array([[h//2, w//2]])
|
28 |
-
input_label = np.array([1])
|
29 |
-
|
30 |
-
masks, scores, logits = predictor.predict(
|
31 |
-
point_coords=input_point,
|
32 |
-
point_labels=input_label,
|
33 |
-
multimask_output=True,
|
34 |
-
)
|
35 |
-
|
36 |
masks_bbox, scores_bbox, logits_bbox = predictor.predict(
|
37 |
box=bbox,
|
38 |
multimask_output=True
|
39 |
)
|
40 |
|
41 |
print(f"SAM Time: {time.time() - start_time:.3f}s")
|
42 |
-
opt_idx = np.argmax(scores)
|
43 |
-
mask = masks[opt_idx]
|
44 |
out_image = np.zeros((image.shape[0], image.shape[1], 4), dtype=np.uint8)
|
45 |
out_image[:, :, :3] = image
|
46 |
out_image_bbox = out_image.copy()
|
47 |
-
out_image[:, :, 3] = mask.astype(np.uint8) * 255
|
48 |
out_image_bbox[:, :, 3] = masks_bbox[-1].astype(np.uint8) * 255 # np.argmax(scores_bbox)
|
49 |
torch.cuda.empty_cache()
|
50 |
return Image.fromarray(out_image_bbox, mode='RGBA')
|
|
|
23 |
start_time = time.time()
|
24 |
predictor.set_image(image)
|
25 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
26 |
masks_bbox, scores_bbox, logits_bbox = predictor.predict(
|
27 |
box=bbox,
|
28 |
multimask_output=True
|
29 |
)
|
30 |
|
31 |
print(f"SAM Time: {time.time() - start_time:.3f}s")
|
|
|
|
|
32 |
out_image = np.zeros((image.shape[0], image.shape[1], 4), dtype=np.uint8)
|
33 |
out_image[:, :, :3] = image
|
34 |
out_image_bbox = out_image.copy()
|
|
|
35 |
out_image_bbox[:, :, 3] = masks_bbox[-1].astype(np.uint8) * 255 # np.argmax(scores_bbox)
|
36 |
torch.cuda.empty_cache()
|
37 |
return Image.fromarray(out_image_bbox, mode='RGBA')
|