mxbiflow 0.1.1__py3-none-any.whl

mxbiflow/tasks/cross_modal/models.py

@@ -0,0 +1,57 @@
+from enum import StrEnum, auto
+
+from pydantic import BaseModel
+
+
+class CrossModalOutcome(StrEnum):
+    CORRECT = auto()
+    INCORRECT = auto()
+    TIMEOUT = auto()
+    ABORTED = auto()
+
+
+class CrossModalResultRecord(BaseModel):
+    timestamp: float
+    session_id: int | None
+    subject_id: str
+    partner_id: str
+
+    trial_id: str
+    trial_number: int
+
+    call_identity_id: str
+    call_category: str
+    is_partner_call: bool
+
+    other_identity_id: str
+    other_category: str
+
+    partner_side: str
+    correct_side: str
+
+    audio_identity_id: str
+    audio_index: int
+    audio_path: str
+
+    left_image_identity_id: str
+    left_image_index: int
+    left_image_path: str
+
+    right_image_identity_id: str
+    right_image_index: int
+    right_image_path: str
+
+    chosen_side: str | None
+    chosen_identity_id: str | None
+
+    outcome: CrossModalOutcome
+
+    trial_start_time: float | None
+    choice_time: float | None
+    latency_sec: float | None
+
+    choice_x: int | None
+    choice_y: int | None
+
+    aborted: bool
+    timeout: bool

mxbiflow/tasks/cross_modal/scene.py

