<commit_before>import cv2
import numpy as np
from skimage.color import hsv2rgb, rgb2hsv, gray2rgb
from skimage.transform import resize
from skimage.exposure import adjust_gamma
import matplotlib

from pipeline.stages.stage import PipelineStage
from pipeline.objects import Image, Orientations, IDs, SaliencyOverlay, \
    LocalizerPositions, LocalizerPositionsOverlay, FinalResultOverlay, SaliencyImage, \
    CrownOverlay, LocalizerShapes


class SaliencyVisualizer(PipelineStage):
    requires = [Image, SaliencyImage]
    provides = [SaliencyOverlay]

    def __init__(self, hue=240 / 360., gamma=0.2):
        self.hue = hue
        self.gamma = gamma

    def call(self, image, saliency_image):
        img_resize = resize(image, saliency_image.shape)
        saliency_range = max(0.15, saliency_image.max() - saliency_image.min())
        saliency_norm = (saliency_image - saliency_image.min()) / saliency_range
        saliency_gamma = adjust_gamma(saliency_norm, gamma=self.gamma)
        cmap = matplotlib.cm.get_cmap('viridis')
        cmap_hsv = rgb2hsv(cmap(saliency_gamma)[:, :, :3])
        hsv = np.stack([
            cmap_hsv[:, :, 0],
            saliency_gamma,
            img_resize
        ], axis=-1)
        return hsv2rgb(hsv)


class LocalizerVisualizer(PipelineStage):
    requires = [Image, LocalizerPositions, LocalizerShapes]
    provides = [LocalizerPositionsOverlay]

    def __init__(self, roi_overlay='circle'):
        assert roi_overlay in ('rect', 'circle')
        self.roi_overlay = roi_overlay

    @staticmethod
    def get_roi_overlay(coordinates, image, roi_size):
        def roi_slice(coord):
            return slice(int(np.ceil(coord - roi_size / 2)),
                         int(np.ceil(coord + roi_size / 2)))

        pltim = np.zeros((image.shape[0], image.shape[1], 3))
        pltim[:, :, 0] = image
        pltim[:, :, 1] = image
        pltim[:, :, 2] = image
        for idx, (r, c) in enumerate(coordinates):
            sl_r = roi_slice(r)
            sl_c = roi_slice(c)
            pltim[sl_r, sl_c] = 1.

        return pltim

    @staticmethod
    def get_circle_overlay(coordinates, image, radius=32, line_width=8,
                           color=(1., 0, 0)):
        import cairocffi as cairo

        height, width = image.shape[:2]
        if image.ndim == 2:
            overlay = np.stack([image, image, image], axis=-1)
        elif image.ndim == 3 and image.shape[-1] == 3:
            overlay = image.copy()
        else:
            raise Exception("Did not understand image shape {}.".format(image.shape))

        image_surface = cairo.ImageSurface(cairo.FORMAT_A8, image.shape[1], image.shape[0])
        ctx = cairo.Context(image_surface)
        for y, x in coordinates:
            ctx.save()
            ctx.translate(int(x), int(y))
            ctx.new_path()
            ctx.arc(0, 0, radius + line_width / 2., 0, 2 * np.pi)
            ctx.close_path()
            ctx.set_source_rgba(0, 0, 0, 1)
            ctx.set_line_width(line_width)
            ctx.stroke()
            ctx.restore()

        image_surface.flush()
        circles = np.ndarray(shape=(height, width),
                             buffer=image_surface.get_data(),
                             dtype=np.uint8)
        circles_mask = (circles == 255)
        for i in range(3):
            overlay[circles_mask, i] = color[i]
        image_surface.finish()

        return overlay

    def call(self, image, positions, shapes):
        roi_size = shapes['roi_size']

        if self.roi_overlay == 'rect':
            return self.get_roi_overlay(positions, image / 255., roi_size)
        else:
            return self.get_circle_overlay(positions, image / 255.)


