Merge pull request #68 from conservationtechlab/dev

Update paths

src/animl/__main__.py

@@ -67,7 +67,7 @@ def main(image_dir, detector_file, classifier_file, class_list, sort=True):
         detections = file_management.load_data(working_dir.detections)
     else:
         detector = megadetector.MegaDetector(detector_file, device=device)
-        md_results = detect.detect_MD_batch(detector, all_frames["Frame"],
+        md_results = detect.detect_MD_batch(detector, all_frames, file_col="Frame",
                                             checkpoint_path=working_dir.mdraw, quiet=True)
         # Convert MD JSON to pandas dataframe, merge with manifest
         print("Converting MD JSON to dataframe and merging with manifest...")
@@ -83,7 +83,7 @@ def main(image_dir, detector_file, classifier_file, class_list, sort=True):
     print("Predicting species of animal detections...")
     print(class_list)
     classifier, classes = classifiers.load_model(classifier_file, class_list, device=device)
-    animals = inference.predict_species(animals.reset_index(drop=True), classifier, classes, device=device,
+    animals = inference.predict_species(animals, classifier, classes, device=device,
                                         file_col="Frame", batch_size=4, out_file=working_dir.predictions)
 
     # merge animal and empty, create symlinks

src/animl/classifiers.py

@@ -38,7 +38,7 @@ def save_model(out_dir, epoch, model, stats):
     torch.save(stats, open(f'{out_dir}/{epoch}.pt', 'wb'))
 
 
-def load_model(model_path, class_file, device="cpu", architecture="CTL"):
+def load_model(model_path, class_file, device=None, architecture="CTL"):
     '''
     Creates a model instance and loads the latest model state weights.
 
@@ -58,11 +58,14 @@ def load_model(model_path, class_file, device="cpu", architecture="CTL"):
     classes = pd.read_csv(Path(r""+class_file))
 
     # check to make sure GPU is available if chosen
-    if device != 'cpu' and not torch.cuda.is_available():
-        print(f'WARNING: device set to "{device}" but is not available; falling back to CPU...')
+    if not torch.cuda.is_available():
         device = 'cpu'
+    elif torch.cuda.is_available() and device is None:
+        device = 'cuda:0'
     else:
-        print('Device set to', device)
+        device = device
+
+    print('Device set to', device)
 
     # load latest model state from given folder
     if os.path.isdir(model_path):

src/animl/file_management.py

@@ -76,10 +76,10 @@ def __init__(self, working_dir):
         if not working_dir.is_dir():
             raise FileNotFoundError("The given directory does not exist.")
 
-        self.basedir = working_dir / "Animl-Directory/"
-        self.datadir = self.basedir / "Data/"
-        self.vidfdir = self.basedir / "Frames/"
-        self.linkdir = self.basedir / "Sorted/"
+        self.basedir = working_dir / Path("Animl-Directory/")
+        self.datadir = self.basedir / Path("Data/")
+        self.vidfdir = self.basedir / Path("Frames/")
+        self.linkdir = self.basedir / Path("Sorted/")
 
         # Create directories if they do not already exist
         self.basedir.mkdir(exist_ok=True)

src/animl/inference.py

@@ -56,6 +56,8 @@ def predict_species(detections, model, classes, device='cpu', out_file=None,
 
     if isinstance(detections, pd.DataFrame):
         # initialize lists
+
+        detections = detections.reset_index(drop=True)
         predictions = []
         probabilities = []
         if raw:

src/animl/models/yolo.py

@@ -147,7 +147,7 @@ def _forward_augment(self, x):
         y = self._clip_augmented(y)  # clip augmented tails
         return torch.cat(y, 1), None  # augmented inference, train
 
-    def _forward_once(self, x, profile=False):
+    def _forward_once(self, x, profile=False, visualize=False):
         y, dt = [], []  # outputs
         for m in self.model:
             if m.f != -1:  # if not from previous layer

src/animl/video_processing.py

@@ -37,9 +37,9 @@ def extract_frame_single(file_path, out_dir, fps=None, frames=None):
             if not ret:
                 break
 
-            out_path = (str(out_dir) + filename + "-" +
-                        uniqueid + "-" +
-                        str(frame_capture) + '.jpg')
+            frame_name = filename + "-" + uniqueid + "-" + \
+                        str(frame_capture) + '.jpg'
+            out_path = os.path.join(str(out_dir),frame_name)
             cv2.imwrite(out_path, frame)
             frames_saved.append([out_path, file_path])
             frame_capture += increment
@@ -51,10 +51,9 @@ def extract_frame_single(file_path, out_dir, fps=None, frames=None):
             cap.set(cv2.CAP_PROP_POS_MSEC, (frame_capture * 1000))
             if not ret:
                 break
-
-            out_path = (str(out_dir) + filename + "-" +
-                        '{:05}'.format(randrange(1, 10 ** 5)) + "-" +
-                        str(frame_capture) + '.jpg')
+            frame_name = filename + "-" + '{:05}'.format(randrange(1, 10 ** 5)) + \
+                         "-" + str(frame_capture) + '.jpg'
+            out_path = os.path.join(str(out_dir),frame_name)
             cv2.imwrite(out_path, frame)
             frames_saved.append([out_path, file_path])
             frame_capture += fps