def _augment(self, item, dsize):
image = item['image']
boxes = item['bbox'] # [n, 4]
h, w = image.shape[:2]
dw, dh = dsize
# random choose one face as interest region,
box = random.choice(boxes)
size = np.random.choice(
np.sqrt(np.prod(boxes[:, 2:] - boxes[:, :2], axis=-1)))
max_size = min(size, self.max_face)
max_scale = max_size / size
min_scale = min(dw / w, dh / h)
scale = random.uniform(min_scale, max_scale)
ibox = np.array([0, 0, w, h], dtype=np.float32)
matrix = matrix2d.scale(scale)
tbox = bbox_affine(ibox, matrix)
dbox = np.array([0, 0, dw, dh], dtype=np.float32)
fbox = bbox_affine(box, matrix)
# determine translation bound
tlow = tbox[:2] - dbox[:2]
thigh = tbox[2:] - dbox[2:]
flow = fbox[2:] - dbox[2:]
fhigh = fbox[:2] - dbox[:2]
thigh = np.minimum(thigh, fhigh)
tlow = np.maximum(tlow, flow)
high = np.maximum(tlow, thigh)
low = np.minimum(tlow, thigh)
txty = np.random.uniform(low, high)
translate = matrix2d.translate(-txty)
matrix = translate @ matrix
angle = np.random.uniform(-self.rotation, self.rotation)
rotate = matrix2d.center_rotate_scale_cw([dw / 2, dh / 2], angle, 1.0)
matrix = rotate @ matrix
mirror = random.choice([True, False])
if mirror:
matrix = matrix2d.hflip(dw) @ matrix
item = self._transform(item, matrix, random.choice(self.inters),
mirror)
# remove face exceed image
boxes = item['bbox'].copy()
area = np.prod(boxes[:, 2:] - boxes[:, :2], axis=-1)
boxes[boxes < 0] = 0
boxes[:, 2:] = np.minimum(boxes[:, 2:], [[dw, dh]])
inner = np.prod(boxes[:, 2:] - boxes[:, :2], axis=-1)
ratio = inner / area
discard = ratio < 0.4
item['label'][discard] = 0
return item
def _transform(self, item, matrix, inter, mirror):
image = item['image']
bboxes = item['bbox']
image = warp_affine(image,
matrix,
self.dsize,
flags=inter,
borderMode=cv2.BORDER_REPLICATE)
bboxes = bbox_affine(bboxes, matrix).astype(np.float32)
ret = {
'image': image,
'bbox': bboxes,
'label': np.ones(bboxes.shape[0], dtype=np.int64)
}
if 'shape' in item:
points = item['shape'].reshape(-1, 2)
points = points @ matrix[:2, :2].T + matrix[:2, 2]
points = points.reshape(bboxes.shape[0], -1, 2).astype(np.float32)
if mirror:
points = points[:, self.symmetry, :]
ret['shape'] = points
ret['mask'] = item['mask']
return ret
def meanbbox(self):
meanshape = self.meanshape
bboxes = []
for item in self.data:
shape = item['shape']
bbox = item['bbox'] # n, 4
matrix = umeyama(shape, meanshape)
bbox = bbox_affine(bbox, matrix)
bboxes.append(bbox)
return np.stack(bboxes, 0).mean(0)
def main(args):
workdir = Path(args.workdir)
workdir.mkdir(parents=True, exist_ok=True)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = models.__dict__[args.model.name](phase='train').to(device)
prior = BBoxShapePrior(args.num_classes, 5, args.anchors,
args.iou_threshold, args.encode_mean,
args.encode_std)
detector = Detector(prior, model)
state = Saver.load_best_from_folder(workdir, map_location='cpu')
detector.load_state_dict(state['model'])
predictor = Predictor(detector, args.test.score_threshold,
args.test.iou_threshold, device)
output_dir = workdir / 'result'
image_dir = Path(args.test.image_dir)
for imagefile in tqdm.tqdm(image_dir.rglob('*.jpg')):
relpath = imagefile.relative_to(image_dir)
outpath = output_dir / relpath
outpath.parent.mkdir(parents=True, exist_ok=True)
image = cv2.imread(str(imagefile), cv2.IMREAD_COLOR)
scale = (args.test.long_size / np.array(image.shape[:2])).min()
w, h = (np.array(image.shape[:2]) * scale).astype(np.int32)[::-1]
matrix = matrix2d.scale(scale)
warped = warp_affine(image, matrix, (w, h))
score, boxes, points = predictor.predict(warped)
n = score.size
if n > 0:
matrix = np.linalg.inv(matrix)
boxes = bbox_affine(boxes, matrix)
# k, p, 2
points = points.reshape(-1, 2) @ matrix[:2, :2].T + matrix[:2, 2]
points = points.reshape(score.size, -1, 2)
# draw result
draw_bbox(image, boxes)
draw_points(image, points)
cv2.imwrite(str(outpath), image)
# write output txt
resfile = outpath.with_suffix('.txt')
with open(resfile, 'wt') as f:
f.write('{}\n'.format(imagefile.name))
f.write(f'{n}\n')
if n > 0:
boxes = bbox2rect(boxes)
for score, box in zip(score.tolist(), boxes):
box = box.reshape(-1).tolist()
line = ' '.join(str(v) for v in box)
f.write(f'{line} {score}\n')
def __init__(self,
dsize,
padding,
meanshape,
meanbbox=None,
symmetry=None,
augments=None) -> None:
super().__init__()
dsize = _to_size(dsize)
meanshape = np.array(meanshape)
unit = to_unit_shape(meanshape)
ref = (unit - 0.5) * (1 - padding) * \
np.array(dsize).min() + np.array(dsize) / 2
if meanbbox:
meanbbox = np.array(meanbbox)
matrix = umeyama(meanshape, ref)
meanbbox = bbox_affine(meanbbox, matrix)
self.meanbbox = meanbbox
self.ref = ref
self.dsize = dsize
self.padding = padding
self.augments = augments
self.symmetry = symmetry
def __call__(self, item):
image = item['image']
bboxes = item['bbox'].reshape(-1, 4)
# calc transform matrix
h, w = image.shape[:2]
dw, dh = self.dsize
scale = min(dw / w, dh / h)
matrix = matrix2d.scale(scale)
image = warp_affine(image, matrix, self.dsize)
bboxes = bbox_affine(bboxes, matrix).astype(np.float32)
ret = {
'image': image,
'bbox': bboxes,
'label': np.ones(bboxes.shape[0], dtype=np.int64)
}
if 'shape' in item:
shapes = item['shape'].reshape(-1, 2)
shapes = shapes @ matrix[:2, :2].T + matrix[:2, 2]
ret['shape'] = shapes.reshape(bboxes.shape[0], -1,
2).astype(np.float32)
ret['mask'] = item['mask']
return ret