@@ -0,0 +1,252 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from time import time
+from tkinter import CENTER, Canvas, Event
+from typing import TYPE_CHECKING
+
+import numpy as np
+from mxbi.tasks.cross_modal.config import CrossModalConfig
+from mxbi.utils.logger import logger
+from mxbi.utils.tkinter.components.canvas_with_border import CanvasWithInnerBorder
+from mxbi.utils.tkinter.components.showdata_widget import ShowDataWidget
+from numpy.typing import NDArray
+from PIL import ImageTk
+
+if TYPE_CHECKING:
+    from mxbi.models.animal import AnimalState
+    from mxbi.models.session import ScreenConfig, SessionState
+    from mxbi.tasks.cross_modal.trial_schema import Trial
+    from mxbi.theater import Theater
+    from PIL.Image import Image as PILImage
+
+
+@dataclass
+class CrossModalResult:
+    chosen_side: str | None
+    timeout: bool
+    feedback: bool
+    cancelled: bool
+    trial_start_time: float | None
+    choice_time: float | None
+    choice_x: int | None
+    choice_y: int | None
+
+
+class CrossModalScene:
+    def __init__(
+        self,
+        theater: "Theater",
+        session_state: "SessionState",
+        animal_state: "AnimalState",
+        screen: "ScreenConfig",
+        trial: "Trial",
+        cross_modal_config: CrossModalConfig,
+        left_image: "PILImage",
+        right_image: "PILImage",
+        audio_stimulus: "NDArray[np.int16]",
+    ) -> None:
+        self._theater = theater
+        self._session_state = session_state
+        self._animal_state = animal_state
+        self._screen = screen
+        self._trial = trial
+        self._cross_modal_config = cross_modal_config
+
+        self._left_image_pil = left_image
+        self._right_image_pil = right_image
+        self._audio_stimulus = audio_stimulus
+
+        self._background: CanvasWithInnerBorder | None = None
+        self._left_canvas: Canvas | None = None
+        self._right_canvas: Canvas | None = None
+        self._left_image: ImageTk.PhotoImage | None = None
+        self._right_image: ImageTk.PhotoImage | None = None
+        self._show_data_widget: ShowDataWidget | None = None
+
+        self._chosen_side: str | None = None
+        self._timeout = False
+        self._feedback = False
+        self._cancelled = False
+
+        self._trial_start_time: float | None = None
+        self._choice_time: float | None = None
+        self._choice_x: int | None = None
+        self._choice_y: int | None = None
+
+    def start(self) -> CrossModalResult:
+        self._create_view()
+        self._bind_events()
+        self._play_audio()
+        self._theater.mainloop()
+        return CrossModalResult(
+            chosen_side=self._chosen_side,
+            timeout=self._timeout,
+            feedback=self._feedback,
+            cancelled=self._cancelled,
+            trial_start_time=self._trial_start_time,
+            choice_time=self._choice_time,
+            choice_x=self._choice_x,
+            choice_y=self._choice_y,
+        )
+
+    def cancel(self) -> None:
+        if self._cancelled:
+            return
+        self._cancelled = True
+        self._feedback = False
+        self._timeout = False
+        self._stop_and_close()
+
+    def _create_view(self) -> None:
+        self._background = CanvasWithInnerBorder(
+            master=self._theater.root,
+            bg="black",
+            width=self._screen.width,
+            height=self._screen.height,
+            border_width=40,
+        )
+        self._background.place(relx=0.5, rely=0.5, anchor="center")
+        self._background.focus_set()
+
+        self._show_data_widget = ShowDataWidget(self._background)
+        self._show_data_widget.place(relx=0, rely=1, anchor="sw")
+        self._show_data_widget.show_data(
+            {
+                "name": self._animal_state.name,
+                "id": self._animal_state.trial_id,
+                "level_id": self._animal_state.current_level_trial_id,
+                "level": self._animal_state.level,
+                "rewards": 0,
+                "correct": self._animal_state.correct_trial,
+                "incorrect": 0,
+                "timeout": 0,
+            }
+        )
+
+        self._background.create_text(
+            self._screen.width / 2,
+            self._screen.height / 2,
+            text="+",
+            fill="white",
+            font=("Helvetica", 40),
+        )
+
+        fixation_ms = self._cross_modal_config.timing.fixation_ms
+        self._background.after(fixation_ms, self._show_images)
+
+    def _show_images(self) -> None:
+        if self._background is None:
+            return
+
+        self._background.delete("all")
+
+        img_size = int(self._left_image_pil.size[0])
+
+        self._left_canvas = Canvas(
+            self._background,
+            width=img_size,
+            height=img_size,
+            bg="grey",
+            highlightthickness=0,
+        )
+        self._left_canvas.place(relx=0.25, rely=0.5, anchor=CENTER)
+
+        self._right_canvas = Canvas(
+            self._background,
+            width=img_size,
+            height=img_size,
+            bg="grey",
+            highlightthickness=0,
+        )
+        self._right_canvas.place(relx=0.75, rely=0.5, anchor=CENTER)
+
+        if self._left_canvas is not None:
+            self._left_canvas.bind(
+                "<ButtonPress-1>", lambda e: self._on_choice("left", e)
+            )
+        if self._right_canvas is not None:
+            self._right_canvas.bind(
+                "<ButtonPress-1>", lambda e: self._on_choice("right", e)
+            )
+
+        self._left_image = ImageTk.PhotoImage(self._left_image_pil)
+        self._left_canvas.create_image(
+            img_size // 2, img_size // 2, image=self._left_image
+        )
+
+        self._right_image = ImageTk.PhotoImage(self._right_image_pil)
+        self._right_canvas.create_image(
+            img_size // 2, img_size // 2, image=self._right_image
+        )
+
+        self._background.create_text(
+            self._screen.width * 0.25,
+            self._screen.height * 0.8,
+            text=f"Left: {self._trial.left_image_identity_id}",
+            fill="white",
+        )
+        self._background.create_text(
+            self._screen.width * 0.75,
+            self._screen.height * 0.8,
+            text=f"Right: {self._trial.right_image_identity_id}",
+            fill="white",
+        )
+
+        self._trial_start_time = time()
+        timeout_ms = self._cross_modal_config.timing.trial_timeout_ms
+        self._background.after(timeout_ms, self._on_timeout)
+
+    def _bind_events(self) -> None:
+        if self._background is None:
+            return
+        self._background.bind("<r>", lambda e: self._give_manual_reward())
+        self._background.bind("<s>", lambda e: self._theater.caputre(self._background))
+
+    def _play_audio(self) -> None:
+        try:
+            self._theater.acontroller.set_master_volume(
+                self._cross_modal_config.audio.master_volume
+            )
+            self._theater.acontroller.set_digital_volume(
+                self._cross_modal_config.audio.digital_volume
+            )
+        except Exception:
+            logger.exception("Failed to set cross-modal audio volume")
+
+        self._theater.aplayer.play_stimulus(self._audio_stimulus)
+
+    def _on_choice(self, side: str, event: Event) -> None:
+        if self._cancelled or self._chosen_side is not None:
+            return
+        self._chosen_side = side
+        self._choice_time = time()
+        self._choice_x = event.x_root
+        self._choice_y = event.y_root
+        self._timeout = False
+        self._feedback = side == self._trial.correct_side
+        self._stop_and_close()
+
+    def _on_timeout(self) -> None:
+        if self._cancelled or self._chosen_side is not None:
+            return
+        self._chosen_side = None
+        self._timeout = True
+        self._feedback = False
+        self._stop_and_close()
+
+    def _stop_and_close(self) -> None:
+        self._theater.aplayer.stop()
+        if self._background is not None:
+            try:
+                self._background.destroy()
+            except Exception:
+                pass
+            self._background = None
+        try:
+            self._theater.root.quit()
+        except Exception:
+            pass
+
+    def _give_manual_reward(self) -> None:
+        self._theater.reward.give_reward(500)

