diff --git a/.gitignore b/.gitignore index b46071c..a56370a 100644 --- a/.gitignore +++ b/.gitignore @@ -160,3 +160,4 @@ env *.png *.ipynb .direnv +result diff --git a/dodgeballtracker/logic.py b/dodgeballtracker/logic.py index ca5e11d..b40b1e5 100644 --- a/dodgeballtracker/logic.py +++ b/dodgeballtracker/logic.py @@ -107,9 +107,12 @@ class DodgeballTracker: def run_model_on_frame(self, model, frame: VideoFrame): model_outputs = model(frame)[0] + return model_outputs.boxes.data.cpu().numpy() + # return model_outputs.boxes.data.tolist() + foo = model_outputs.boxes.data for data in model_outputs.boxes.data.tolist(): detected_obj = DetectedObject(*data) - if detected_obj.confidence > 0.5: + if detected_obj.score > 0.5: yield detected_obj def persist_points(self): @@ -130,7 +133,7 @@ class DodgeballTracker: return YOLO("yolo11n.pt", verbose=True) def update_tracker(self, tracker, detected_objs: list[DetectedObject], frame: VideoFrame) -> None: - important_detected_objs = [do for do in detected_objs if do.confidence > 0.5] + important_detected_objs = [do for do in detected_objs if do.score > 0.5] def main(self): # define some constants @@ -141,7 +144,11 @@ class DodgeballTracker: for i, (frame, ts) in enumerate(iter_frames(video_cap)): # run the YOLO model on the frame detected_objs = list(self.run_model_on_frame(model, frame)) - self.video_writer.plot_boxes_to_frame(detected_objs, frame) + data_ = tracker.update( + np.array(detected_objs), frame) + if i == 100: + print(data_) + # self.video_writer.plot_boxes_to_frame(detected_objs, frame) # data_ = tracker.update(np.array([dobj.as_nparray() for dobj in detected_objs]), frame) # if data_: # print("dkld") diff --git a/dodgeballtracker/typing.py b/dodgeballtracker/typing.py index 334802f..d8612e5 100644 --- a/dodgeballtracker/typing.py +++ b/dodgeballtracker/typing.py @@ -4,7 +4,7 @@ import numpy as np XCoord = NewType('XCoord', float) YCoord = NewType('YCoord', float) -Confidence = NewType('Confidence', float) +Score = NewType('Confidence', float) Category = NewType('Category', float) @@ -23,14 +23,14 @@ class DetectedObject: ymin: YCoord xmax: XCoord ymax: YCoord - confidence: Confidence + score: Score category: Category def as_int_box(self) -> tuple[tuple[int, int], tuple[int, int]]: return (int(self.xmin), int(self.ymin)), (int(self.xmax), int(self.ymax)) def as_nparray(self): - return np.array([self.xmin, self.ymax, self.xmax, self.ymax, self.confidence, self.category]) + return np.array([self.xmin, self.ymax, self.xmax, self.ymax, self.score, self.category]) diff --git a/dodgeballtracker/video_writer.py b/dodgeballtracker/video_writer.py index 293e8f6..0c0273f 100644 --- a/dodgeballtracker/video_writer.py +++ b/dodgeballtracker/video_writer.py @@ -81,7 +81,7 @@ class MP4VVideoWriter(VideoWriter): self._video_writer.write(frame) def plot_box_to_frame(self, data: DetectedObject, frame: VideoFrame) -> None: - confidence = float(data.confidence) + confidence = float(data.score) for threshold in self._thresholds: if threshold.min_threshold < confidence: # xmin, ymin, xmax, ymax = self._extract_box(data) diff --git a/flake.nix b/flake.nix index 359e704..5d92b2a 100644 --- a/flake.nix +++ b/flake.nix @@ -22,8 +22,42 @@ system = system; config.allowUnfree = true; }; + pythonPackages = pkgs.python312Packages // opencv-python-related.packages.${system}.python312Packages; + + pandera = pkgs.python3Packages.buildPythonPackage { + version = "0.5.11"; + pname = "lapx"; + pyproject = true; + src = pkgs.python3Packages.fetchPypi { + pname = "pandera"; + version = "0.20.4"; + hash = "sha256-zPYXgpPvnUOT3Ed25HR349LdUcgA7N/t7Gf/9Q9K08I="; + + }; + + propagatedBuildInputs = with pythonPackages; [ multimethod + numpy + pandas + pydantic + typeguard + typing-inspect + wrapt + ]; + + nativeBuildInputs = with pythonPackages; [ + setuptools + ]; + + meta = with pkgs.lib; { + description = "The Open-source Framework for Precision Data Testing"; + homepage = "https://github.com/unionai-oss/pandera"; + license = pkgs.lib.licenses.mit; + maintainers = with maintainers; [ ]; + }; + }; + virtual-env = pkgs.python312.withPackages (ps: [ ps.matplotlib ps.torchWithCuda @@ -32,10 +66,11 @@ pythonPackages.boxmot pythonPackages.ultralytics pythonPackages.click + pandera ]); in { - packages = { env = virtual-env; }; + packages = { env = virtual-env; pandera = pandera; }; devShells.default = pkgs.mkShell { buildInputs = [ virtual-env pkgs.yt-dlp pkgs.ffmpeg_7 ]; };