def test__given_small_image__when_downscaling_img__then_returns_same_img():
img = np.zeros((10, 20, 3))
scaler = ImgScaler(img_length_limit=100)
scaled_img = scaler.downscale_img(img)
assert (scaled_img == img).all()
def test__given_big_horizontal_image__when_downscaling_img__then_returns_downscaled_by_width_img():
img = np.zeros((100, 300, 3))
scaler = ImgScaler(img_length_limit=100)
scaled_img = scaler.downscale_img(img)
assert scaled_img.shape == (33, 100, 3)
def test__given_big_vertical_image__when_downscaling_img__then_returns_downscaled_by_height_img():
img = np.zeros((300, 100, 3))
scaler = ImgScaler(img_length_limit=100)
scaled_img = scaler.downscale_img(img)
assert scaled_img.shape == (100, 33, 3)
def find_faces(self,
img: Array3D,
det_prob_threshold: float = None) -> List[BoundingBoxDTO]:
if det_prob_threshold is None:
det_prob_threshold = self.det_prob_threshold
assert 0 <= det_prob_threshold <= 1
scaler = ImgScaler(self.IMG_LENGTH_LIMIT)
img = scaler.downscale_img(img)
results = self._detection_model.get(img, det_thresh=det_prob_threshold)
boxes = []
for result in results:
downscaled_box_array = result.bbox.astype(np.int).flatten()
downscaled_box = BoundingBoxDTO(x_min=downscaled_box_array[0],
y_min=downscaled_box_array[1],
x_max=downscaled_box_array[2],
y_max=downscaled_box_array[3],
probability=result.det_score,
np_landmarks=result.landmark)
box = downscaled_box.scaled(scaler.upscale_coefficient)
if box.probability <= det_prob_threshold:
logger.debug(
f'Box Filtered out because below threshold ({det_prob_threshold}: {box})'
)
continue
logger.debug(f"Found: {box}")
boxes.append(box)
return boxes
def find_faces(self, img: Array3D, det_prob_threshold: float = None) -> List[BoundingBoxDTO]:
if det_prob_threshold is None:
det_prob_threshold = self.det_prob_threshold
assert 0 <= det_prob_threshold <= 1
scaler = ImgScaler(self.IMG_LENGTH_LIMIT)
img = scaler.downscale_img(img)
fdn = self._face_detection_net
detect_face_result = fdn.detect_faces(img)
img_size = np.asarray(img.shape)[0:2]
bounding_boxes = []
for face in detect_face_result:
x, y, w, h = face['box']
box = BoundingBoxDTO(
x_min=int(np.maximum(x - (self.left_margin * w), 0)),
y_min=int(np.maximum(y - (self.top_margin * h), 0)),
x_max=int(np.minimum(x + w + (self.right_margin * w), img_size[1])),
y_max=int(np.minimum(y + h + (self.bottom_margin * h), img_size[0])),
np_landmarks=np.array([list(value) for value in face['keypoints'].values()]),
probability=face['confidence']
)
logger.debug(f"Found: {box}")
bounding_boxes.append(box)
filtered_bounding_boxes = []
for box in bounding_boxes:
box = box.scaled(scaler.upscale_coefficient)
if box.probability <= det_prob_threshold:
logger.debug(f'Box filtered out because below threshold ({det_prob_threshold}): {box}')
continue
filtered_bounding_boxes.append(box)
return filtered_bounding_boxes
def test__given_not_downscaled_image__when_upscaling_box__then_returns_same_box(
):
img = np.zeros((20, 20, 3))
scaler = ImgScaler(img_length_limit=100)
scaler.downscale_img(img)
output_box = BoundingBoxDTO(10, 10, 20, 20,
1).scaled(scaler.upscale_coefficient)
assert output_box == BoundingBoxDTO(10, 10, 20, 20, 1)
def find_faces(self,
img: Array3D,
det_prob_threshold: float = None) -> List[BoundingBoxDTO]:
if det_prob_threshold is None:
det_prob_threshold = self.det_prob_threshold
assert 0 <= det_prob_threshold <= 1
scaler = ImgScaler(self.IMG_LENGTH_LIMIT)
img = scaler.downscale_img(img)
fdn = self._face_detection_nets
detect_face_result = detect_face.detect_face(
img, self.FACE_MIN_SIZE, fdn.pnet, fdn.rnet, fdn.onet,
[self.det_threshold_a, self.det_threshold_b, self.det_threshold_c],
self.SCALE_FACTOR)
img_size = np.asarray(img.shape)[0:2]
bounding_boxes = []
detect_face_result = list(
zip(detect_face_result[0], detect_face_result[1].transpose()))
for result_item, landmarks in detect_face_result:
result_item = np.squeeze(result_item)
margin = self.BOX_MARGIN / 2
box = BoundingBoxDTO(
x_min=int(np.maximum(result_item[0] - margin, 0)),
y_min=int(np.maximum(result_item[1] - margin, 0)),
x_max=int(np.minimum(result_item[2] + margin, img_size[1])),
y_max=int(np.minimum(result_item[3] + margin, img_size[0])),
np_landmarks=landmarks.reshape(2, 5).transpose(),
probability=result_item[4])
logger.debug(f"Found: {box}")
bounding_boxes.append(box)
filtered_bounding_boxes = []
for box in bounding_boxes:
box = box.scaled(scaler.upscale_coefficient)
if box.probability <= det_prob_threshold:
logger.debug(
f'Box filtered out because below threshold ({det_prob_threshold}): {box}'
)
continue
filtered_bounding_boxes.append(box)
return filtered_bounding_boxes
def save_img(img: Array3D,
boxes: List[BoundingBoxDTO] = None,
noses: List[Tuple[int, int]] = None,
filepath: Union[Path, str] = None):
filepath = _get_filepath(filepath)
box_line_width = 3
font_size = 20
font_size_smaller = 15
radius = 7
img_length_limit = 1200
cross_half_length = 20
green_color = _to_rgb255(Color('#36cc36'))
error_color = _to_rgb255(Color('#ff4444'))
error_line_width = 3
def _draw_detection_box(text, box: BoundingBoxDTO, color):
img_draw.rectangle(box.xy, outline=color, width=box_line_width)
img_draw.text(text=text,
xy=(box.x_min, box.y_min - font_size - 1),
fill=color,
font=_get_font(font_size))
img_draw.text(text=f"{box.probability:.4f}",
xy=(box.x_min, box.y_max + 1),
fill=color,
font=_get_font(font_size_smaller))
scaler = ImgScaler(img_length_limit)
img = scaler.downscale_img(img)
pil_img = Image.fromarray(img, 'RGB')
img_draw = ImageDraw.Draw(pil_img)
noses_given = noses is not None
noses = [scaler.downscale_nose(nose) for nose in noses or ()]
boxes = [box.scaled(scaler.downscale_coefficient) for box in boxes or ()]
boxes = sorted(boxes, key=lambda box: (box.x_min, box.y_min))
draw_boxes = []
draw_error_boxes = []
draw_boxnoses = []
for box in boxes:
dot_drawn = False
if noses:
nearest_nose_idx = get_nearest_point_idx(box.center, noses)
nearest_nose = noses[nearest_nose_idx]
if box.is_point_inside(nearest_nose):
draw_boxnoses.append((box, nearest_nose))
noses.pop(nearest_nose_idx)
dot_drawn = True
if noses_given and not dot_drawn:
draw_error_boxes.append(box)
if not noses_given:
draw_boxes.append(box)
color_iter = _bright_color_gen()
no_errors_found = len(noses) == 0 and len(draw_error_boxes) == 0
for i, boxnose in enumerate(draw_boxnoses):
box, nose = boxnose
color = next(color_iter) if no_errors_found else green_color
_draw_detection_box(text=str(i + 1), box=box, color=color)
_draw_dot(img_draw, xy=nose, radius=radius, color=color)
for i, box in enumerate(draw_boxes):
color = next(color_iter) if no_errors_found else green_color
_draw_detection_box(text=str(i + 1), box=box, color=color)
for box in draw_error_boxes:
_draw_detection_box(text='Error', box=box, color=error_color)
for nose in noses:
_draw_cross(img_draw,
xy=nose,
half_length=cross_half_length,
color=error_color,
width=error_line_width)
pil_img.save(filepath, 'PNG')