Module facetorch.analyzer.utilizer.draw

Expand source code
import torch
import torchvision
from codetiming import Timer
from facetorch.base import BaseUtilizer
from facetorch.datastruct import ImageData
from facetorch.logger import LoggerJsonFile
from torchvision import transforms

logger = LoggerJsonFile().logger


class BoxDrawer(BaseUtilizer):
    def __init__(
        self,
        transform: transforms.Compose,
        device: torch.device,
        optimize_transform: bool,
        color: str,
        line_width: int,
    ):
        """Initializes the BoxDrawer class. This class is used to draw the face boxes to the image tensor.

        Args:
            transform (Compose): Composed Torch transform object.
            device (torch.device): Torch device cpu or cuda object.
            optimize_transform (bool): Whether to optimize the transform.
            color (str): Color of the boxes.
            line_width (int): Line width of the boxes.

        """
        super().__init__(transform, device, optimize_transform)
        self.color = color
        self.line_width = line_width

    @Timer("BoxDrawer.run", "{name}: {milliseconds:.2f} ms", logger=logger.debug)
    def run(self, data: ImageData) -> ImageData:
        """Draws face boxes to the image tensor.

        Args:
            data (ImageData): ImageData object containing the image tensor and face locations.
        Returns:
            ImageData: ImageData object containing the image tensor with face boxes.
        """
        loc_tensor = data.aggregate_loc_tensor()
        labels = [str(face.indx) for face in data.faces]
        data.img = torchvision.utils.draw_bounding_boxes(
            image=data.img,
            boxes=loc_tensor,
            labels=labels,
            colors=self.color,
            width=self.line_width,
        )

        return data


class LandmarkDrawerTorch(BaseUtilizer):
    def __init__(
        self,
        transform: transforms.Compose,
        device: torch.device,
        optimize_transform: bool,
        width: int,
        color: str,
    ):
        """Initializes the LandmarkDrawer class. This class is used to draw the 3D face landmarks to the image tensor.

        Args:
            transform (Compose): Composed Torch transform object.
            device (torch.device): Torch device cpu or cuda object.
            optimize_transform (bool): Whether to optimize the transform.
            width (int): Marker keypoint width.
            color (str): Marker color.

        """
        super().__init__(transform, device, optimize_transform)
        self.width = width
        self.color = color

    @Timer("LandmarkDrawer.run", "{name}: {milliseconds:.2f} ms", logger=logger.debug)
    def run(self, data: ImageData) -> ImageData:
        """Draws 3D face landmarks to the image tensor.

        Args:
            data (ImageData): ImageData object containing the image tensor and 3D face landmarks.
        Returns:
            ImageData: ImageData object containing the image tensor with 3D face landmarks.
        """
        data = self._draw_landmarks(data)

        return data

    def _draw_landmarks(self, data: ImageData) -> ImageData:
        """Draws 3D face landmarks to the image tensor.

        Args:
            data (ImageData): ImageData object containing the image tensor, 3D face landmarks, and faces.

        Returns:
            (ImageData): ImageData object containing the image tensor with 3D face landmarks.
        """

        if len(data.faces) > 0:
            pts = [face.preds["align"].other["lmk3d"].cpu() for face in data.faces]

            img_in = data.img.clone()
            pts = torch.stack(pts)
            pts = torch.swapaxes(pts, 2, 1)

            img_out = torchvision.utils.draw_keypoints(
                img_in,
                pts,
                colors=self.color,
                radius=self.width,
            )
            data.img = img_out

        return data

Classes

class BoxDrawer (transform: torchvision.transforms.transforms.Compose, device: torch.device, optimize_transform: bool, color: str, line_width: int)

Initializes the BoxDrawer class. This class is used to draw the face boxes to the image tensor.

Args

transform : Compose
Composed Torch transform object.
device : torch.device
Torch device cpu or cuda object.
optimize_transform : bool
Whether to optimize the transform.
color : str
Color of the boxes.
line_width : int
Line width of the boxes.
Expand source code
class BoxDrawer(BaseUtilizer):
    def __init__(
        self,
        transform: transforms.Compose,
        device: torch.device,
        optimize_transform: bool,
        color: str,
        line_width: int,
    ):
        """Initializes the BoxDrawer class. This class is used to draw the face boxes to the image tensor.

        Args:
            transform (Compose): Composed Torch transform object.
            device (torch.device): Torch device cpu or cuda object.
            optimize_transform (bool): Whether to optimize the transform.
            color (str): Color of the boxes.
            line_width (int): Line width of the boxes.

        """
        super().__init__(transform, device, optimize_transform)
        self.color = color
        self.line_width = line_width

    @Timer("BoxDrawer.run", "{name}: {milliseconds:.2f} ms", logger=logger.debug)
    def run(self, data: ImageData) -> ImageData:
        """Draws face boxes to the image tensor.

        Args:
            data (ImageData): ImageData object containing the image tensor and face locations.
        Returns:
            ImageData: ImageData object containing the image tensor with face boxes.
        """
        loc_tensor = data.aggregate_loc_tensor()
        labels = [str(face.indx) for face in data.faces]
        data.img = torchvision.utils.draw_bounding_boxes(
            image=data.img,
            boxes=loc_tensor,
            labels=labels,
            colors=self.color,
            width=self.line_width,
        )

        return data

