<commit_before>from typing import Callable, Dict, DefaultDict, Generic, Optional, TypeVar

from collections import defaultdict
from dataclasses import dataclass
from functools import partial

import trio
import cv2
import numpy as np

from hedgehog.protocol.errors import FailedCommandError
from hedgehog.protocol.messages import ack, vision as vision_msg

from . import CommandHandler, CommandRegistry
from .. import vision


T = TypeVar('T')


def pack_faces(feature) -> vision_msg.Feature:
    return vision_msg.FacesFeature([
        vision_msg.Face(rect)
        for rect in feature
    ])


def pack_blobs(feature) -> vision_msg.Feature:
    blobs, _ = feature
    return vision_msg.BlobsFeature([
        vision_msg.Blob(rect, centroid, confidence)
        for rect, centroid, confidence in blobs
    ])


@dataclass
class Channel(Generic[T]):
    msg: vision_msg.Channel
    detect: Callable[[np.ndarray], T]
    highlight: Callable[[np.ndarray, T], np.ndarray]
    pack: Callable[[T], vision_msg.Feature]


class ChannelData(Generic[T]):
    feature: Optional[T] = None
    highlight: Optional[np.ndarray] = None
    packed: Optional[vision_msg.Feature] = None

    def get_feature(self, img: np.ndarray, channel: Channel[T]):
        if self.feature is None:
            self.feature = channel.detect(img)
        return self.feature

    def get_highlight(self, img: np.ndarray, channel: Channel[T]):
        if self.highlight is None:
            feature = self.get_feature(img, channel)
            self.highlight = channel.highlight(img, feature)
        return self.highlight

    def get_packed(self, img: np.ndarray, channel: Channel[T]):
        if self.packed is None:
            feature = self.get_feature(img, channel)
            self.packed = channel.pack(feature)
        return self.packed


@dataclass
class Camera:
    grabber: vision.Grabber
    scope: trio.CancelScope


class VisionHandler(CommandHandler):
    _commands = CommandRegistry()

    open_count: int = 0
    _camera: Optional[Camera] = None

    _channels: Dict[str, Channel]
    _img: Optional[np.ndarray] = None
    _channel_data: Optional[DefaultDict[str, ChannelData]] = None

    def __init__(self) -> None:
        self._channels = {}

    async def _handle_camera(self, *, task_status=trio.TASK_STATUS_IGNORED):
        async with trio.open_nursery() as nursery:
            grabber = vision.Grabber(vision.camera())
            await nursery.start(grabber.run)
            task_status.started(Camera(grabber, nursery.cancel_scope))

    @_commands.register(vision_msg.OpenCameraAction)
    async def open_camera_action(self, server, ident, msg):
        if self.open_count == 0:
            self._camera = await server.add_task(self._handle_camera)
        self.open_count += 1

        return ack.Acknowledgement()

    @_commands.register(vision_msg.CloseCameraAction)
    async def close_camera_action(self, server, ident, msg):
        if self.open_count == 0:
            raise FailedCommandError("trying to close already closed camera")

        self.open_count -= 1
        if self.open_count == 0:
            self._camera.scope.cancel()
            self._camera = None
            self._img = None
            self._channel_data = None

        return ack.Acknowledgement()

    def set_channels(self, channels: Dict[str, vision_msg.Channel]):
        for key, channel in channels.items():
            if isinstance(channel, vision_msg.FacesChannel):
                self._channels[key] = Channel(
                    channel,
                    partial(vision.detect_faces, vision.haar_face_cascade),
                    vision.highlight_faces,
                    pack_faces,
                )
            elif isinstance(channel, vision_msg.BlobsChannel):
                self._channels[key] = Channel(
                    channel,
                    partial(vision.detect_blobs, min_hsv=channel.hsv_min, max_hsv=channel.hsv_max),
                    vision.highlight_blobs,
                    pack_blobs,
                )
            else:  # pragma: nocover
                assert False

    @_commands.register(vision_msg.CreateChannelAction)
    async def create_channel_action(self, server, ident, msg):<commit_msg>change server logic for vision: store which clients are connected. Each client can only open the camera once (because after the first close the camera might be closed for good). Unset all state if the camera task terminates, even if not closed on purpose (e.g. camera unplugged)<commit_after>from typing import Callable, Dict, DefaultDict, Generic, Optional, Set, TypeVar

from collections import defaultdict
from dataclasses import dataclass
from functools import partial

import trio
import cv2
import numpy as np

