Module facetorch.analyzer.utilizer.draw

Expand source code
import torch
import torchvision
import numpy as np
from codetiming import Timer
from facetorch.base import BaseUtilizer
from facetorch.datastruct import ImageData
from facetorch.logger import LoggerJsonFile
from torchvision import transforms
from matplotlib import pyplot as plt
from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas

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.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 LandmarkDrawer(BaseUtilizer):
    def __init__(
        self,
        transform: transforms.Compose,
        device: torch.device,
        optimize_transform: bool,
        marker: str,
        markersize: float,
        opacity: float,
        line_width: float,
        color: str,
        markeredgecolor: 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.
            marker (str): Marker type.
            markersize (float): Marker size.
            opacity (float): Opacity (transparency) of landmarks.
            line_width (float): Line width.
            color (str): Marker color.
            markeredgecolor (str): Marker edge color.

        """
        super().__init__(transform, device, optimize_transform)
        self.marker = marker
        self.markersize = markersize
        self.opacity = opacity
        self.line_width = line_width
        self.color = color
        self.markeredgecolor = markeredgecolor

    @Timer("LandmarkDrawer.run", "{name}: {milliseconds:.2f} ms", 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.
        """

        def _plot_close(pts_i, _i1, _i2):
            """Plots a line between two points."""
            plt.plot(
                [pts_i[0, _i1], pts_i[0, _i2]],
                [pts_i[1, _i1], pts_i[1, _i2]],
                color=self.color,
                lw=self.line_width,
                alpha=self.opacity - 0.1,
            )

        img = data.img.cpu().numpy().transpose(1, 2, 0)
        pts = [face.preds["align"].other["lmk3d"].cpu() for face in data.faces]

        if len(pts) > 0:

            height, width = img.shape[:2]
            fig = plt.figure(figsize=(height / 100, width / 100))
            plt.imshow(img)
            plt.subplots_adjust(left=0, right=1, top=1, bottom=0)
            plt.axis("off")

            if not type(pts) in [tuple, list]:
                pts = [pts]
            for i in range(len(pts)):

                nums = [0, 17, 22, 27, 31, 36, 42, 48, 60, 68]

                # close eyes and mouths
                _plot_close(pts[i], 41, 36)
                _plot_close(pts[i], 47, 42)
                _plot_close(pts[i], 59, 48)
                _plot_close(pts[i], 67, 60)

                for ind in range(len(nums) - 1):
                    l, r = nums[ind], nums[ind + 1]
                    plt.plot(
                        pts[i][0, l:r],
                        pts[i][1, l:r],
                        color=self.color,
                        lw=self.line_width,
                        alpha=self.opacity - 0.1,
                    )

                    plt.plot(
                        pts[i][0, l:r],
                        pts[i][1, l:r],
                        marker=self.marker,
                        linestyle="None",
                        markersize=self.markersize,
                        color=self.color,
                        markeredgecolor=self.markeredgecolor,
                        alpha=self.opacity,
                    )

            canvas = FigureCanvas(fig)
            canvas.draw()
            img_np = np.frombuffer(canvas.tostring_rgb(), dtype=np.uint8)
            img_np = img_np.reshape(canvas.get_width_height()[::-1] + (3,))
            img_np = img_np.transpose(2, 0, 1)
            data.img = torch.from_numpy(np.array(img_np))

            plt.close()

        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.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.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 LandmarkDrawer (transform: torchvision.transforms.transforms.Compose, device: torch.device, optimize_transform: bool, marker: str, markersize: float, opacity: float, line_width: float, color: str, markeredgecolor: 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.
marker : str
Marker type.
markersize : float
Marker size.
opacity : float
Opacity (transparency) of landmarks.
line_width : float
Line width.
color : str
Marker color.
markeredgecolor : str
Marker edge color.
Expand source code
class LandmarkDrawer(BaseUtilizer):
    def __init__(
        self,
        transform: transforms.Compose,
        device: torch.device,
        optimize_transform: bool,
        marker: str,
        markersize: float,
        opacity: float,
        line_width: float,
        color: str,
        markeredgecolor: 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.
            marker (str): Marker type.
            markersize (float): Marker size.
            opacity (float): Opacity (transparency) of landmarks.
            line_width (float): Line width.
            color (str): Marker color.
            markeredgecolor (str): Marker edge color.

        """
        super().__init__(transform, device, optimize_transform)
        self.marker = marker
        self.markersize = markersize
        self.opacity = opacity
        self.line_width = line_width
        self.color = color
        self.markeredgecolor = markeredgecolor

    @Timer("LandmarkDrawer.run", "{name}: {milliseconds:.2f} ms", 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.
        """

        def _plot_close(pts_i, _i1, _i2):
            """Plots a line between two points."""
            plt.plot(
                [pts_i[0, _i1], pts_i[0, _i2]],
                [pts_i[1, _i1], pts_i[1, _i2]],
                color=self.color,
                lw=self.line_width,
                alpha=self.opacity - 0.1,
            )

        img = data.img.cpu().numpy().transpose(1, 2, 0)
        pts = [face.preds["align"].other["lmk3d"].cpu() for face in data.faces]

        if len(pts) > 0:

            height, width = img.shape[:2]
            fig = plt.figure(figsize=(height / 100, width / 100))
            plt.imshow(img)
            plt.subplots_adjust(left=0, right=1, top=1, bottom=0)
            plt.axis("off")

            if not type(pts) in [tuple, list]:
                pts = [pts]
            for i in range(len(pts)):

                nums = [0, 17, 22, 27, 31, 36, 42, 48, 60, 68]

                # close eyes and mouths
                _plot_close(pts[i], 41, 36)
                _plot_close(pts[i], 47, 42)
                _plot_close(pts[i], 59, 48)
                _plot_close(pts[i], 67, 60)

                for ind in range(len(nums) - 1):
                    l, r = nums[ind], nums[ind + 1]
                    plt.plot(
                        pts[i][0, l:r],
                        pts[i][1, l:r],
                        color=self.color,
                        lw=self.line_width,
                        alpha=self.opacity - 0.1,
                    )

                    plt.plot(
                        pts[i][0, l:r],
                        pts[i][1, l:r],
                        marker=self.marker,
                        linestyle="None",
                        markersize=self.markersize,
                        color=self.color,
                        markeredgecolor=self.markeredgecolor,
                        alpha=self.opacity,
                    )

            canvas = FigureCanvas(fig)
            canvas.draw()
            img_np = np.frombuffer(canvas.tostring_rgb(), dtype=np.uint8)
            img_np = img_np.reshape(canvas.get_width_height()[::-1] + (3,))
            img_np = img_np.transpose(2, 0, 1)
            data.img = torch.from_numpy(np.array(img_np))

            plt.close()

        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.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