import cv2 import numpy as np import re from skimage import transform as trans def parse_lst_line(line): vec = line.strip().split("\t") assert len(vec)>=3 aligned = int(vec[0]) image_path = vec[2] label = int(re.sub("\D","",vec[1])) bbox = None landmark = None #print(vec) if len(vec)>3: bbox = np.zeros( (4,), dtype=np.int32) for i in range(3,7): bbox[i-3] = int(vec[i]) landmark = None if len(vec)>7: _l = [] for i in range(7,17): _l.append(float(vec[i])) landmark = np.array(_l).reshape( (2,5) ).T #print(aligned) return image_path, label, bbox, landmark, aligned def read_image(img_path, **kwargs): mode = kwargs.get('mode', 'rgb') layout = kwargs.get('layout', 'HWC') if mode=='gray': img = cv2.imread(img_path, cv2.CV_LOAD_IMAGE_GRAYSCALE) else: img = cv2.imread(img_path, cv2.CV_LOAD_IMAGE_COLOR) if mode=='rgb': #print('to rgb') img = img[...,::-1] if layout=='CHW': img = np.transpose(img, (2,0,1)) return img def preprocess(img, bbox=None, landmark=None, **kwargs): #print('-' * 50) #print('landmark: ',landmark) if isinstance(img, str): img = read_image(img, **kwargs) M = None image_size = [] str_image_size = kwargs.get('image_size', '') if len(str_image_size)>0: image_size = [int(x) for x in str_image_size.split(',')] if len(image_size)==1: image_size = [image_size[0], image_size[0]] assert len(image_size)==2 assert image_size[0]==112 assert image_size[0]==112 or image_size[1]==96 if landmark is not None: assert len(image_size)==2 src = np.array([ [30.2946, 51.6963], [65.5318, 51.5014], [48.0252, 71.7366], [33.5493, 92.3655], [62.7299, 92.2041] ], dtype=np.float32 ) if image_size[1]==112: src[:,0] += 8.0 dst = landmark.astype(np.float32) tform = trans.SimilarityTransform() tform.estimate(dst, src) M = tform.params[0:2,:] #print(M) # = cv2.estimateRigidTransform( dst.reshape(1,5,2), src.reshape(1,5,2), False) if M is None: if bbox is None: #use center crop det = np.zeros(4, dtype=np.int32) det[0] = int(img.shape[1]*0.0625) det[1] = int(img.shape[0]*0.0625) det[2] = img.shape[1] - det[0] det[3] = img.shape[0] - det[1] else: det = bbox margin = kwargs.get('margin', 44) bb = np.zeros(4, dtype=np.int32) bb[0] = np.maximum(det[0]-margin/2, 0) bb[1] = np.maximum(det[1]-margin/2, 0) bb[2] = np.minimum(det[2]+margin/2, img.shape[1]) bb[3] = np.minimum(det[3]+margin/2, img.shape[0]) ret = img[bb[1]:bb[3],bb[0]:bb[2],:] if len(image_size)>0: ret = cv2.resize(ret, (image_size[1], image_size[0])) return ret else: #do align using landmark assert len(image_size)==2 #src = src[0:3,:] #dst = dst[0:3,:] #print(src.shape, dst.shape) #print('src: ',src) #print('dst: ',dst) #print(M) #print('-' * 50) warped = cv2.warpAffine(img,M,(image_size[1],image_size[0]), borderValue = 0.0) #tform3 = trans.ProjectiveTransform() #tform3.estimate(src, dst) #warped = trans.warp(img, tform3, output_shape=_shape) return warped