from hedgehog.protocol import Header
from hedgehog.protocol.errors import FailedCommandError
from hedgehog.protocol.messages import ack, vision as vision_msg

from . import CommandHandler, CommandRegistry
from .. import vision


T = TypeVar('T')


def pack_faces(feature) -> vision_msg.Feature:
    return vision_msg.FacesFeature([
        vision_msg.Face(rect)
        for rect in feature
    ])


def pack_blobs(feature) -> vision_msg.Feature:
    blobs, _ = feature
    return vision_msg.BlobsFeature([
        vision_msg.Blob(rect, centroid, confidence)
        for rect, centroid, confidence in blobs
    ])


@dataclass
class Channel(Generic[T]):
    msg: vision_msg.Channel
    detect: Callable[[np.ndarray], T]
    highlight: Callable[[np.ndarray, T], np.ndarray]
    pack: Callable[[T], vision_msg.Feature]


class ChannelData(Generic[T]):
    feature: Optional[T] = None
    highlight: Optional[np.ndarray] = None
    packed: Optional[vision_msg.Feature] = None

    def get_feature(self, img: np.ndarray, channel: Channel[T]):
        if self.feature is None:
            self.feature = channel.detect(img)
        return self.feature

    def get_highlight(self, img: np.ndarray, channel: Channel[T]):
        if self.highlight is None:
            feature = self.get_feature(img, channel)
            self.highlight = channel.highlight(img, feature)
        return self.highlight

    def get_packed(self, img: np.ndarray, channel: Channel[T]):
        if self.packed is None:
            feature = self.get_feature(img, channel)
            self.packed = channel.pack(feature)
        return self.packed


@dataclass
class Camera:
    grabber: vision.Grabber
    scope: trio.CancelScope
    closed: trio.Event


class VisionHandler(CommandHandler):
    _commands = CommandRegistry()

    _connections: Set[Header] = set()
    _camera: Optional[Camera] = None

    _channels: Dict[str, Channel]
    _img: Optional[np.ndarray] = None
    _channel_data: Optional[DefaultDict[str, ChannelData]] = None

    def __init__(self) -> None:
        self._channels = {}

    async def _handle_camera(self, *, task_status=trio.TASK_STATUS_IGNORED):
        try:
            async with trio.open_nursery() as nursery:
                grabber = vision.Grabber(vision.camera())
                await nursery.start(grabber.run)
                self._camera = Camera(grabber, nursery.cancel_scope, trio.Event())
                task_status.started()
        finally:
            self._camera.closed.set()
            self._connections.clear()
            self._camera = None
            self._img = None
            self._channel_data = None

    async def _close_camera(self):
        self._camera.scope.cancel()
        await self._camera.closed.wait()

    @_commands.register(vision_msg.OpenCameraAction)
    async def open_camera_action(self, server, ident, msg):
        if ident in self._connections:
            raise FailedCommandError("trying to open already opened camera")

        if not self._connections:
            await server.add_task(self._handle_camera)

        self._connections.add(ident)

        return ack.Acknowledgement()

    @_commands.register(vision_msg.CloseCameraAction)
    async def close_camera_action(self, server, ident, msg):
        if ident not in self._connections:
            # don't fail on closing if we're not open
            return ack.Acknowledgement()

        self._connections.remove(ident)

        if not self._connections and self._camera:
            await self._close_camera()

        return ack.Acknowledgement()

    def set_channels(self, channels: Dict[str, vision_msg.Channel]):
        for key, channel in channels.items():
            if isinstance(channel, vision_msg.FacesChannel):
                self._channels[key] = Channel(
                    channel,
                    partial(vision.detect_faces, vision.haar_face_cascade),
                    vision.highlight_faces,
                    pack_faces,
                )
            elif isinstance(channel, vision_msg.BlobsChannel):
                self._channels[key] = Channel(
                    channel,
                    partial(vision.detect_blobs, min_hsv=channel.hsv_min, max_hsv=channel.hsv_max),
                    vision.highlight_blobs,
                    pack_blobs,
                )
            else:  # pragma: nocover
                assert False

    @_commands.register(vision_msg.CreateChannelAction)
    async def create_channel_action(self, server, ident, msg):
{'fraction_non_alphanumeric': 0.08837811589511169, 'fraction_numerical': 0.0008632782993417503, 'mean_word_length': 2.5769972983404092, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 9, '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': '1060878', 'n_tokens_mistral': 2707, 'n_tokens_neox': 2540, 'n_words': 655}