utils/annot.py

# annot.py
from io import StringIO
from random import sample
from collections import OrderedDict
import timecode as tc
from .behavior import Behavior
from sortedcontainers import SortedKeyList
from qtpy.QtCore import QObject, QRectF, Signal, Slot
from qtpy.QtGui import QColor
from qtpy.QtWidgets import QGraphicsItem

class Bout(object):
    """
    """

    def __init__(self, start, end, behavior):
        self._start = start
        self._end = end
        self._behavior = behavior

    def __lt__(self, b):
        if type(b) is tc.Timecode:
            return self._start.float < b.float
        elif type(b) is Bout:
            return self._start < b._start
        else:
            raise NotImplementedError(f"Comparing Bout with {type(b)} not supported")

    def __le__(self, b):
        if type(b) is tc.Timecode:
            return self._start.float <= b.is_float
        elif type(b) is Bout:
            return self._start <= b._start
        else:
            raise NotImplementedError(f'Comparing Bout with {type(b)} not supported')

    def is_at(self, t):
        return self._start <= t and self._end >= t

    def start(self):
        return self._start

    def set_start(self, start):
        self._start = start

    def end(self):
        return self._end

    def set_end(self, end):
        self._end = end

    def len(self):
        return self._end - self._start + tc.Timecode(self._start.framerate, frames=1)

    def behavior(self):
        return self._behavior

    def name(self):
        return self._behavior.get_name()

    def color(self):
        return self._behavior.get_color()

    def is_active(self):
        return self._behavior.is_active()

    def is_visible(self):
        return self._behavior.is_visible() and self._behavior.is_active()

    def __repr__(self):
        return f"Bout: start = {self.start()}, end = {self.end()}, behavior: {self.behavior()}"

