work in progress
This commit is contained in:
SkalskiP
parent
88d63473fe
commit
a5b1079e39
|
|
@ -1847,7 +1847,14 @@ class CropAnnotator(BaseAnnotator):
|
|||
A class for drawing scaled up crops of detections on the scene.
|
||||
"""
|
||||
|
||||
def __init__(self, position: Position = Position.TOP_CENTER, scale_factor: int = 2):
|
||||
def __init__(
|
||||
self,
|
||||
position: Position = Position.TOP_CENTER,
|
||||
scale_factor: int = 2,
|
||||
color: Union[Color, ColorPalette] = ColorPalette.DEFAULT,
|
||||
thickness: int = 2,
|
||||
color_lookup: ColorLookup = ColorLookup.CLASS
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
position (Position): The anchor position for placing the cropped and scaled
|
||||
|
|
@ -1858,13 +1865,17 @@ class CropAnnotator(BaseAnnotator):
|
|||
"""
|
||||
self.position: Position = position
|
||||
self.scale_factor: int = scale_factor
|
||||
self.color: Union[Color, ColorPalette] = color
|
||||
self.thickness: int = thickness
|
||||
self.color_lookup: ColorLookup = color_lookup
|
||||
|
||||
@scene_to_annotator_img_type
|
||||
def annotate(
|
||||
self,
|
||||
scene: np.ndarray,
|
||||
scene: ImageType,
|
||||
detections: Detections,
|
||||
) -> np.ndarray:
|
||||
custom_color_lookup: Optional[np.ndarray] = None,
|
||||
) -> ImageType:
|
||||
"""
|
||||
Annotates the provided scene with scaled and cropped parts of the image based
|
||||
on the provided detections. Each detection is cropped from the original scene
|
||||
|
|
@ -1873,8 +1884,12 @@ class CropAnnotator(BaseAnnotator):
|
|||
|
||||
|
||||
Args:
|
||||
scene (np.ndarray): The image where cropped detection will be placed.
|
||||
scene (ImageType): The image where cropped detection will be placed.
|
||||
`ImageType` is a flexible type, accepting either `numpy.ndarray`
|
||||
or `PIL.Image.Image`.
|
||||
detections (Detections): Object detections to annotate.
|
||||
custom_color_lookup (Optional[np.ndarray]): Custom color lookup array.
|
||||
Allows to override the default color mapping strategy.
|
||||
|
||||
Returns:
|
||||
The annotated image.
|
||||
|
|
@ -1901,37 +1916,79 @@ class CropAnnotator(BaseAnnotator):
|
|||
]
|
||||
anchors = detections.get_anchors_coordinates(anchor=self.position).astype(int)
|
||||
|
||||
for resized_crop, anchor in zip(resized_crops, anchors):
|
||||
for idx, (resized_crop, anchor) in enumerate(zip(resized_crops, anchors)):
|
||||
crop_wh = resized_crop.shape[1], resized_crop.shape[0]
|
||||
crop_anchor = self.calculate_crop_coordinates(
|
||||
(x1, y1), (x2, y2) = self.calculate_crop_coordinates(
|
||||
anchor=anchor, crop_wh=crop_wh, position=self.position
|
||||
)
|
||||
scene = place_image(scene=scene, image=resized_crop, anchor=crop_anchor)
|
||||
scene = place_image(scene=scene, image=resized_crop, anchor=(x1, y1))
|
||||
color = resolve_color(
|
||||
color=self.color,
|
||||
detections=detections,
|
||||
detection_idx=idx,
|
||||
color_lookup=self.color_lookup
|
||||
if custom_color_lookup is None
|
||||
else custom_color_lookup,
|
||||
)
|
||||
cv2.rectangle(
|
||||
img=scene,
|
||||
pt1=(x1, y1),
|
||||
pt2=(x2, y2),
|
||||
color=color.as_bgr(),
|
||||
thickness=self.thickness,
|
||||
)
|
||||
|
||||
return scene
|
||||
|
||||
@staticmethod
|
||||
def calculate_crop_coordinates(
|
||||
anchor: Tuple[int, int], crop_wh: Tuple[int, int], position: Position
|
||||
) -> Tuple[int, int]:
|
||||
) -> Tuple[Tuple[int, int], Tuple[int, int]]:
|
||||
anchor_x, anchor_y = anchor
|
||||
width, height = crop_wh
|
||||
|
||||
if position == Position.TOP_LEFT:
|
||||
return anchor_x - width, anchor_y - height
|
||||
return (
|
||||
(anchor_x - width, anchor_y - height),
|
||||
(anchor_x, anchor_y)
|
||||
)
|
||||
elif position == Position.TOP_CENTER:
|
||||
return anchor_x - width // 2, anchor_y - height
|
||||
return (
|
||||
(anchor_x - width // 2, anchor_y - height),
|
||||
(anchor_x + width // 2, anchor_y)
|
||||
)
|
||||
elif position == Position.TOP_RIGHT:
|
||||
return anchor_x, anchor_y - height
|
||||
return (
|
||||
(anchor_x, anchor_y - height),
|
||||
(anchor_x + width, anchor_y)
|
||||
)
|
||||
elif position == Position.CENTER_LEFT:
|
||||
return anchor_x - width, anchor_y - height // 2
|
||||
return (
|
||||
(anchor_x - width, anchor_y - height // 2),
|
||||
(anchor_x, anchor_y + height // 2)
|
||||
)
|
||||
elif position == Position.CENTER or position == Position.CENTER_OF_MASS:
|
||||
return anchor_x - width // 2, anchor_y - height // 2
|
||||
return (
|
||||
(anchor_x - width // 2, anchor_y - height // 2),
|
||||
(anchor_x + width // 2, anchor_y + height // 2)
|
||||
)
|
||||
elif position == Position.CENTER_RIGHT:
|
||||
return anchor_x, anchor_y - height // 2
|
||||
return (
|
||||
(anchor_x, anchor_y - height // 2),
|
||||
(anchor_x + width, anchor_y + height // 2)
|
||||
)
|
||||
elif position == Position.BOTTOM_LEFT:
|
||||
return anchor_x - width, anchor_y
|
||||
return (
|
||||
(anchor_x - width, anchor_y),
|
||||
(anchor_x, anchor_y + height)
|
||||
)
|
||||
elif position == Position.BOTTOM_CENTER:
|
||||
return anchor_x - width // 2, anchor_y
|
||||
return (
|
||||
(anchor_x - width // 2, anchor_y),
|
||||
(anchor_x + width // 2, anchor_y + height)
|
||||
)
|
||||
elif position == Position.BOTTOM_RIGHT:
|
||||
return anchor_x, anchor_y
|
||||
return (
|
||||
(anchor_x, anchor_y),
|
||||
(anchor_x + width, anchor_y + height)
|
||||
)
|
||||
|
|
|
|||
Loading…
Reference in New Issue