From bd776b0a14ab58974e382ff21f7db9ecd635d67d Mon Sep 17 00:00:00 2001
From: didier amyot <didier.amyot@all-dressed-programming.com>
Date: Sat, 2 Nov 2024 07:28:44 -0400
Subject: [PATCH] modidied stuff

---
 dodgeballtracker/logic.py | 22 ++++++++++++++++------
 flake.nix                 | 36 ++++++++++++++++++++++++++----------
 2 files changed, 42 insertions(+), 16 deletions(-)

diff --git a/dodgeballtracker/logic.py b/dodgeballtracker/logic.py
index e1e9b50..2477fae 100644
--- a/dodgeballtracker/logic.py
+++ b/dodgeballtracker/logic.py
@@ -1,6 +1,6 @@
 import json
 from pathlib import Path
-
+import logging
 import torch
 from ultralytics import YOLO
 import cv2
@@ -10,10 +10,22 @@ import numpy as np
 from boxmot import DeepOcSort as DeepOCSORT
 
 
+logger = logging.getLogger(__name__)
 
 Color = tuple[int, int, int]
 
 
+def get_best_available_torch_device():
+    if torch.cuda.is_available():
+        logger.debug("Using the cuda (gpu) torch device.")
+        return 'cuda'
+    if torch.backends.mps.is_available():
+        logger.debug("Using the mps (apple silicon) torch device.")
+        return 'mps'
+    logger.debug("Default to cpu (no apple silicon or cuda) torch device.")
+    return 'cpu'
+
+
 class DetectedObjectBox:
     def __init__(self, min_threshold: float, color: Color):
         self.min_threshold = min_threshold
@@ -93,7 +105,7 @@ class DodgeballTracker:
         for d in tracker_data:
             self.print_box(d, frame_idx, frame_ts)
 
-    def plot_boxe_to_frame(self, data, frame):
+    def plot_box_to_frame(self, data, frame):
         confidence = float(data[4])
         for detected_box_object in self.detected_box_objects:
 
@@ -126,7 +138,7 @@ class DodgeballTracker:
     def _get_tracker(self):
         return DeepOCSORT(
             Path('lmbn_n_cuhk03_d.pt'),  # which ReID model to use
-            torch.device("cpu"),
+            torch.device(get_best_available_torch_device()),
             True,
         )
 
@@ -151,10 +163,8 @@ class DodgeballTracker:
             # run the YOLO model on the frame
             res = []
             detections = self.run_model_on_frame(model, frame)
-            print("ici?")
             for data in detections.boxes.data.tolist():
-                print("la?")
-                is_good_box = self.plot_boxe_to_frame(data, frame)
+                is_good_box = self.plot_box_to_frame(data, frame)
                 if is_good_box:
                     res.append(data)
 
diff --git a/flake.nix b/flake.nix
index f279941..807b09c 100644
--- a/flake.nix
+++ b/flake.nix
@@ -2,24 +2,40 @@
   description = "Pre-built CPU-only OpenCV packages for Python.  ";
   inputs = {
     nixpkgs.url = "github:NixOS/nixpkgs/nixpkgs-unstable";
-    opencv-python-related.url = "gitlab:all-dressed-programming/opencv-python-related";
+    opencv-python-related.url =
+      "gitlab:all-dressed-programming/opencv-python-related";
+    opencv-python-related.home-manager.inputs.nixpkgs.follows = "nixpkgs";
     flake-utils.url = "github:numtide/flake-utils";
   };
 
+  nixConfig = {
+    extra-substituters = [ "https://cuda-maintainers.cachix.org" ];
+    extra-trusted-public-keys = [
+      "cuda-maintainers.cachix.org-1:0dq3bujKpuEPMCX6U4WylrUDZ9JyUG0VpVZa7CNfq5E="
+    ];
+  };
+
   outputs = { self, nixpkgs, flake-utils, opencv-python-related }:
     flake-utils.lib.eachDefaultSystem (system:
       let
-        pkgs = nixpkgs.legacyPackages.${system};
-        pythonPackages = pkgs.python312Packages // opencv-python-related.packages.${system}.python312Packages;
-        virtual-env = pkgs.python312.withPackages (ps: [ ps.matplotlib
-                                                         pythonPackages.opencv4
-                                                         pythonPackages.boxmot
-                                                         pythonPackages.ultralytics ]);
+        pkgs = import nixpkgs {
+          system = system;
+          config.allowUnfree = true;
+        };
+        pythonPackages = pkgs.python312Packages
+          // opencv-python-related.packages.${system}.python312Packages;
+        virtual-env = pkgs.python312.withPackages (ps: [
+          ps.matplotlib
+          ps.torchWithCuda
+          # pythonPackages.opencv4
+          # pythonPackages.boxmot
+          # pythonPackages.ultralytics
+        ]);
 
       in {
-        packages = {
-          env = virtual-env;
+        packages = { env = virtual-env; };
+        devShells.default = pkgs.mkShell {
+          buildInputs = [ virtual-env pkgs.yt-dlp pkgs.ffmpeg_7 ];
         };
-        devShells.default = pkgs.mkShell { buildInputs = [ virtual-env pkgs.yt-dlp pkgs.ffmpeg_7]; };
       });
 }