Pipeline¶

def reset(self):
    """Clear the vote cache and any state the primary keeps between frames."""
    self.refiner.reset()
    reset = getattr(self.primary, "reset", None)
    if callable(reset):
        reset()

def wait(self):
    """Block until every background request has come back. See ``workers``."""
    self.refiner.wait()

def close(self):
    """Stop the background workers, dropping anything still in flight."""
    self.refiner.close()

def step(self, img, preds=None):
    """Process one frame and return the refined [Tracks][vizor.boxes.Tracks]."""
    tracks = self.primary.track(img)
    if not isinstance(tracks, Tracks):
        tracks = Tracks(tracks)
    if tracks.names is None:
        tracks.names = getattr(self.primary, "names", None)
    return self.refiner.run(tracks, img=img, preds=preds)

def run(self, src, save=None, show=False, fourcc="mp4v"):
    """Yield refined tracks for every frame of ``src``.

    ``src`` is anything [Video][vizor.utils.video.Video] opens: a file path,
    a camera index, or a stream url. Pass ``save`` to also write an annotated
    video, and ``show`` to display it in a window.

    Each yielded ``Tracks`` carries its frame, so ``out.draw()`` needs no
    argument. Drawing happens on the frame itself, so copy it first if you
    want the original.
    """
    video = Video(src)
    writer = Writer(save, fps=video.fps, fourcc=fourcc) if save else None
    try:
        for frame in video:
            out = self.step(frame)
            yield out
            if writer or show:
                drawn = out.draw(frame)
                if writer:
                    writer.write(drawn)
                if show:
                    import cv2

                    cv2.imshow("vizor", drawn)
                    if cv2.waitKey(1) & 0xFF in (27, ord("q")):
                        break
    finally:
        video.close()
        if writer:
            writer.close()
        if show:
            import cv2

            try:
                cv2.destroyWindow("vizor")
            except cv2.error:
                pass

def save(self, src, out, show=False):
    """Run over ``src`` and write the annotated video to ``out``."""
    for _ in self.run(src, save=out, show=show):
        pass
    return out

def wait(self):
    """Block until every request in flight has come back, then bank the votes.

    Only useful with ``workers``. The votes land too late for the frames that
    triggered them, but they are there for whatever you refine next.
    """
    for _, fut in list(self.jobs):
        fut.exception()
    self._drain()

def close(self):
    """Stop the workers. Anything still in flight is dropped."""
    self.jobs.clear()
    self.busy.clear()
    self.shots.clear()
    if self.pool is not None:
        self.pool.shutdown(wait=False, cancel_futures=True)
        self.pool = None

def reset(self):
    """Forget every vote and drop anything in flight. Call this between videos."""
    self.cache.clear()
    for _, fut in self.jobs:
        fut.cancel()
    self.jobs.clear()
    self.busy.clear()
    self.shots.clear()
    self.frame = -1

def run(self, tracks, img=None, preds=None):
    """Refine ``tracks`` in place and return them.

    ``img`` is the current frame. ``preds`` lets you supply the secondary's
    output yourself, which skips the model call entirely.
    """
    if not isinstance(tracks, Tracks):
        tracks = Tracks(tracks)
    if self.names is None:
        self.names = tracks.names
    self.frame += 1
    if len(tracks):
        if self.mode == "full":
            self._full(tracks, img, preds)
        elif self.mode == "crop":
            self._crop(tracks, img)
        else:
            self._collage(tracks, img)
        self._apply(tracks)
    tracks.names = self.names if self.names is not None else tracks.names
    tracks.img = img if img is not None else tracks.img
    return tracks

def add(self, id, cls):
    """Record one vote of ``cls`` for track ``id``."""
    id = int(id)
    if id < 0:
        return
    dq = self.data.get(id)
    if dq is None:
        dq = self.data[id] = deque(maxlen=self.hist)
        while len(self.data) > self.size:
            self.data.popitem(last=False)
    self.data.move_to_end(id)
    dq.append(int(cls))

def get(self, id, default=None):
    """Most common class voted for ``id``, or ``default`` if it has no votes."""
    id = int(id)
    dq = self.data.get(id) if id >= 0 else None
    if not dq:
        return default
    self.data.move_to_end(id)
    return Counter(dq).most_common(1)[0][0]

def count(self, id):
    """How many votes ``id`` has, capped at ``hist``."""
    dq = self.data.get(int(id))
    return len(dq) if dq else 0

def drop(self, id):
    """Forget one track id. Does nothing if it was never seen."""
    self.data.pop(int(id), None)

def clear(self):
    """Forget every track id."""
    self.data.clear()