class Channel(QGraphicsItem):
    """
    """

    contentChanged = Signal()

    def __init__(self, chan = None):
        super().__init__()
        if not chan is None:
            self._bouts_by_start = chan._bouts_by_start
            self._bouts_by_end = chan._bouts_by_end
            self._top = chan._top
        else:
            self._bouts_by_start = SortedKeyList(key=lambda bout: bout.start().float)
            self._bouts_by_end = SortedKeyList(key=lambda bout: bout.end().float)
            self._top = 0.
        self.cur_ix = 0
        self.fakeFirstBout = Bout(
            tc.Timecode('30.0', '0:0:0:0'),
            tc.Timecode('30.0', '0:0:0:0'),
            Behavior('none', '', QColor.fromRgbF(0., 0., 0.), active = True)
        )
        self.fakeLastBout = Bout(
            tc.Timecode('30.0', '23:59:59:29'),
            tc.Timecode('30.0', '23:59:59:29'),
            Behavior('none', '', QColor.fromRgbF(0., 0., 0.), active = True)
        )

    def add(self, b):
        if isinstance(b, Bout):
            self._bouts_by_start.add(b)
            self._bouts_by_end.add(b)
        else:
            raise TypeError("Can only add Bouts to Channel")

    def remove(self, b):
        # can raise ValueError if b is not in the channel
        self._bouts_by_start.remove(b)
        self._bouts_by_end.remove(b)

    def update_start(self, b, new_start):
        """
        Update the starting time of a bout while
        preserving the _bouts_by_start access order
        """
        self._bouts_by_start.remove(b)
        b.set_start(new_start)
        self._bouts_by_start.add(b)

    def update_end(self, b, new_end):
        """
        Update the ending time of a bout while
        preserving the _bouts_by_end access order
        """
        self._bouts_by_end.remove(b)
        b.set_end(new_end)
        self._bouts_by_end.add(b)

    def __add__(self, b):
        self.add(b)

    def _get_next(self, t, sortedlist):
        ix = sortedlist.bisect_key_right(t.float)
        l = len(sortedlist)
        if ix == l:
            return self.fakeLastBout, t.next()
        return sortedlist[ix], t.next()

    def _get_prev(self, t, sortedlist):
        ix = sortedlist.bisect_key_left(t.float)
        if ix == 0:
            return self.fakeFirstBout, t.back()
        return sortedlist[ix-1], t.back()

    def _get_inner(self, t, sortedList, op):
        t_local = t + 0 # kludgy copy constructor!
        visible = False
        while not visible:
            # no need to check for the end, because the fake first and last bouts are visible
            bout, t_local = op(t_local, sortedList)
            visible = bout.is_visible()
        return bout

    def get_next_start(self, t):
        return self._get_inner(t, self._bouts_by_start, self._get_next)

    def get_next_end(self, t):
        return self._get_inner(t, self._bouts_by_end, self._get_next)

    def get_prev_start(self, t):
        return self._get_inner(t, self._bouts_by_start, self._get_prev)

    def get_prev_end(self, t):
        return self._get_inner(t, self._bouts_by_end, self._get_prev)

    def get_in_range(self, start, end):
        """
        get all bouts that intersect the range [start, end]
        """
        return [bout for bout in self._bouts_by_start
            if bout.start().float <= end.float and bout.end().float >= start.float]

    def get_at(self, t):
        """
        get all bouts that span time t
        """
        return self.get_in_range(t, t)

    def __iter__(self):
        return iter(self._bouts_by_start)

    def irange(self, start_time, end_time):
        if isinstance(start_time, tc.Timecode):
            start_time = start_time.float
        if isinstance(end_time, tc.Timecode):
            end_time = end_time.float
        if not isinstance(start_time, float):
            raise TypeError(f"Can't handle start_time of type {type(start_time)}")
        if not isinstance(end_time, float):
            raise TypeError(f"Can't handle end_time of type {type(end_time)}")
        return self._bouts_by_start.irange_key(start_time, end_time)

    def top(self):
        return self._top

    def set_top(self, top):
        self._top = top

    def boundingRect(self):
        width = self.fakeLastBout.end().float
        return QRectF(0., self.top(), width, 1.)

    def paint(self, painter, option, widget=None):
        boundingRect = option.rect
        in_range_bouts = self._bouts_by_start.irange_key(boundingRect.left(), boundingRect.right())
        while True:
            try:
                bout = next(in_range_bouts)
            except StopIteration:
                break
            if bout.is_visible():
                painter.fillRect(
                    QRectF(bout.start().float, self.top(), bout.len().float, 1.),
                    bout.color()
                    )

    def _delete_all_inner(self, predicate):
        to_delete = list()
        # can't alter the bouts within the iterator
        for bout in iter(self):
            if predicate(bout):
                to_delete.append(bout)
        deleted_names = set()
        for bout in to_delete:
            deleted_names.add(bout.name())
            self.remove(bout)
        return deleted_names

    def delete_bouts_by_name(self, behavior_name):
        return self._delete_all_inner(lambda bout: bout.name() == behavior_name)

    def delete_inactive_bouts(self):
        return self._delete_all_inner(lambda bout: not bout.is_active())

    def truncate_or_remove_bouts(self, behavior, start, end, delete_all=False):
        """
        Delete bouts entirely within the range [start, end], and
        truncate bouts that extend outside the range.
        If behavior matches _deleteBehavior, the activity affects
        all bouts.  Otherwise, it only affects bouts with matching behavior.
        """
        items = self.get_in_range(start, end)
        for item in items:
            if not delete_all and behavior.get_name() != item.name():
                continue
            # Delete bouts that are entirely within the range
            if item.start() >= start and item.end() < end:
                print(f"removing {item} from active channel")
                self.remove(item)

            # Truncate and duplicate bouts that extend out both sides of the range
            if item.start() < start and item.end() > end:
                self.add(Bout(end, item.end(), item.behavior()))
                self.update_end(item, start)

            # Truncate bouts at the start boundary that start before the range
            elif item.start() < start and item.end() <= end:
                self.update_end(item, start)

            # Truncate bouts at the end boundary that end after the range
            elif item.start() >= start and item.end() > end:
                self.update_start(item, end)

            else:
                print(f"truncate_or_delete_bouts: Unexpected bout {item}")

    def coalesce_bouts(self, start, end):
        """
        combine overlapping bouts of the same behavior within [start, end]
        """
        to_delete = []
        items = self.get_in_range(start, end)
        # items will be ordered by start time
        for ix, first in enumerate(items):
            if first in to_delete:
                # previously coalesced
                continue
            if ix == len(items)-1:
                break
            for second in items[ix+1:]:
                if (first.name() == second.name() and
                    first.end() >= second.start()):
                    if first.end() < second.end():
                        self.update_end(first, second.end())
                    to_delete.append(second)
        for item in to_delete:
            self.remove(item)