mxbiflow/tasks/cross_modal/stage.py

@@ -0,0 +1,218 @@
+import time
+from pathlib import Path
+from typing import TYPE_CHECKING, Final
+
+from mxbi.data_logger import DataLogger
+from mxbi.tasks.cross_modal.bundle_dir import CrossModalBundleDir
+from mxbi.tasks.cross_modal.config import CrossModalConfig, load_cross_modal_config
+from mxbi.tasks.cross_modal.media import load_wav_as_int16
+from mxbi.tasks.cross_modal.models import CrossModalOutcome, CrossModalResultRecord
+from mxbi.tasks.cross_modal.scene import CrossModalResult, CrossModalScene
+from mxbi.tasks.cross_modal.trial_io import TrialCursor
+from mxbi.utils.logger import logger
+from PIL import Image
+
+if TYPE_CHECKING:
+    import numpy as np
+    from mxbi.models.animal import AnimalState
+    from mxbi.models.session import SessionState
+    from mxbi.models.task import Feedback
+    from mxbi.theater import Theater
+    from numpy.typing import NDArray
+    from PIL.Image import Image as PILImage
+
+
+class CrossModalTask:
+    STAGE_NAME: Final[str] = "cross_modal_task"
+
+    def __init__(
+        self,
+        theater: "Theater",
+        session_state: "SessionState",
+        animal_state: "AnimalState",
+    ) -> None:
+        self._theater = theater
+        self._session_state = session_state
+        self._animal_state = animal_state
+        self._screen = session_state.session_config.screen_type
+
+        bundle_dir_str = getattr(
+            session_state.session_config, "cross_modal_bundle_dir", None
+        )
+        if not bundle_dir_str:
+            raise RuntimeError(
+                "No cross-modal bundle directory configured. "
+                "Set session_config.cross_modal_bundle_dir via the launcher."
+            )
+
+        bundle_root = Path(bundle_dir_str).expanduser().resolve()
+        self._bundle_dir = CrossModalBundleDir.from_dir_path(bundle_root)
+
+        subject_id = self._animal_state.name
+        trials = self._bundle_dir.load_trials(subject_id)
+        if not trials:
+            raise RuntimeError(
+                f"Bundle contains zero trials for subject '{subject_id}'."
+            )
+
+        self._cursor = TrialCursor(
+            bundle_root=self._bundle_dir.root_dir, subject_id=subject_id
+        )
+        self._trial_index = self._cursor.next_index(len(trials))
+        self._trial = trials[self._trial_index]
+
+        self._cross_modal_config: CrossModalConfig = load_cross_modal_config()
+
+        image_size = int(
+            min(self._screen.width, self._screen.height)
+            * self._cross_modal_config.visual.image_scale
+        )
+
+        left_image_path = self._bundle_dir.resolve_media_path(
+            self._trial.left_image_path
+        )
+        right_image_path = self._bundle_dir.resolve_media_path(
+            self._trial.right_image_path
+        )
+        audio_path = self._bundle_dir.resolve_media_path(self._trial.audio_path)
+
+        left_image = self._prepare_image(
+            left_image_path,
+            image_size=image_size,
+        )
+        right_image = self._prepare_image(
+            right_image_path,
+            image_size=image_size,
+        )
+
+        audio_stimulus = load_wav_as_int16(
+            audio_path,
+            rate_policy=self._cross_modal_config.audio.wav_rate_policy,
+            gain=self._cross_modal_config.audio.gain,
+        )
+
+        self._scene = CrossModalScene(
+            theater=self._theater,
+            session_state=self._session_state,
+            animal_state=self._animal_state,
+            screen=self._screen,
+            trial=self._trial,
+            cross_modal_config=self._cross_modal_config,
+            left_image=left_image,
+            right_image=right_image,
+            audio_stimulus=audio_stimulus,
+        )
+
+        self._data_logger = DataLogger(
+            self._session_state, self._animal_state.name, self.STAGE_NAME
+        )
+        self._feedback = False
+
+    def start(self) -> "Feedback":
+        result = self._scene.start()
+        self._feedback = result.feedback
+
+        if not result.cancelled:
+            self._log_trial(result)
+
+        self._cursor.advance(self._trial_index)
+
+        logger.debug(
+            "cross_modal_task: session_id=%s, subject=%s, level=%s, "
+            "trial_index=%s, is_partner=%s, feedback=%s",
+            getattr(self._session_state, "session_id", None),
+            self._animal_state.name,
+            self._animal_state.level,
+            self._trial_index + 1,
+            self._trial.is_partner_call,
+            self._feedback,
+        )
+        return self._feedback
+
+    def quit(self) -> None:
+        self._scene.cancel()
+
+    def on_idle(self) -> None:
+        self._scene.cancel()
+
+    def on_return(self) -> None:
+        self._scene.cancel()
+
+    @property
+    def condition(self):
+        return None
+
+    def _log_trial(self, result: CrossModalResult) -> None:
+        try:
+            if result.timeout:
+                outcome = CrossModalOutcome.TIMEOUT
+            elif result.cancelled:
+                outcome = CrossModalOutcome.ABORTED
+            elif result.feedback:
+                outcome = CrossModalOutcome.CORRECT
+            else:
+                outcome = CrossModalOutcome.INCORRECT
+
+            if (
+                result.trial_start_time is not None
+                and result.choice_time is not None
+                and not result.timeout
+            ):
+                latency = result.choice_time - result.trial_start_time
+            else:
+                latency = None
+
+            if result.chosen_side == "left":
+                chosen_identity = self._trial.left_image_identity_id
+            elif result.chosen_side == "right":
+                chosen_identity = self._trial.right_image_identity_id
+            else:
+                chosen_identity = None
+
+            rec = CrossModalResultRecord(
+                timestamp=time.time(),
+                session_id=getattr(self._session_state, "session_id", None),
+                subject_id=self._trial.subject_id,
+                partner_id=self._trial.partner_id,
+                trial_id=self._trial.trial_id,
+                trial_number=self._trial.trial_number,
+                call_identity_id=self._trial.call_identity_id,
+                call_category=self._trial.call_category,
+                is_partner_call=self._trial.is_partner_call,
+                other_identity_id=self._trial.other_identity_id,
+                other_category=self._trial.other_category,
+                partner_side=self._trial.partner_side,
+                correct_side=self._trial.correct_side,
+                audio_identity_id=self._trial.audio_identity_id,
+                audio_index=self._trial.audio_index,
+                audio_path=self._trial.audio_path,
+                left_image_identity_id=self._trial.left_image_identity_id,
+                left_image_index=self._trial.left_image_index,
+                left_image_path=self._trial.left_image_path,
+                right_image_identity_id=self._trial.right_image_identity_id,
+                right_image_index=self._trial.right_image_index,
+                right_image_path=self._trial.right_image_path,
+                chosen_side=result.chosen_side,
+                chosen_identity_id=chosen_identity,
+                outcome=outcome,
+                trial_start_time=result.trial_start_time,
+                choice_time=result.choice_time,
+                latency_sec=latency,
+                choice_x=result.choice_x,
+                choice_y=result.choice_y,
+                aborted=result.cancelled,
+                timeout=result.timeout,
+            )
+
+            payload = rec.model_dump()
+            self._data_logger.save_jsonl(payload)
+            self._data_logger.save_csv_row(payload)
+        except Exception:
+            logger.exception("Failed to log cross-modal trial")
+
+    @staticmethod
+    def _prepare_image(image_path: Path, *, image_size: int) -> "PILImage":
+        with Image.open(image_path) as img:
+            prepared = img.convert("RGB").rotate(90, expand=True)
+            prepared = prepared.resize((image_size, image_size), Image.LANCZOS)
+            return prepared

