chaoxu commited on
Commit
3055ceb
1 Parent(s): 14550cf

update sam

Browse files
Files changed (1) hide show
  1. 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')