class Annotations(QObject):
    """
    """

    # Signals
    annotations_changed = Signal()
    active_annotations_changed = Signal()

    def __init__(self, behaviors):
        super().__init__()
        self._channels = OrderedDict()
        self._behaviors = behaviors
        self._movies = []
        self._start_frame = None
        self._end_frame = None
        self._sample_rate = None
        self._stimulus = None
        self._format = None
        self.annotation_names = []
        behaviors.behaviors_changed.connect(self.note_annotations_changed)

    def read(self, fn):
        with open(fn, "r") as f:
            line = f.readline()
            line = line.strip().lower()
            if line.endswith("annotation file"):
                self._format = 'Caltech'
                self._read_caltech(f)
            elif line.startswith("scorevideo log"):
                self._format = 'Ethovision'
                self._read_ethovision(f)
            else:
                print("Unsupported annotation file format")

    def read_io(self, uploaded_file):
        text_str = uploaded_file.getvalue().decode("utf-8")
        with StringIO(text_str) as f:
            f.__setattr__('name', uploaded_file.name)
            line = f.readline()
            line = line.strip().lower()
            if line.endswith("annotation file"):
                self._format = 'Caltech'
                self._read_caltech(f)
            elif line.startswith("scorevideo log"):
                self._format = 'Ethovision'
                self._read_ethovision(f)
            else:
                print("Unsupported annotation file format")

    def _read_caltech(self, f):
        found_movies = False
        found_timecode = False
        found_stimulus = False
        found_channel_names = False
        found_annotation_names = False
        found_all_channels = False
        found_all_annotations = False
        new_behaviors_activated = False
        reading_channel = False
        to_activate = []
        channel_names = []
        current_channel = None
        current_bout = None

        self._format = 'Caltech'

        line = f.readline()
        while line:
            if found_annotation_names and not new_behaviors_activated:
                self.ensure_and_activate_behaviors(to_activate)
                new_behaviors_activated = True

            line.strip()

            if not line:
                if reading_channel:
                    reading_channel = False
                    current_channel = None
                    current_bout = None
            elif line.lower().startswith("movie file"):
                items = line.split()
                for item in items:
                    if item.lower().startswith("movie"):
                        continue
                    if item.lower().startswith("file"):
                        continue
                    self._movies.append(item)
                found_movies = True
            elif line.lower().startswith("stimulus name"):
                # TODO: do something when we know what one of these looks like
                found_stimulus = True
            elif line.lower().startswith("annotation start frame") or line.lower().startswith("annotation start time"):
                items = line.split()
                if len(items) > 3:
                    try:
                        self._start_frame = int(items[3])
                    except:
                        self._start_frame = int(float(items[3]))
                    if self._end_frame and self._sample_rate:
                        found_timecode = True
            elif line.lower().startswith("annotation stop frame") or  line.lower().startswith("annotation stop time"):
                items = line.split()
                if len(items) > 3:
                    try:
                        self._end_frame = int(items[3])
                    except:
                        self._end_frame = int(float(items[3]))
                    if self._start_frame and self._sample_rate:
                        found_timecode = True
            elif line.lower().startswith("annotation framerate"):
                items = line.split()
                if len(items) > 2:
                    self._sample_rate = float(items[2])
                    if self._start_frame and self._end_frame:
                        found_timecode = True
            elif line.lower().startswith("list of channels"):
                line = f.readline()
                while line:
                    line = line.strip()
                    if not line:
                        break # blank line -- end of section
                    channel_names.append(line)
                    line = f.readline()
                found_channel_names = True
            elif line.lower().startswith("list of annotations"):
                line = f.readline()
                while line:
                    line = line.strip()
                    if not line:
                        break # blank line -- end of section
                    to_activate.append(line)
                    line = f.readline().strip()
                found_annotation_names = True
            elif line.strip().lower().endswith("---"):
                for ch_name in channel_names:
                    if line.startswith(ch_name):
                        self._channels[ch_name] = Channel()
                        current_channel = ch_name
                        reading_channel = True
                        break
                if reading_channel:
                    reading_annot = False
                    line = f.readline()
                    while line:
                        line = line.strip()
                        if not line: # blank line
                            if reading_annot:
                                reading_annot = False
                                current_bout = None
                            else:
                                # second blank line, so done with channel
                                reading_channel = False
                                current_channel = None
                                break
                        elif line.startswith(">"):
                            current_bout = line[1:]
                            reading_annot = True
                        elif line.lower().startswith("start"):
                            pass # skip a header line
                        else:
                            items = line.split()
                            is_float = '.' in items[0] or '.' in items[1] or '.' in items[2]
                            bout_start = items[0]
                            if float(bout_start) < 1:
                                bout_start = 1
                            bout_end = items[1]
                            if float(bout_end) < 1:
                                bout_end = 1
                            self.add_bout(
                                Bout(
                                    tc.Timecode(self._sample_rate, start_seconds=float(bout_start)) if is_float
                                        else tc.Timecode(self._sample_rate, frames=int(bout_start)),
                                    tc.Timecode(self._sample_rate, start_seconds=float(bout_end)) if is_float
                                        else tc.Timecode(self._sample_rate, frames=int(bout_end)),
                                    self._behaviors.get(current_bout)),
                                current_channel)
                        line = f.readline()
            line = f.readline()
        print(f"Done reading Caltech annotation file {f.name}")
        self.note_annotations_changed()

    def write_caltech(self, f, video_files, stimulus):
        if not f.writable():
            raise RuntimeError("File not writable")
        f.write("Bento annotation file\n")
        f.write("Movie file(s):")
        for file in video_files:
            f.write(' ' + file)
        f.write('\n\n')

        f.write(f"Stimulus name: {stimulus}\n")
        f.write(f"Annotation start frame: {self._start_frame}\n")
        f.write(f"Annotation stop frame: {self._end_frame}\n")
        f.write(f"Annotation framerate: {self._sample_rate}\n")
        f.write("\n")

        f.write("List of Channels:\n")
        for ch in self.channel_names():
            f.write(ch + "\n")
        f.write("\n")

        f.write("List of annotations:\n")
        for annot in self.annotation_names:
            f.write(annot + "\n")
        f.write("\n")

        for ch in self.channel_names():
            by_name = {}
            f.write(f"{ch}----------\n")

            for bout in self.channel(ch):
                if not by_name.get(bout.name()):
                    by_name[bout.name()] = []
                by_name[bout.name()].append(bout)

            for annot in by_name:
                f.write(f">{annot}\n")
                f.write("Start\tStop\tDuration\n")
                for bout in by_name[annot]:
                    start = bout.start().frames
                    end = bout.end().frames
                    f.write(f"{start}\t{end}\t{end - start}\n")
                f.write("\n")

            f.write("\n")

        f.close()
        print(f"Done writing Caltech annotation file {f.name}")

    def _read_ethovision(self, f):
        print("Ethovision annotations not yet supported")

    def clear_channels(self):
        self._channels.clear()

    def channel_names(self):
        return list(self._channels.keys())

    def channel(self, ch: str) -> Channel:
        return self._channels[ch]

    def addEmptyChannel(self, ch: str):
        if ch not in self.channel_names():
            self._channels[ch] = Channel()

    def add_bout(self, bout, channel):
        if bout.name() not in self.annotation_names:
            self.annotation_names.append(bout.name())
        self._channels[channel].add(bout)
        if bout.end() > self.end_time():
            self.set_end_frame(bout.end())

    def start_time(self):
        """
        At some point we will need to support a start time distinct from
        frame number, perhaps derived from the OS file modify time
        or the start time of the corresponding video (or other media) file
        """
        if not self._start_frame or not self._sample_rate:
            return tc.Timecode('30.0', '0:0:0:0')
        return tc.Timecode(self._sample_rate, frames=self._start_frame)

    def start_frame(self):
        return self._start_frame

    def set_start_frame(self, t):
        if isinstance(t, int):
            self._start_frame = t
        elif isinstance(t, tc.Timecode):
            self._start_frame = t.frames
        else:
            raise TypeError("Expected a frame number or Timecode")

    def end_time(self):
        if not self._end_frame or not self._sample_rate:
            return tc.Timecode('30.0', '23:59:59:29')
        return tc.Timecode(self._sample_rate, frames=self._end_frame)

    def end_frame(self):
        return self._end_frame

    def set_end_frame(self, t):
        if isinstance(t, int):
            self._end_frame = t
        elif isinstance(t, tc.Timecode):
            self._end_frame = t.frames
        else:
            raise TypeError("Expected a frame number or Timecode")

    def sample_rate(self):
        return self._sample_rate

    def set_sample_rate(self, sample_rate):
        self._sample_rate = sample_rate

    def format(self):
        return self._format

    def set_format(self, format):
        self._format = format

    def delete_bouts_by_name(self, behavior_name):
        deleted_names = set()
        for chan_name in self.channel_names():
            deleted_names.update(self.channel(chan_name).delete_bouts_by_name(behavior_name))
        for name in deleted_names:
            self.annotation_names.remove(name)

    def delete_inactive_bouts(self):
        deleted_names = set()
        for chan_name in self.channel_names():
            deleted_names.update(self.channel(chan_name).delete_inactive_bouts())
        for name in deleted_names:
            self.annotation_names.remove(name)

    def ensure_and_activate_behaviors(self, toActivate):
        behaviorSetUpdated = False
        for behaviorName in toActivate:
            behaviorSetUpdated |= self._behaviors.addIfMissing(behaviorName)
            self.annotation_names.append(behaviorName)
            self._behaviors.get(behaviorName).set_active(True)
        if behaviorSetUpdated:
            self.annotations_changed.emit()
        self.active_annotations_changed.emit()

    def ensure_active_behaviors(self):
        for behavior in self._behaviors:
            if behavior.is_active() and behavior.get_name() not in self.annotation_names:
                self.annotation_names.append(behavior.get_name())

    def truncate_or_remove_bouts(self, behavior, start, end, chan):
        """
        Delete bouts entirely within the range [start, end], or
        truncate bouts that extend outside the range.
        If behavior matches _deleteBehavior, the activity affects
        all bouts.  Otherwise, it only affects bouts with matching behavior.
        """
        delete_all = (behavior.get_name() == self._behaviors.getDeleteBehavior().get_name())
        self._channels[chan].truncate_or_remove_bouts(behavior, start, end, delete_all)
        self.note_annotations_changed()

    def coalesce_bouts(self, start, end, chan):
        """
        combine overlapping bouts of the same behavior within [start, end]
        """
        self._channels[chan].coalesce_bouts(start, end)
        self.note_annotations_changed()

    @Slot()
    def note_annotations_changed(self):
        self.annotations_changed.emit()