mxbiflow/tasks/cross_modal/trial_io.py

@@ -0,0 +1,23 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+
+_TRIAL_INDEX_CACHE: dict[tuple[Path, str], int] = {}
+
+
+@dataclass(frozen=True)
+class TrialCursor:
+    bundle_root: Path
+    subject_id: str
+
+    def next_index(self, trial_count: int) -> int:
+        if trial_count <= 0:
+            raise ValueError("trial_count must be > 0")
+        key = (self.bundle_root.resolve(), self.subject_id)
+        idx = _TRIAL_INDEX_CACHE.get(key, 0)
+        return idx % trial_count
+
+    def advance(self, last_index: int) -> None:
+        key = (self.bundle_root.resolve(), self.subject_id)
+        _TRIAL_INDEX_CACHE[key] = last_index + 1

mxbiflow/tasks/cross_modal/trial_schema.py

@@ -0,0 +1,113 @@
+from pathlib import Path
+from typing import Any, Literal
+
+from pydantic import BaseModel, field_validator, model_validator
+
+
+class Trial(BaseModel):
+    trial_id: str
+    trial_number: int
+    subject_id: str
+    partner_id: str
+
+    call_identity_id: str
+    call_category: str
+    is_partner_call: bool
+
+    other_identity_id: str
+    other_category: str
+
+    partner_side: Literal["left", "right"]
+    correct_side: Literal["left", "right"]
+
+    audio_identity_id: str
+    audio_index: int
+    audio_path: str
+
+    left_image_identity_id: str
+    left_image_index: int
+    left_image_path: str
+
+    right_image_identity_id: str
+    right_image_index: int
+    right_image_path: str
+
+    seed: str
+
+    @model_validator(mode="before")
+    @classmethod
+    def _from_bundle_trial(cls, values: Any) -> Any:
+        if not isinstance(values, dict):
+            return values
+
+        if "trialId" not in values:
+            return values
+
+        audio = values.get("audio") if isinstance(values.get("audio"), dict) else {}
+        left_image = (
+            values.get("leftImage") if isinstance(values.get("leftImage"), dict) else {}
+        )
+        right_image = (
+            values.get("rightImage")
+            if isinstance(values.get("rightImage"), dict)
+            else {}
+        )
+
+        return {
+            "trial_id": values.get("trialId"),
+            "trial_number": values.get("trialNumber"),
+            "subject_id": values.get("subjectId"),
+            "partner_id": values.get("partnerId"),
+            "call_identity_id": values.get("callIdentityId"),
+            "call_category": values.get("callCategory"),
+            "is_partner_call": values.get("isPartnerCall"),
+            "other_identity_id": values.get("otherIdentityId"),
+            "other_category": values.get("otherCategory"),
+            "partner_side": values.get("partnerSide"),
+            "correct_side": values.get("correctSide"),
+            "audio_identity_id": audio.get("identityId"),
+            "audio_index": audio.get("exemplarIndex"),
+            "audio_path": audio.get("path"),
+            "left_image_identity_id": left_image.get("identityId"),
+            "left_image_index": left_image.get("exemplarIndex"),
+            "left_image_path": left_image.get("path"),
+            "right_image_identity_id": right_image.get("identityId"),
+            "right_image_index": right_image.get("exemplarIndex"),
+            "right_image_path": right_image.get("path"),
+            "seed": values.get("seed"),
+        }
+
+    @field_validator("partner_side", "correct_side")
+    @classmethod
+    def _normalize_sides(cls, v: str) -> str:
+        lv = v.strip().lower()
+        if lv not in {"left", "right"}:
+            raise ValueError("side must be 'left' or 'right'")
+        return lv
+
+    @field_validator("is_partner_call", mode="before")
+    @classmethod
+    def _boolify(cls, v):
+        if isinstance(v, bool):
+            return v
+        if isinstance(v, int):
+            return bool(v)
+        if isinstance(v, str):
+            lv = v.strip().lower()
+            if lv in {"true", "t", "yes", "y", "1"}:
+                return True
+            if lv in {"false", "f", "no", "n", "0"}:
+                return False
+        raise ValueError("is_partner_call must be boolean-like")
+
+    def audio_path_obj(self, base: Path | None = None) -> Path:
+        p = Path(self.audio_path)
+        return (base / p) if base and not p.is_absolute() else p
+
+    def left_image_path_obj(self, base: Path | None = None) -> Path:
+        p = Path(self.left_image_path)
+        return (base / p) if base and not p.is_absolute() else p
+
+    def right_image_path_obj(self, base: Path | None = None) -> Path:
+        p = Path(self.right_image_path)
+        return (base / p) if base and not p.is_absolute() else p