Spaces:
Runtime error
Runtime error
File size: 2,461 Bytes
c7f097c |
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 |
import os
import cv2
import numpy as np
from pathlib import Path
import argparse
def get_bbox(msk):
rows = np.any(msk, axis=1)
cols = np.any(msk, axis=0)
rmin, rmax = np.where(rows)[0][[0,-1]]
cmin, cmax = np.where(cols)[0][[0,-1]]
return rmin, rmax, cmin, cmax
def process_img(img, msk, bbox=None):
if bbox is None:
bbox = get_bbox(msk > 100)
cx = (bbox[3] + bbox[2])//2
cy = (bbox[1] + bbox[0])//2
w = img.shape[1]
h = img.shape[0]
height = int(1.138*(bbox[1] - bbox[0]))
hh = height//2
# crop
dw = min(cx, w-cx, hh)
if cy-hh < 0:
img = cv2.copyMakeBorder(img,hh-cy,0,0,0,cv2.BORDER_CONSTANT,value=[0,0,0])
msk = cv2.copyMakeBorder(msk,hh-cy,0,0,0,cv2.BORDER_CONSTANT,value=0)
cy = hh
if cy+hh > h:
img = cv2.copyMakeBorder(img,0,cy+hh-h,0,0,cv2.BORDER_CONSTANT,value=[0,0,0])
msk = cv2.copyMakeBorder(msk,0,cy+hh-h,0,0,cv2.BORDER_CONSTANT,value=0)
img = img[cy-hh:(cy+hh),cx-dw:cx+dw,:]
msk = msk[cy-hh:(cy+hh),cx-dw:cx+dw]
dw = img.shape[0] - img.shape[1]
if dw != 0:
img = cv2.copyMakeBorder(img,0,0,dw//2,dw//2,cv2.BORDER_CONSTANT,value=[0,0,0])
msk = cv2.copyMakeBorder(msk,0,0,dw//2,dw//2,cv2.BORDER_CONSTANT,value=0)
img = cv2.resize(img, (512, 512))
msk = cv2.resize(msk, (512, 512))
kernel = np.ones((3,3),np.uint8)
msk = cv2.erode((255*(msk > 100)).astype(np.uint8), kernel, iterations = 1)
return img, msk
def main():
'''
given foreground mask, this script crops and resizes an input image and mask for processing.
'''
parser = argparse.ArgumentParser()
parser.add_argument('-i', '--input_image', type=str, help='if the image has alpha channel, it will be used as mask')
parser.add_argument('-m', '--input_mask', type=str)
parser.add_argument('-o', '--out_path', type=str, default='./sample_images')
args = parser.parse_args()
img = cv2.imread(args.input_image, cv2.IMREAD_UNCHANGED)
if img.shape[2] == 4:
msk = img[:,:,3:]
img = img[:,:,:3]
else:
msk = cv2.imread(args.input_mask, cv2.IMREAD_GRAYSCALE)
img_new, msk_new = process_img(img, msk)
img_name = Path(args.input_image).stem
cv2.imwrite(os.path.join(args.out_path, img_name + '.png'), img_new)
cv2.imwrite(os.path.join(args.out_path, img_name + '_mask.png'), msk_new)
if __name__ == "__main__":
main() |