Skip to content

Commit

Permalink
bugfix: Extract - rotation in detection
Browse files Browse the repository at this point in the history
  • Loading branch information
@torzdf
torzdf committed Jun 15, 2023
1 parent 77723a6 commit ac8206e
Showing 1 changed file with 31 additions and 20 deletions.
51 changes: 31 additions & 20 deletions plugins/extract/detect/_base.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -174,7 +174,7 @@ def get_batch(self, queue: "Queue") -> Tuple[bool, DetectorBatch]:
{k: len(v) if isinstance(v, (list, np.ndarray)) else v
for k, v in batch.__dict__.items()})
else:
logger.trace(item) # type:ignore
logger.trace(item) # type:ignore[attr-defined]

if not exhausted and not batch.filename:
# This occurs when face filter is fed aligned faces.
Expand All @@ -201,7 +201,7 @@ def finalize(self, batch: BatchType) -> Generator[ExtractMedia, None, None]:
for the detected faces found in the frame.
"""
assert isinstance(batch, DetectorBatch)
logger.trace("Item out: %s", # type:ignore
logger.trace("Item out: %s", # type:ignore[attr-defined]
{k: len(v) if isinstance(v, (list, np.ndarray)) else v
for k, v in batch.__dict__.items()})

Expand Down Expand Up @@ -235,13 +235,13 @@ def finalize(self, batch: BatchType) -> Generator[ExtractMedia, None, None]:
output = self._extract_media.pop(filename)
output.add_detected_faces(batch.detected_faces[idx])

logger.trace("final output: (filename: '%s', image shape: %s, " # type:ignore
"detected_faces: %s, item: %s", output.filename, output.image_shape,
output.detected_faces, output)
logger.trace("final output: (filename: '%s', " # type:ignore[attr-defined]
"image shape: %s, detected_faces: %s, item: %s",
output.filename, output.image_shape, output.detected_faces, output)
yield output

@staticmethod
def _to_detected_face(left: float, top: float, right: float, bottom: float):
def _to_detected_face(left: float, top: float, right: float, bottom: float) -> DetectedFace:
""" Convert a bounding box to a detected face object
Parameters
Expand Down Expand Up @@ -276,9 +276,15 @@ def _predict(self, batch: BatchType) -> DetectorBatch:
# Rotate the batch and insert placeholders for already found faces
self._rotate_batch(batch, angle)
try:
batch.prediction = self.predict(batch.feed)
logger.trace("angle: %s, filenames: %s, prediction: %s", # type:ignore
angle, batch.filename, batch.prediction)
pred = self.predict(batch.feed)
if angle == 0:
batch.prediction = pred
else:
batch.prediction = np.array([b if b.any() else p
for b, p in zip(batch.prediction, pred)])
logger.trace("angle: %s, filenames: %s, " # type:ignore[attr-defined]
"prediction: %s",
angle, batch.filename, pred)
except tf_errors.ResourceExhaustedError as err:
msg = ("You do not have enough GPU memory available to run detection at the "
"selected batch size. You can try a number of things:"
Expand Down Expand Up @@ -308,20 +314,22 @@ def _predict(self, batch: BatchType) -> DetectorBatch:
raise

if angle != 0 and any(face.any() for face in batch.prediction):
logger.verbose("found face(s) by rotating image %s degrees", # type:ignore
logger.verbose("found face(s) by rotating image %s " # type:ignore[attr-defined]
"degrees",
angle)

found_faces = cast(List[np.ndarray], ([face if not found.any() else found
for face, found in zip(batch.prediction,
found_faces)]))

if all(face.any() for face in found_faces):
logger.trace("Faces found for all images") # type:ignore
logger.trace("Faces found for all images") # type:ignore[attr-defined]
break

batch.prediction = np.array(found_faces, dtype="object")
logger.trace("detect_prediction output: (filenames: %s, prediction: %s, " # type:ignore
"rotmat: %s)", batch.filename, batch.prediction, batch.rotation_matrix)
logger.trace("detect_prediction output: (filenames: %s, " # type:ignore[attr-defined]
"prediction: %s, rotmat: %s)",
batch.filename, batch.prediction, batch.rotation_matrix)
return batch

# <<< DETECTION IMAGE COMPILATION METHODS >>> #
Expand Down Expand Up @@ -349,7 +357,8 @@ def _compile_detection_image(self, item: ExtractMedia

image = self._scale_image(image, item.image_size, scale)
image = self._pad_image(image)
logger.trace("compiled: (images shape: %s, scale: %s, pad: %s)", # type:ignore
logger.trace("compiled: (images shape: %s, " # type:ignore[attr-defined]
"scale: %s, pad: %s)",
image.shape, scale, pad)
return image, scale, pad

Expand All @@ -367,7 +376,7 @@ def _set_scale(self, image_size: Tuple[int, int]) -> float:
The scaling factor from original image size to model input size
"""
scale = self.input_size / max(image_size)
logger.trace("Detector scale: %s", scale) # type:ignore
logger.trace("Detector scale: %s", scale) # type:ignore[attr-defined]
return scale

def _set_padding(self, image_size: Tuple[int, int], scale: float) -> Tuple[int, int]:
Expand Down Expand Up @@ -410,11 +419,12 @@ def _scale_image(image: np.ndarray, image_size: Tuple[int, int], scale: float) -
interpln = cv2.INTER_CUBIC if scale > 1.0 else cv2.INTER_AREA
if scale != 1.0:
dims = (int(image_size[1] * scale), int(image_size[0] * scale))
logger.trace("Resizing detection image from %s to %s. Scale=%s", # type:ignore
logger.trace("Resizing detection image from %s to %s. " # type:ignore[attr-defined]
"Scale=%s",
"x".join(str(i) for i in reversed(image_size)),
"x".join(str(i) for i in dims), scale)
image = cv2.resize(image, dims, interpolation=interpln)
logger.trace("Resized image shape: %s", image.shape) # type:ignore
logger.trace("Resized image shape: %s", image.shape) # type:ignore[attr-defined]
return image

def _pad_image(self, image: np.ndarray) -> np.ndarray:
Expand Down Expand Up @@ -442,7 +452,7 @@ def _pad_image(self, image: np.ndarray) -> np.ndarray:
pad_l,
pad_r,
cv2.BORDER_CONSTANT)
logger.trace("Padded image shape: %s", image.shape) # type:ignore
logger.trace("Padded image shape: %s", image.shape) # type:ignore[attr-defined]
return image

# <<< FINALIZE METHODS >>> #
Expand Down Expand Up @@ -634,7 +644,7 @@ def _rotate_image_by_angle(self,
https://stackoverflow.com/questions/22041699
"""

logger.trace("Rotating image: (image: %s, angle: %s)", # type:ignore
logger.trace("Rotating image: (image: %s, angle: %s)", # type:ignore[attr-defined]
image.shape, angle)
channels_first = image.shape[0] <= 4
if channels_first:
Expand All @@ -645,7 +655,8 @@ def _rotate_image_by_angle(self,
rotation_matrix = cv2.getRotationMatrix2D(image_center, -1.*angle, 1.)
rotation_matrix[0, 2] += self.input_size / 2 - image_center[0]
rotation_matrix[1, 2] += self.input_size / 2 - image_center[1]
logger.trace("Rotated image: (rotation_matrix: %s", rotation_matrix) # type:ignore
logger.trace("Rotated image: (rotation_matrix: %s", # type:ignore[attr-defined]
rotation_matrix)
image = cv2.warpAffine(image, rotation_matrix, (self.input_size, self.input_size))
if channels_first:
image = np.moveaxis(image, 2, 0)
Expand Down

0 comments on commit ac8206e

Please sign in to comment.