going somewhere.
This commit is contained in:
1 parent
2c798c8ca5
commit
35c5ab9a32
5 files changed
+51
-8
No files matched your search
@@ -160,3 +160,4 @@ env
|
||||
*.png
|
||||
*.ipynb
|
||||
.direnv
|
||||
result
|
||||
@@ -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")
|
||||
|
||||
@@ -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])
|
||||
|
||||
|
||||
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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 ];
|
||||
};
|
||||
|
||||
Reference in new issue
Block a user