Update detect_frames.py

scripts/detect_frames.py

@@ -38,27 +38,26 @@ def main(args, config={}, verbose=True):
 
     print(config)
 
-    dirname = args.frames
-
     model, device = setup_model(args.weights)
     
     locations = [
-        "elwha-val"
+        "kenai-val"
     ]
     for loc in locations:
 
-        in_loc_dir = os.path.join(dirname, loc)
+        in_loc_dir = os.path.join(args.frames, loc)
         out_loc_dir = os.path.join(args.output, loc)
         print(in_loc_dir)
         print(out_loc_dir)
 
-        seq_list = os.listdir(in_loc_dir)
+        detect_location(in_loc_dir, out_loc_dir, config, model, device, verbose)
 
-        detect_location(in_loc_dir, out_loc_dir, config, model, device, seq_list, verbose)
 
                 
 def detect_location(in_loc_dir, out_loc_dir, config, model, device, seq_list, verbose):
 
+    seq_list = os.listdir(in_loc_dir)
+
     with tqdm(total=len(seq_list), desc="...", ncols=0) as pbar:
         for seq in seq_list:
 
@@ -118,7 +117,7 @@ def detect(in_dir, config, model, device, verbose):
             confs = pred[:, 4].clone().tolist()
             scale_boxes(image_shape, boxes, original_shape[0], original_shape[1])  # to original shape
 
-            frame = [ [*bb, conf, file_name] for bb, conf, file_name in zip(boxes.tolist(), confs, file_names[batch_i*32+si]) ]
+            frame = [ [*bb, conf] for bb, conf in zip(boxes.tolist(), confs) ]
             
             file_name = file_names[batch_i*32 + si]
             for ann in frame: