120 lines
3.5 KiB
Python
120 lines
3.5 KiB
Python
import abc
|
|
from pathlib import Path
|
|
from dodgeballtracker.typing import VideoFrame, DetectedObject, Color
|
|
import matplotlib
|
|
import cv2
|
|
|
|
|
|
class DetectedObjectBox:
|
|
def __init__(self, min_threshold: float, color: Color):
|
|
self.min_threshold = min_threshold
|
|
self.color = color
|
|
|
|
def __lt__(self, other):
|
|
return self.min_threshold < other.min_threshold
|
|
|
|
@classmethod
|
|
def from_viridis(cls, min_threshold):
|
|
cm = matplotlib.colormaps['viridis']
|
|
color: Color = tuple((int(255 * cm(min_threshold)[i]) for i in range(3)))
|
|
return cls(min_threshold, color)
|
|
|
|
|
|
def create_video_writer(video_cap: cv2.VideoCapture, output_filename: Path) -> cv2.VideoWriter:
|
|
|
|
# grab the width, height, and fps of the frames in the video stream.
|
|
frame_width = int(video_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
|
|
frame_height = int(video_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
|
|
fps = int(video_cap.get(cv2.CAP_PROP_FPS))
|
|
|
|
# initialize the FourCC and a video writer object
|
|
fourcc = cv2.VideoWriter_fourcc(*'mp4v')
|
|
writer = cv2.VideoWriter(str(output_filename.absolute()),
|
|
fourcc,
|
|
fps,
|
|
(frame_width, frame_height))
|
|
|
|
return writer
|
|
|
|
|
|
class VideoWriter(abc.ABC):
|
|
|
|
@abc.abstractmethod
|
|
def write_frame(self, frame: VideoFrame) -> None:
|
|
pass
|
|
|
|
@abc.abstractmethod
|
|
def plot_boxes_to_frame(self, data: list[DetectedObject], frame: VideoFrame) -> None:
|
|
pass
|
|
|
|
@abc.abstractmethod
|
|
def close(self) -> None:
|
|
pass
|
|
|
|
@abc.abstractmethod
|
|
def wait(self) -> bool:
|
|
pass
|
|
|
|
|
|
class NoopVideoWriter(VideoWriter):
|
|
|
|
def plot_boxes_to_frame(self, boxes: list[DetectedObject], frame: VideoFrame) -> None:
|
|
pass
|
|
|
|
def write_frame(self, frame: VideoFrame) -> None:
|
|
pass
|
|
|
|
def close(self) -> None:
|
|
pass
|
|
|
|
def wait(self) -> bool:
|
|
return False
|
|
|
|
|
|
class MP4VVideoWriter(VideoWriter):
|
|
|
|
def __init__(self, video_writer: cv2.VideoWriter, thresholds: list[DetectedObjectBox]) -> None:
|
|
self._video_writer = video_writer
|
|
self._thresholds = thresholds
|
|
|
|
def write_frame(self, frame: VideoFrame) -> None:
|
|
self._video_writer.write(frame)
|
|
|
|
def plot_box_to_frame(self, data: DetectedObject, frame: VideoFrame) -> None:
|
|
confidence = float(data.confidence)
|
|
for threshold in self._thresholds:
|
|
if threshold.min_threshold < confidence:
|
|
# xmin, ymin, xmax, ymax = self._extract_box(data)
|
|
cv2.rectangle(frame, *data.as_int_box(), threshold.color, 2)
|
|
return
|
|
|
|
def plot_boxes_to_frame(self, boxes: list[DetectedObject], frame: VideoFrame) -> None:
|
|
for box in boxes:
|
|
self.plot_box_to_frame(box, frame)
|
|
|
|
|
|
def close(self) -> None:
|
|
self._video_writer.release()
|
|
cv2.destroyAllWindows()
|
|
|
|
def wait(self) -> bool:
|
|
return False
|
|
|
|
@classmethod
|
|
def from_frame(cls, input_: Path, output: Path, thresholds: list[float]) -> VideoWriter:
|
|
object_drawers = [DetectedObjectBox.from_viridis(t) for t in thresholds]
|
|
video_cap = cv2.VideoCapture(str(input_))
|
|
try:
|
|
return cls(create_video_writer(video_cap, output), object_drawers)
|
|
finally:
|
|
video_cap.release()
|
|
|
|
|
|
class MP4VVideoWriterWithFrame(MP4VVideoWriter):
|
|
def write_frame(self, frame: VideoFrame) -> None:
|
|
super().write_frame(frame)
|
|
cv2.imshow("Frame", frame)
|
|
|
|
def wait(self) -> bool:
|
|
return cv2.waitKey(1) == ord("q")
|