File size: 1,563 Bytes
e7334c8 |
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 |
import base64, os, io, random, time
from PIL import Image
import numpy as np
def b64_mask_encode(mask_np_arr, tmp_dir = '/tmp/miro/mask_encoding/'):
'''
turn a binary mask in numpy into a base64 string
'''
mask_im = Image.fromarray(np.array(mask_np_arr).astype(np.uint8)*255)
mask_im = mask_im.convert(mode = '1') # convert to 1bit image
if not os.path.isdir(tmp_dir):
print(f'b64_mask_encode: making tmp dir for mask encoding...')
os.makedirs(tmp_dir)
timestr = time.strftime("%Y%m%d-%H%M%S")
hash_str = random.getrandbits(128)
tmp_fname = tmp_dir + f'{timestr}_{hash_str}_mask.png'
mask_im.save(tmp_fname)
return base64.b64encode(open(tmp_fname, 'rb').read())
def b64_mask_decode(b64_string):
'''
decode a base64 string back to a binary mask numpy array
'''
im_bytes = base64.b64decode(b64_string)
im_decode = Image.open(io.BytesIO(im_bytes))
return np.array(im_decode)
def get_true_mask(mask_arr, im_w_h:tuple, x0, y0, x1, y1):
'''
decode the mask of CM output to get a mask that's the same size as source im
'''
if x0 > im_w_h[0] or x1 > im_w_h[0] or y0 > im_w_h[1] or y1 > im_w_h[1]:
raise ValueError(f'get_true_mask: Xs and Ys exceeded im_w_h bound: {im_w_h}')
if mask_arr.shape != (y1 - y0, x1 - x0):
raise ValueError(f'get_true_mask: Bounding Box h: {y1-y0} w: {x1-x0} does not match mask shape: {mask_arr.shape}')
w, h = im_w_h
mask = np.zeros((h,w), dtype = np.uint8)
mask[y0:y1, x0:x1] = mask_arr
return mask
|