dentist / masks_prepare.py
theparit's picture
Upload 146 files
aecbb99
# -*- coding: utf-8 -*-
"""
@author: serdarhelli
"""
import os
import sys
import numpy as np
from PIL import Image
from zipfile import ZipFile
from natsort import natsorted
script_dir=os.path.abspath(os.path.dirname(sys.argv[0]))
default_path=script_dir+'/Original_Masks/'
def convert_one_channel(img):
#some images have 3 channels , although they are grayscale image
if len(img.shape)>2:
img=img[:,:,0]
return img
else:
return img
def pre_masks(resize_shape=(512,512),path=default_path):
ZipFile(path+"/Orig_Masks.zip").extractall(path+'/Masks/')
path=path+'/Masks/'
dirs=natsorted(os.listdir(path))
masks=img=Image.open(path+dirs[0])
masks=(masks.resize((resize_shape),Image.LANCZOS))
masks=convert_one_channel(np.asarray(masks))
for i in range (1,len(dirs)):
img=Image.open(path+dirs[i])
img=img.resize((resize_shape),Image.LANCZOS)
img=convert_one_channel(np.asarray(img))
masks=np.concatenate((masks,img))
masks=np.reshape(masks,(len(dirs),resize_shape[0],resize_shape[1],1))
return masks
default_path=script_dir+'/Custom_Masks/'
#CustomMasks 512x512
def pre_splitted_masks(path=default_path):
ZipFile(path+"/splitted_masks.zip").extractall(path+'/Masks/')
path=path+'/Masks/'
dirs=natsorted(os.listdir(path))
masks=img=Image.open(path+dirs[0])
masks=convert_one_channel(np.asarray(masks))
for i in range (1,len(dirs)):
img=Image.open(path+dirs[i])
img=convert_one_channel(np.asarray(img))
masks=np.concatenate((masks,img))
masks=np.reshape(masks,(len(dirs),512,512,1))
return masks