mirror of
https://github.com/deepinsight/insightface.git
synced 2026-05-13 12:00:46 +00:00
546 lines
20 KiB
Python
546 lines
20 KiB
Python
# pylint: skip-file
|
|
import mxnet as mx
|
|
import numpy as np
|
|
import sys, os
|
|
import random
|
|
import math
|
|
import scipy.misc
|
|
import cv2
|
|
import logging
|
|
import sklearn
|
|
import datetime
|
|
import img_helper
|
|
from mxnet.io import DataIter
|
|
from mxnet import ndarray as nd
|
|
from mxnet import io
|
|
from mxnet import recordio
|
|
from PIL import Image
|
|
|
|
|
|
class FaceSegIter0(DataIter):
|
|
def __init__(self, batch_size,
|
|
path_imgrec = None,
|
|
data_name = "data",
|
|
label_name = "softmax_label"):
|
|
self.batch_size = batch_size
|
|
self.data_name = data_name
|
|
self.label_name = label_name
|
|
assert path_imgrec
|
|
logging.info('loading recordio %s...',
|
|
path_imgrec)
|
|
path_imgidx = path_imgrec[0:-4]+".idx"
|
|
self.imgrec = mx.recordio.MXIndexedRecordIO(path_imgidx, path_imgrec, 'r') # pylint: disable=redefined-variable-type
|
|
self.seq = list(self.imgrec.keys)
|
|
self.cur = 0
|
|
self.reset()
|
|
self.num_classes = 68
|
|
self.record_img_size = 384
|
|
self.input_img_size = self.record_img_size
|
|
self.data_shape = (3, self.input_img_size, self.input_img_size)
|
|
self.label_shape = (self.num_classes, 2)
|
|
self.provide_data = [(data_name, (batch_size,) + self.data_shape)]
|
|
self.provide_label = [(label_name, (batch_size,) + self.label_shape)]
|
|
|
|
def reset(self):
|
|
print('reset')
|
|
"""Resets the iterator to the beginning of the data."""
|
|
self.cur = 0
|
|
|
|
def next_sample(self):
|
|
"""Helper function for reading in next sample."""
|
|
if self.cur >= len(self.seq):
|
|
raise StopIteration
|
|
idx = self.seq[self.cur]
|
|
self.cur += 1
|
|
s = self.imgrec.read_idx(idx)
|
|
header, img = recordio.unpack(s)
|
|
img = mx.image.imdecode(img).asnumpy()
|
|
#label = np.zeros( (self.num_classes, self.record_img_size, self.record_img_size), dtype=np.uint8)
|
|
hlabel = np.array(header.label).reshape( (self.num_classes,2) )
|
|
return img, hlabel
|
|
|
|
def next(self):
|
|
"""Returns the next batch of data."""
|
|
#print('next')
|
|
batch_size = self.batch_size
|
|
batch_data = nd.empty((batch_size,)+self.data_shape)
|
|
batch_label = nd.empty((batch_size,)+self.label_shape)
|
|
i = 0
|
|
#self.cutoff = random.randint(800,1280)
|
|
try:
|
|
while i < batch_size:
|
|
#print('N', i)
|
|
data, label = self.next_sample()
|
|
data = nd.array(data)
|
|
data = nd.transpose(data, axes=(2, 0, 1))
|
|
label = nd.array(label)
|
|
#print(data.shape, label.shape)
|
|
batch_data[i][:] = data
|
|
batch_label[i][:] = label
|
|
i += 1
|
|
except StopIteration:
|
|
if not i:
|
|
raise StopIteration
|
|
|
|
return mx.io.DataBatch([batch_data], [batch_label], batch_size - i)
|
|
|
|
class FaceSegIter(DataIter):
|
|
def __init__(self, batch_size,
|
|
per_batch_size = 0,
|
|
path_imgrec = None,
|
|
aug_level = 0,
|
|
force_mirror = False,
|
|
use_coherent = 0,
|
|
args = None,
|
|
data_name = "data",
|
|
label_name = "softmax_label"):
|
|
self.aug_level = aug_level
|
|
self.force_mirror = force_mirror
|
|
self.use_coherent = use_coherent
|
|
self.batch_size = batch_size
|
|
self.per_batch_size = per_batch_size
|
|
self.data_name = data_name
|
|
self.label_name = label_name
|
|
assert path_imgrec
|
|
logging.info('loading recordio %s...',
|
|
path_imgrec)
|
|
path_imgidx = path_imgrec[0:-4]+".idx"
|
|
self.imgrec = mx.recordio.MXIndexedRecordIO(path_imgidx, path_imgrec, 'r') # pylint: disable=redefined-variable-type
|
|
self.seq = list(self.imgrec.keys)
|
|
print('train size', len(self.seq))
|
|
self.cur = 0
|
|
self.reset()
|
|
self.num_classes = args.num_classes
|
|
self.record_img_size = 384
|
|
self.input_img_size = args.input_img_size
|
|
self.data_shape = (3, self.input_img_size, self.input_img_size)
|
|
self.label_classes = self.num_classes
|
|
if aug_level>0:
|
|
self.output_label_size = args.output_label_size
|
|
self.label_shape = (self.label_classes, self.output_label_size, self.output_label_size)
|
|
else:
|
|
self.output_label_size = args.input_img_size
|
|
#self.label_shape = (self.num_classes, 2)
|
|
self.label_shape = (self.num_classes, self.output_label_size, self.output_label_size)
|
|
self.provide_data = [(data_name, (batch_size,) + self.data_shape)]
|
|
self.provide_label = [(label_name, (batch_size,) + self.label_shape)]
|
|
if self.use_coherent>0:
|
|
#self.provide_label += [("softmax_label2", (batch_size,)+self.label_shape)]
|
|
self.provide_label += [("coherent_label", (batch_size,6))]
|
|
self.img_num = 0
|
|
self.invalid_num = 0
|
|
self.mode = 1
|
|
self.vis = 0
|
|
self.stats = [0,0]
|
|
self.mirror_set = [
|
|
(22,23),
|
|
(21,24),
|
|
(20,25),
|
|
(19,26),
|
|
(18,27),
|
|
(40,43),
|
|
(39,44),
|
|
(38,45),
|
|
(37,46),
|
|
(42,47),
|
|
(41,48),
|
|
(33,35),
|
|
(32,36),
|
|
(51,53),
|
|
(50,54),
|
|
(62,64),
|
|
(61,65),
|
|
(49,55),
|
|
(49,55),
|
|
(68,66),
|
|
(60,56),
|
|
(59,57),
|
|
(1,17),
|
|
(2,16),
|
|
(3,15),
|
|
(4,14),
|
|
(5,13),
|
|
(6,12),
|
|
(7,11),
|
|
(8,10),
|
|
]
|
|
|
|
def get_data_shape(self):
|
|
return self.data_shape
|
|
|
|
#def get_label_shape(self):
|
|
# return self.label_shape
|
|
|
|
def get_shape_dict(self):
|
|
D = {}
|
|
for (k,v) in self.provide_data:
|
|
D[k] = v
|
|
for (k,v) in self.provide_label:
|
|
D[k] = v
|
|
return D
|
|
|
|
def get_label_names(self):
|
|
D = []
|
|
for (k,v) in self.provide_label:
|
|
D.append(k)
|
|
return D
|
|
|
|
def reset(self):
|
|
print('reset')
|
|
"""Resets the iterator to the beginning of the data."""
|
|
if self.aug_level>0:
|
|
random.shuffle(self.seq)
|
|
self.cur = 0
|
|
|
|
def next_sample(self):
|
|
"""Helper function for reading in next sample."""
|
|
if self.cur >= len(self.seq):
|
|
raise StopIteration
|
|
idx = self.seq[self.cur]
|
|
self.cur += 1
|
|
s = self.imgrec.read_idx(idx)
|
|
header, img = recordio.unpack(s)
|
|
img = mx.image.imdecode(img).asnumpy()
|
|
#label = np.zeros( (self.num_classes, self.record_img_size, self.record_img_size), dtype=np.uint8)
|
|
hlabel = np.array(header.label).reshape( (self.num_classes,2) )
|
|
annot = {}
|
|
ul = np.array( (50000,50000), dtype=np.int32)
|
|
br = np.array( (0,0), dtype=np.int32)
|
|
for i in xrange(hlabel.shape[0]):
|
|
#hlabel[i] = hlabel[i][::-1]
|
|
h = int(hlabel[i][0])
|
|
w = int(hlabel[i][1])
|
|
key = np.array((h,w))
|
|
#print(key.shape, ul.shape, br.shape)
|
|
ul = np.minimum(key, ul)
|
|
br = np.maximum(key, br)
|
|
#label[h][w] = i+1
|
|
#label[i][h][w] = 1.0
|
|
#print(h,w,i+1)
|
|
|
|
if self.mode==0:
|
|
ul_margin = np.array( (60,30) )
|
|
br_margin = np.array( (30,30) )
|
|
crop_ul = ul
|
|
crop_ul = ul-ul_margin
|
|
crop_ul = np.maximum(crop_ul, np.array( (0,0), dtype=np.int32) )
|
|
crop_br = br
|
|
crop_br = br+br_margin
|
|
crop_br = np.minimum(crop_br, np.array( (img.shape[0],img.shape[1]), dtype=np.int32 ) )
|
|
elif self.mode==1:
|
|
#mm = (self.record_img_size - 256)//2
|
|
#crop_ul = (mm, mm)
|
|
#crop_br = (self.record_img_size-mm, self.record_img_size-mm)
|
|
crop_ul = (0,0)
|
|
crop_br = (self.record_img_size, self.record_img_size)
|
|
annot['scale'] = 256
|
|
#annot['center'] = np.array( (self.record_img_size/2+10, self.record_img_size/2) )
|
|
else:
|
|
mm = (48, 64)
|
|
#mm = (64, 80)
|
|
crop_ul = mm
|
|
crop_br = (self.record_img_size-mm[0], self.record_img_size-mm[1])
|
|
crop_ul = np.array(crop_ul)
|
|
crop_br = np.array(crop_br)
|
|
img = img[crop_ul[0]:crop_br[0],crop_ul[1]:crop_br[1],:]
|
|
#print(img.shape, crop_ul, crop_br)
|
|
invalid = False
|
|
for i in xrange(hlabel.shape[0]):
|
|
if (hlabel[i]<crop_ul).any() or (hlabel[i] >= crop_br).any():
|
|
invalid = True
|
|
hlabel[i] -= crop_ul
|
|
|
|
#mm = np.amin(ul)
|
|
#mm2 = self.record_img_size - br
|
|
#mm2 = np.amin(mm2)
|
|
#mm = min(mm, mm2)
|
|
#print('mm',mm, ul, br)
|
|
#print('invalid', invalid)
|
|
if invalid:
|
|
self.invalid_num+=1
|
|
|
|
annot['invalid'] = invalid
|
|
#annot['scale'] = (self.record_img_size - mm*2)*1.1
|
|
#if self.mode==1:
|
|
# annot['scale'] = 256
|
|
#print(annot)
|
|
|
|
#center = ul+br
|
|
#center /= 2.0
|
|
#annot['center'] = center
|
|
|
|
#img = img[ul[0]:br[0],ul[1]:br[1],:]
|
|
#scale = br-ul
|
|
#scale = max(scale[0], scale[1])
|
|
#print(img.shape)
|
|
return img, hlabel, annot
|
|
|
|
def do_aug(self, data, label, annot):
|
|
if self.vis:
|
|
self.img_num+=1
|
|
#if self.img_num<=self.vis:
|
|
# filename = './vis/raw_%d.jpg' % (self.img_num)
|
|
# print('save', filename)
|
|
# draw = data.copy()
|
|
# for i in xrange(label.shape[0]):
|
|
# cv2.circle(draw, (label[i][1], label[i][0]), 1, (0, 0, 255), 2)
|
|
# scipy.misc.imsave(filename, draw)
|
|
|
|
rotate = 0
|
|
#scale = 1.0
|
|
if 'scale' in annot:
|
|
scale = annot['scale']
|
|
else:
|
|
scale = max(data.shape[0], data.shape[1])
|
|
if 'center' in annot:
|
|
center = annot['center']
|
|
else:
|
|
center = np.array( (data.shape[0]/2, data.shape[1]/2) )
|
|
max_retry = 3
|
|
if self.aug_level==0:
|
|
max_retry = 6
|
|
retry = 0
|
|
found = False
|
|
_scale = scale
|
|
while retry<max_retry:
|
|
retry+=1
|
|
succ = True
|
|
if self.aug_level>0:
|
|
rotate = np.random.randint(-40, 40)
|
|
#rotate2 = np.random.randint(-40, 40)
|
|
rotate2 = 0
|
|
scale_config = 0.2
|
|
#rotate = 0
|
|
#scale_config = 0.0
|
|
_scale = min(1+scale_config, max(1-scale_config, (np.random.randn() * scale_config) + 1))
|
|
_scale *= scale
|
|
_scale = int(_scale)
|
|
#translate = np.random.randint(-5, 5, size=(2,))
|
|
#center += translate
|
|
if self.mode==1:
|
|
cropped = img_helper.crop2(data, center, _scale, (self.input_img_size, self.input_img_size), rot=rotate)
|
|
if self.use_coherent==2:
|
|
cropped2 = img_helper.crop2(data, center, _scale, (self.input_img_size, self.input_img_size), rot=rotate2)
|
|
else:
|
|
cropped = img_helper.crop(data, center, _scale, (self.input_img_size, self.input_img_size), rot=rotate)
|
|
#print('cropped', cropped.shape)
|
|
label_out = np.zeros(self.label_shape, dtype=np.float32)
|
|
label2_out = np.zeros(self.label_shape, dtype=np.float32)
|
|
G = 0
|
|
#if self.use_coherent:
|
|
# G = 1
|
|
_g = G
|
|
if G==0:
|
|
_g = 1
|
|
#print('shape', label.shape, label_out.shape)
|
|
for i in xrange(label.shape[0]):
|
|
pt = label[i].copy()
|
|
pt = pt[::-1]
|
|
#print('before gaussian', label_out[i].shape, pt.shape)
|
|
_pt = pt.copy()
|
|
trans = img_helper.transform(_pt, center, _scale, (self.output_label_size, self.output_label_size), rot=rotate)
|
|
#print(trans.shape)
|
|
if not img_helper.gaussian(label_out[i], trans, _g):
|
|
succ = False
|
|
break
|
|
if self.use_coherent==2:
|
|
_pt = pt.copy()
|
|
trans2 = img_helper.transform(_pt, center, _scale, (self.output_label_size, self.output_label_size), rot=rotate2)
|
|
if not img_helper.gaussian(label2_out[i], trans2, _g):
|
|
succ = False
|
|
break
|
|
if not succ:
|
|
if self.aug_level==0:
|
|
_scale+=20
|
|
continue
|
|
|
|
|
|
if self.use_coherent==1:
|
|
cropped2 = np.copy(cropped)
|
|
for k in xrange(cropped2.shape[2]):
|
|
cropped2[:,:,k] = np.fliplr(cropped2[:,:,k])
|
|
label2_out = np.copy(label_out)
|
|
for k in xrange(label2_out.shape[0]):
|
|
label2_out[k,:,:] = np.fliplr(label2_out[k,:,:])
|
|
new_label2_out = np.copy(label2_out)
|
|
for item in self.mirror_set:
|
|
mir = (item[0]-1, item[1]-1)
|
|
new_label2_out[mir[1]] = label2_out[mir[0]]
|
|
new_label2_out[mir[0]] = label2_out[mir[1]]
|
|
label2_out = new_label2_out
|
|
elif self.use_coherent==2:
|
|
pass
|
|
elif ((self.aug_level>0 and np.random.rand() < 0.5) or self.force_mirror): #flip aug
|
|
for k in xrange(cropped.shape[2]):
|
|
cropped[:,:,k] = np.fliplr(cropped[:,:,k])
|
|
for k in xrange(label_out.shape[0]):
|
|
label_out[k,:,:] = np.fliplr(label_out[k,:,:])
|
|
new_label_out = np.copy(label_out)
|
|
for item in self.mirror_set:
|
|
mir = (item[0]-1, item[1]-1)
|
|
new_label_out[mir[1]] = label_out[mir[0]]
|
|
new_label_out[mir[0]] = label_out[mir[1]]
|
|
label_out = new_label_out
|
|
|
|
if G==0:
|
|
for k in xrange(label.shape[0]):
|
|
ind = np.unravel_index(np.argmax(label_out[k], axis=None), label_out[k].shape)
|
|
label_out[k,:,:] = 0.0
|
|
label_out[k,ind[0],ind[1]] = 1.0
|
|
if self.use_coherent:
|
|
ind = np.unravel_index(np.argmax(label2_out[k], axis=None), label2_out[k].shape)
|
|
label2_out[k,:,:] = 0.0
|
|
label2_out[k,ind[0],ind[1]] = 1.0
|
|
found = True
|
|
break
|
|
|
|
|
|
#self.stats[0]+=1
|
|
if not found:
|
|
#self.stats[1]+=1
|
|
#print('find aug error', retry)
|
|
#print(self.stats)
|
|
return None
|
|
|
|
if self.vis>0 and self.img_num<=self.vis:
|
|
print('crop', data.shape, center, _scale, rotate, cropped.shape)
|
|
filename = './vis/cropped_%d.jpg' % (self.img_num)
|
|
print('save', filename)
|
|
draw = cropped.copy()
|
|
alabel = label_out.copy()
|
|
for i in xrange(label.shape[0]):
|
|
a = cv2.resize(alabel[i], (self.input_img_size, self.input_img_size))
|
|
ind = np.unravel_index(np.argmax(a, axis=None), a.shape)
|
|
cv2.circle(draw, (ind[1], ind[0]), 1, (0, 0, 255), 2)
|
|
scipy.misc.imsave(filename, draw)
|
|
if not self.use_coherent:
|
|
return cropped, label_out
|
|
else:
|
|
rotate2 = 0
|
|
r = rotate - rotate2
|
|
#r = rotate2 - rotate
|
|
r = math.pi*r/180
|
|
cos_r = math.cos(r)
|
|
sin_r = math.sin(r)
|
|
#c = cropped2.shape[0]//2
|
|
#M = cv2.getRotationMatrix2D( (c, c), rotate2-rotate, 1)
|
|
M = np.array( [
|
|
[cos_r, -1*sin_r, 0.0],
|
|
[sin_r, cos_r, 0.0]
|
|
] )
|
|
#print(M)
|
|
#M=np.array([16, 15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0,
|
|
# 26, 25, 24, 23, 22, 21, 20, 19, 18, 17, 27, 28, 29, 30, 35, 34, 33,
|
|
# 32, 31, 45, 44, 43, 42, 47, 46, 39, 38, 37, 36, 41, 40, 54, 53, 52,
|
|
# 51, 50, 49, 48, 59, 58, 57, 56, 55, 64, 63, 62, 61, 60, 67, 66, 65])
|
|
return cropped, label_out, cropped2, label2_out, M.flatten()
|
|
|
|
def next(self):
|
|
"""Returns the next batch of data."""
|
|
#print('next')
|
|
batch_size = self.batch_size
|
|
batch_data = nd.empty((batch_size,)+self.data_shape)
|
|
batch_label = nd.empty((batch_size,)+self.label_shape)
|
|
if self.use_coherent:
|
|
batch_label2 = nd.empty((batch_size,)+self.label_shape)
|
|
batch_coherent_label = nd.empty((batch_size,6))
|
|
i = 0
|
|
#self.cutoff = random.randint(800,1280)
|
|
try:
|
|
while i < batch_size:
|
|
#print('N', i)
|
|
data, label, annot = self.next_sample()
|
|
if not self.use_coherent:
|
|
R = self.do_aug(data, label, annot)
|
|
if R is None:
|
|
continue
|
|
data, label = R
|
|
#data, label, data2, label2, M = R
|
|
#ind = np.unravel_index(np.argmax(label[0], axis=None), label[0].shape)
|
|
#print(label.shape, np.count_nonzero(label[0]), ind)
|
|
#print(label[0,25:35,0:10])
|
|
data = nd.array(data)
|
|
data = nd.transpose(data, axes=(2, 0, 1))
|
|
label = nd.array(label)
|
|
#print(data.shape, label.shape)
|
|
try:
|
|
self.check_valid_image(data)
|
|
except RuntimeError as e:
|
|
logging.debug('Invalid image, skipping: %s', str(e))
|
|
continue
|
|
batch_data[i][:] = data
|
|
batch_label[i][:] = label
|
|
i += 1
|
|
else:
|
|
R = self.do_aug(data, label, annot)
|
|
if R is None:
|
|
continue
|
|
data, label, data2, label2, M = R
|
|
data = nd.array(data)
|
|
data = nd.transpose(data, axes=(2, 0, 1))
|
|
label = nd.array(label)
|
|
data2 = nd.array(data2)
|
|
data2 = nd.transpose(data2, axes=(2, 0, 1))
|
|
label2 = nd.array(label2)
|
|
M = nd.array(M)
|
|
#print(data.shape, label.shape)
|
|
try:
|
|
self.check_valid_image(data)
|
|
except RuntimeError as e:
|
|
logging.debug('Invalid image, skipping: %s', str(e))
|
|
continue
|
|
batch_data[i][:] = data
|
|
batch_label[i][:] = label
|
|
#batch_label2[i][:] = label2
|
|
batch_coherent_label[i][:] = M
|
|
#i+=1
|
|
j = i+self.per_batch_size//2
|
|
batch_data[j][:] = data2
|
|
batch_label[j][:] = label2
|
|
batch_coherent_label[j][:] = M
|
|
i += 1
|
|
if j%self.per_batch_size==self.per_batch_size-1:
|
|
i = j+1
|
|
except StopIteration:
|
|
if not i:
|
|
raise StopIteration
|
|
|
|
#return {self.data_name : batch_data,
|
|
# self.label_name : batch_label}
|
|
#print(batch_data.shape, batch_label.shape)
|
|
if not self.use_coherent:
|
|
return mx.io.DataBatch([batch_data], [batch_label], batch_size - i)
|
|
else:
|
|
#return mx.io.DataBatch([batch_data], [batch_label, batch_label2, batch_coherent_label], batch_size - i)
|
|
return mx.io.DataBatch([batch_data], [batch_label, batch_coherent_label], batch_size - i)
|
|
|
|
def check_data_shape(self, data_shape):
|
|
"""Checks if the input data shape is valid"""
|
|
if not len(data_shape) == 3:
|
|
raise ValueError('data_shape should have length 3, with dimensions CxHxW')
|
|
if not data_shape[0] == 3:
|
|
raise ValueError('This iterator expects inputs to have 3 channels.')
|
|
|
|
def check_valid_image(self, data):
|
|
"""Checks if the input data is valid"""
|
|
if len(data[0].shape) == 0:
|
|
raise RuntimeError('Data shape is wrong')
|
|
|
|
def imdecode(self, s):
|
|
"""Decodes a string or byte string to an NDArray.
|
|
See mx.img.imdecode for more details."""
|
|
return imdecode(s)
|
|
|
|
def read_image(self, fname):
|
|
"""Reads an input image `fname` and returns the decoded raw bytes.
|
|
|
|
Example usage:
|
|
----------
|
|
>>> dataIter.read_image('Face.jpg') # returns decoded raw bytes.
|
|
"""
|
|
with open(os.path.join(self.path_root, fname), 'rb') as fin:
|
|
img = fin.read()
|
|
return img
|
|
|
|
|