renderer_neural_result.py

#!/usr/bin/env python3

# Public licence for commercial/non-commercial use RRA (internationally)
# 
# Each Project IP Owner grants to, or must obtain for, each other Party and any member of the public a perpetual, irrevocable, worldwide, non-exclusive, royalty-free, non-transferable licence (including a right of sub-license to any person (in the case of GBRF including, but not limited to, the Department)) to Use the Project IP and Project Improvements, in the field of reef restoration and adaptation, for:
# 
#   (a) non-commercial purposes, educational and/or research purposes (including for the performance of Core Commonwealth Functions by the Department); and/or
#   (b) commercial purposes, whether in Australia or elsewhere.
# 
# Each Project IP Owner acknowledges that any licence granted to the Department for Core Commonwealth Functions will not be restricted to use in the field of reef restoration and adaptation.
# 
# Derivative works must be distributed with a copy of this licence which does not further restrict the rights of licensees.
# 
# Amendment providing additional Limitation of Liability for QUT
# 
# QUT does not warrant that:
# 
#   (a) software/code is fit for the Approved Purpose, or that it has any particular qualities or characteristics;
#   (b) the software/code is free from errors, viruses, worms, or similar defects;
#   (c) the use of the software/code by the Licensee will lead to any particular result; or
#   (d) the use of the software/code will not infringe the rights (including Intellectual Property rights) of any person.
# 
# Copyright (C) 2025 Queensland University of Technology
# 

# Author:   Alec Tutin
# Date:     2024-06-17

import colorsys, cv2
from cslics_mqtt.comms import Box
from datetime import datetime
from logging import Logger
from numpy import ndarray
from typing import List

class RendererNeuralResult:
    """A class to handle rendering of neural processing results onto images."""

    def __init__(self, labels: List[str], logger: Logger, colour_timestamp: cv2.typing.Scalar = (255, 255, 255), label_saturation: float = 0.8, label_value: float = 0.8):
        """
        Args:
            labels: A list of the labels in the model.
            logger: The logger to use to report errors and warnings.
            colour_timestamp: The colour to use when drawing the timestamp.
            label_saturation: The colour saturation to use when drawing labels and boxes.
            label_value: The colour value to use when drawing labels and boxes.
        """

        self.__logger: Logger = logger.getChild(RendererNeuralResult.__name__)

        #: The model labels.
        self.labels: List[str] = labels
        #: The font to use when drawing text.
        self.font: int = cv2.FONT_HERSHEY_SIMPLEX
        #: The scale to use for the font when drawing text.
        self.font_scale: float = 5.0
        #: The thickness to render the fonts in when drawing text.
        self.font_thickness: int = 4
        #: The colour to use when drawing the timestamp.
        self.colour_timestamp: cv2.typing.Scalar = colour_timestamp
        text_size, _ = cv2.getTextSize('I', self.font, self.font_scale, self.font_thickness)
        #: The height of the font used for drawing text.
        self.font_height: int = text_size[1]
        #: The colour saturation to use when drawing labels and boxes.
        self.label_saturation: float = label_saturation
        #: The colour value to use when drawing labels and boxes.
        self.label_value: float = label_value

    def get_label_colour(self, label: int) -> cv2.typing.Scalar:
        """Get the colour for a label by index.

        Args:
            label: The index of the label.

        Returns:
            The colour to use as a OpenCV2 Scalar.
        """

        label_count: int = len(self.labels)
        label_t: float = label / label_count
        rgb = colorsys.hsv_to_rgb(label_t, self.label_saturation, self.label_value)
        return (int(rgb[2] * 255), int(rgb[1] * 255), int(rgb[0] * 255))

    def draw_labels(self, mat: ndarray, counts: List[int]) -> None:
        """Draw the `self.labels` onto the supplied image.

        Args:
            mat: The image.
            counts: The counts for each label.
        """

        height, _, _ = mat.shape

        if len(self.labels) != len(counts):
            self.__logger.warning(f'Number of counts ({len(counts)}) does not match model labels ({len(self.labels)})!')

        label_count: int = len(self.labels)
        
        margin: int = int(self.font_height * 0.5)
        pixels_per_label: int = int(self.font_height * 1.2)

        for i in range(label_count):
            label: str = self.labels[i]

            label_position: int = (label_count - i - 1)
            origin: cv2.typing.Point = (10, height - margin - label_position * pixels_per_label)
            colour: cv2.typing.Scalar = self.get_label_colour(i)

            count: int = counts[i] if len(counts) > i else 0

            cv2.putText(mat, f'{label} - {count}', origin, self.font, self.font_scale, colour, self.font_thickness)
    
    def draw_boxes(self, mat: ndarray, boxes: List[Box]) -> None:
        """Draw the boxes onto the supplied image.

        Args:
            mat: The image.
            boxes: A list of boxes for the detections in the image.
        """
        
        height, width, _ = mat.shape

        if width == 0 or height == 0:
            return
        
        for box in boxes:
            bl = (int(box.min_x * width), int(box.min_y * height))
            tr = (int(box.max_x * width), int(box.max_y * height))

            cv2.rectangle(mat, bl, tr, self.get_label_colour(box.label), 4)
    
    def draw_timestamp(self, mat: ndarray, timestamp: datetime) -> None:
        """Draw a timestamp onto the supplied image.

        Args:
            mat: The image.
            timestamp: The timestamp to draw.
        """

        height, width, _ = mat.shape

        if width == 0 or height == 0:
            return
        
        timestamp_string: str = timestamp.strftime('%Y-%m-%d %H:%M:%S')
        text_size, _ = cv2.getTextSize(timestamp_string, self.font, self.font_scale, self.font_thickness)
        origin: cv2.typing.Point = (width - text_size[0] - 10, height - 40)
        cv2.putText(mat, timestamp_string, origin, self.font, self.font_scale, self.colour_timestamp, self.font_thickness)