class ResultCrownVisualizer(PipelineStage):
    requires = [Image, LocalizerPositions, Orientations, IDs]
    provides = [CrownOverlay]

    def __init__(self, outer_radius=60, inner_radius=32,
                 orientation_radius=72,
                 true_hue=120 / 360., false_hue=240 / 360., orientation_hue=0 / 360.,
                 line_width=3):
        self.outer_radius = outer_radius
        self.inner_radius = inner_radius
        self.orientation_radius = orientation_radius
        self.true_hue = true_hue
        self.false_hue = false_hue
        self.orientation_hue = orientation_hue
        self.line_width = line_width

    def call(self, image, positions, orientations, ids):
        import cairocffi as cairo
        z_rots = orientations[:, 0]
        positions = np.stack([positions[:, 1], positions[:, 0]], axis=-1)
        height, width = image.shape
        # hsva
        # avsh
        image_surface = cairo.ImageSurface(cairo.FORMAT_ARGB32, width, height)
        ctx = cairo.Context(image_surface)
        ctx.set_antialias(cairo.ANTIALIAS_NONE)
        for z, pos, id in zip(z_rots, positions, ids):
            self._draw_crown(ctx, z, pos, id)
        image_surface.flush()
        overlay = np.ndarray(shape=(height, width, 4),
                             buffer=image_surface.get_data(),
                             dtype=np.uint8)
        overlay_hsva = overlay / 255.
        image_surface.finish()
        return overlay_hsva

    @staticmethod
    def _hsv2rgba(h, s, v, alpha=0):
        r, g, b = hsv2rgb(np.array([[[h, s, v]]])).flatten().tolist()
        return (b, g, r, alpha)

    def _draw_crown(self, ctx, angle, pos, bits):
        def x(w):<commit_msg>Avoid exception when generating CrownOverlay without any tag detections<commit_after>import cv2
import numpy as np
from skimage.color import hsv2rgb, rgb2hsv, gray2rgb
from skimage.transform import resize
from skimage.exposure import adjust_gamma
import matplotlib

from pipeline.stages.stage import PipelineStage
from pipeline.objects import Image, Orientations, IDs, SaliencyOverlay, \
    LocalizerPositions, LocalizerPositionsOverlay, FinalResultOverlay, SaliencyImage, \
    CrownOverlay, LocalizerShapes


class SaliencyVisualizer(PipelineStage):
    requires = [Image, SaliencyImage]
    provides = [SaliencyOverlay]

    def __init__(self, hue=240 / 360., gamma=0.2):
        self.hue = hue
        self.gamma = gamma

    def call(self, image, saliency_image):
        img_resize = resize(image, saliency_image.shape)
        saliency_range = max(0.15, saliency_image.max() - saliency_image.min())
        saliency_norm = (saliency_image - saliency_image.min()) / saliency_range
        saliency_gamma = adjust_gamma(saliency_norm, gamma=self.gamma)
        cmap = matplotlib.cm.get_cmap('viridis')
        cmap_hsv = rgb2hsv(cmap(saliency_gamma)[:, :, :3])
        hsv = np.stack([
            cmap_hsv[:, :, 0],
            saliency_gamma,
            img_resize
        ], axis=-1)
        return hsv2rgb(hsv)


