SAM_DataAnnotation/src/plotting.py at main · BioIntelligence-Lab/SAM_DataAnnotation · GitHub

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
import matplotlib.pyplot as plt
import matplotlib.patches as patches
import matplotlib.cm as cm
from .bbox import *

def find_and_visualize_bboxes(ax, img, mask):
  # Plot image
  ax.imshow(img, cmap='gray')
  ax.imshow(mask, alpha=0.5)
  bboxes = find_bboxes(mask)
  # Add bboxes as patches
  for bbox in bboxes:
    x1, y1, x2, y2 = bbox
    xwidth = x2 - x1
    ywidth = y2 - y1
    p = patches.Rectangle(
      (x1, y1),
      xwidth,
      ywidth,
      fc = 'none',
      ec = 'red'
    )
    ax.add_patch(p)
  ax.axis('off')

def overlay_masks(*masks, cmap=cm.viridis):
  masks = np.array(masks)
  targets = np.arange(masks.shape[0]) + 1
  overlay = np.sum(masks * targets[:,None,None], axis=0)
  color_overlay = cmap(rescale_intensities(overlay).astype(np.uint8))
  return (rescale_intensities(color_overlay)*overlay[:,:,None]).astype(np.uint8)

def visualize(ax, img, mask, points = None, bbox = None, bg_points = None):
  # Plot image
  ax.imshow(img, cmap='gray')
  ax.imshow(mask, alpha=0.6)
  if isinstance(points, np.ndarray):
    ax.plot(points[:,0], points[:,1], 'r.')
  if isinstance(bg_points, np.ndarray):
    ax.plot(bg_points[:,0], bg_points[:,1], 'rx')
  if isinstance(bbox, np.ndarray):
    # Add bboxes as patches
    for b in bbox:
      x1, y1, x2, y2 = b
      xwidth = x2 - x1
      ywidth = y2 - y1
      p = patches.Rectangle(
        (x1, y1),
        xwidth,
        ywidth,
        fc = 'none',
        ec = 'red'
      )
      ax.add_patch(p)
  ax.axis('off')