Ancestors

Methods

def run(self, data: ImageData) ‑> ImageData

Draws face boxes to the image tensor.

Args

data : ImageData
ImageData object containing the image tensor and face locations.

Returns

ImageData
ImageData object containing the image tensor with face boxes.
Expand source code
@Timer("BoxDrawer.run", "{name}: {milliseconds:.2f} ms", logger=logger.debug)
def run(self, data: ImageData) -> ImageData:
    """Draws face boxes to the image tensor.

    Args:
        data (ImageData): ImageData object containing the image tensor and face locations.
    Returns:
        ImageData: ImageData object containing the image tensor with face boxes.
    """
    loc_tensor = data.aggregate_loc_tensor()
    labels = [str(face.indx) for face in data.faces]
    data.img = torchvision.utils.draw_bounding_boxes(
        image=data.img,
        boxes=loc_tensor,
        labels=labels,
        colors=self.color,
        width=self.line_width,
    )

    return data

Inherited members

class LandmarkDrawerTorch (transform: torchvision.transforms.transforms.Compose, device: torch.device, optimize_transform: bool, width: int, color: str)

Initializes the LandmarkDrawer class. This class is used to draw the 3D face landmarks to the image tensor.

Args

transform : Compose
Composed Torch transform object.
device : torch.device
Torch device cpu or cuda object.
optimize_transform : bool
Whether to optimize the transform.
width : int
Marker keypoint width.
color : str
Marker color.
Expand source code
class LandmarkDrawerTorch(BaseUtilizer):
    def __init__(
        self,
        transform: transforms.Compose,
        device: torch.device,
        optimize_transform: bool,
        width: int,
        color: str,
    ):
        """Initializes the LandmarkDrawer class. This class is used to draw the 3D face landmarks to the image tensor.

        Args:
            transform (Compose): Composed Torch transform object.
            device (torch.device): Torch device cpu or cuda object.
            optimize_transform (bool): Whether to optimize the transform.
            width (int): Marker keypoint width.
            color (str): Marker color.

        """
        super().__init__(transform, device, optimize_transform)
        self.width = width
        self.color = color

    @Timer("LandmarkDrawer.run", "{name}: {milliseconds:.2f} ms", logger=logger.debug)
    def run(self, data: ImageData) -> ImageData:
        """Draws 3D face landmarks to the image tensor.

        Args:
            data (ImageData): ImageData object containing the image tensor and 3D face landmarks.
        Returns:
            ImageData: ImageData object containing the image tensor with 3D face landmarks.
        """
        data = self._draw_landmarks(data)

        return data

    def _draw_landmarks(self, data: ImageData) -> ImageData:
        """Draws 3D face landmarks to the image tensor.

        Args:
            data (ImageData): ImageData object containing the image tensor, 3D face landmarks, and faces.

        Returns:
            (ImageData): ImageData object containing the image tensor with 3D face landmarks.
        """

        if len(data.faces) > 0:
            pts = [face.preds["align"].other["lmk3d"].cpu() for face in data.faces]

            img_in = data.img.clone()
            pts = torch.stack(pts)
            pts = torch.swapaxes(pts, 2, 1)

            img_out = torchvision.utils.draw_keypoints(
                img_in,
                pts,
                colors=self.color,
                radius=self.width,
            )
            data.img = img_out

        return data

Ancestors

Methods

def run(self, data: ImageData) ‑> ImageData

Draws 3D face landmarks to the image tensor.

Args

data : ImageData
ImageData object containing the image tensor and 3D face landmarks.

Returns

ImageData
ImageData object containing the image tensor with 3D face landmarks.
Expand source code
@Timer("LandmarkDrawer.run", "{name}: {milliseconds:.2f} ms", logger=logger.debug)
def run(self, data: ImageData) -> ImageData:
    """Draws 3D face landmarks to the image tensor.

    Args:
        data (ImageData): ImageData object containing the image tensor and 3D face landmarks.
    Returns:
        ImageData: ImageData object containing the image tensor with 3D face landmarks.
    """
    data = self._draw_landmarks(data)

    return data

Inherited members