class LocalizerVisualizer(PipelineStage):
    requires = [Image, LocalizerPositions, LocalizerShapes]
    provides = [LocalizerPositionsOverlay]

    def __init__(self, roi_overlay='circle'):
        assert roi_overlay in ('rect', 'circle')
        self.roi_overlay = roi_overlay

    @staticmethod
    def get_roi_overlay(coordinates, image, roi_size):
        def roi_slice(coord):
            return slice(int(np.ceil(coord - roi_size / 2)),
                         int(np.ceil(coord + roi_size / 2)))

        pltim = np.zeros((image.shape[0], image.shape[1], 3))
        pltim[:, :, 0] = image
        pltim[:, :, 1] = image
        pltim[:, :, 2] = image
        for idx, (r, c) in enumerate(coordinates):
            sl_r = roi_slice(r)
            sl_c = roi_slice(c)
            pltim[sl_r, sl_c] = 1.

        return pltim

    @staticmethod
    def get_circle_overlay(coordinates, image, radius=32, line_width=8,
                           color=(1., 0, 0)):
        import cairocffi as cairo

        height, width = image.shape[:2]
        if image.ndim == 2:
            overlay = np.stack([image, image, image], axis=-1)
        elif image.ndim == 3 and image.shape[-1] == 3:
            overlay = image.copy()
        else:
            raise Exception("Did not understand image shape {}.".format(image.shape))

        image_surface = cairo.ImageSurface(cairo.FORMAT_A8, image.shape[1], image.shape[0])
        ctx = cairo.Context(image_surface)
        for y, x in coordinates:
            ctx.save()
            ctx.translate(int(x), int(y))
            ctx.new_path()
            ctx.arc(0, 0, radius + line_width / 2., 0, 2 * np.pi)
            ctx.close_path()
            ctx.set_source_rgba(0, 0, 0, 1)
            ctx.set_line_width(line_width)
            ctx.stroke()
            ctx.restore()

        image_surface.flush()
        circles = np.ndarray(shape=(height, width),
                             buffer=image_surface.get_data(),
                             dtype=np.uint8)
        circles_mask = (circles == 255)
        for i in range(3):
            overlay[circles_mask, i] = color[i]
        image_surface.finish()

        return overlay

    def call(self, image, positions, shapes):
        roi_size = shapes['roi_size']

        if self.roi_overlay == 'rect':
            return self.get_roi_overlay(positions, image / 255., roi_size)
        else:
            return self.get_circle_overlay(positions, image / 255.)


class ResultCrownVisualizer(PipelineStage):
    requires = [Image, LocalizerPositions, Orientations, IDs]
    provides = [CrownOverlay]

    def __init__(self, outer_radius=60, inner_radius=32,
                 orientation_radius=72,
                 true_hue=120 / 360., false_hue=240 / 360., orientation_hue=0 / 360.,
                 line_width=3):
        self.outer_radius = outer_radius
        self.inner_radius = inner_radius
        self.orientation_radius = orientation_radius
        self.true_hue = true_hue
        self.false_hue = false_hue
        self.orientation_hue = orientation_hue
        self.line_width = line_width

    def call(self, image, positions, orientations, ids):
        import cairocffi as cairo
        has_detections = len(orientations) > 0
        if has_detections:
            z_rots = orientations[:, 0]
            positions = np.stack([positions[:, 1], positions[:, 0]], axis=-1)
        height, width = image.shape
        # hsva
        # avsh
        image_surface = cairo.ImageSurface(cairo.FORMAT_ARGB32, width, height)
        ctx = cairo.Context(image_surface)
        ctx.set_antialias(cairo.ANTIALIAS_NONE)
        if has_detections:
            for z, pos, id in zip(z_rots, positions, ids):
                self._draw_crown(ctx, z, pos, id)
        image_surface.flush()
        overlay = np.ndarray(shape=(height, width, 4),
                             buffer=image_surface.get_data(),
                             dtype=np.uint8)
        overlay_hsva = overlay / 255.
        image_surface.finish()
        return overlay_hsva

    @staticmethod
    def _hsv2rgba(h, s, v, alpha=0):
        r, g, b = hsv2rgb(np.array([[[h, s, v]]])).flatten().tolist()
        return (b, g, r, alpha)

    def _draw_crown(self, ctx, angle, pos, bits):
        def x(w):
{'fraction_non_alphanumeric': 0.10977047990565182, 'fraction_numerical': 0.01841603919078291, 'mean_word_length': 2.265402843601896, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 4, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3691700', 'n_tokens_mistral': 3632, 'n_tokens_neox': 3424, 'n_words': 760}