Updated detection

dataloader.py

@@ -138,7 +138,7 @@ class YOLOFrameDataset(Dataset):
         return img, (h0, w0), (h1, w1)  # img, hw_original, hw_resized
 
     def __len__(self):
-        return len(self.batches)
+        return len(self.batch_indices)
 
     def __iter__(self):
         for batch_idx in self.batch_indices:

scripts/detect_frames.py

@@ -39,6 +39,8 @@ def main(args, config={}, verbose=True):
     print(config)
 
     dirname = args.frames
+
+    model, device = setup_model(args.weights)
     
     locations = [
         "elwha-val"
@@ -46,84 +48,56 @@ def main(args, config={}, verbose=True):
     for loc in locations:
 
         in_loc_dir = os.path.join(dirname, loc)
-        out_dir = os.path.join(args.output, loc, "tracker", "data")
+        out_loc_dir = os.path.join(args.output, loc)
         print(in_loc_dir)
-        print(out_dir)
-
-        # run detection + tracking
+        print(out_loc_dir)
 
         seq_list = os.listdir(in_loc_dir)
 
-        iterate_sequences(in_loc_dir, out_dir, config, args.weights, seq_list, verbose)
+        detect_location(in_loc_dir, out_loc_dir, config, model, device, seq_list, verbose)
 
                 
-def iterate_sequences(in_dir, out_dir, config, weights, seq_list, verbose):
-    model, device = setup_model(weights)
+def detect_location(in_loc_dir, out_loc_dir, config, model, device, seq_list, verbose):
 
-    idx = 1
     with tqdm(total=len(seq_list), desc="...", ncols=0) as pbar:
-            for seq in seq_list:
-                ann_list = []
-                pbar.update(1)
-                if (seq.startswith(".")): continue
-                pbar.set_description("Processing " + seq)
-                if verbose:
-                    print(" ")
-                    print("(" + str(idx) + "/" + str(len(seq_list)) + ") " + seq)
-                    print(" ")
-                idx += 1
-                in_seq_dir = os.path.join(in_dir, seq)
-                frame_list = detect(in_seq_dir, out_dir, config, seq, model, device, verbose)
-                for frame in frame_list:
-                    if frame is not None:
-                        for ann in frame:
-                            ann_list.append({
-                                'image_id': ann[5],
-                                'category_id': 0,
-                                'bbox': [ann[0], ann[1], ann[2] - ann[0], ann[3] - ann[1]],
-                                'score': ann[4]
-                            })
-                result = json.dumps(ann_list)
-
-                out_seq_dir = os.path.join(out_dir, seq)
-                os.makedirs(out_seq_dir, exist_ok=True)
-                with open(os.path.join(out_seq_dir, 'pred.json'), 'w') as f:
-                    f.write(result)
-
-def iterate_files(in_dir, out_dir, config, weights, verbose):
-    model, device = setup_model(weights)
+        for seq in seq_list:
 
-    ann_list = []
-    frame_list = detect(in_dir, out_dir, config, "specified folder", model, device, verbose)
-    with tqdm(total=len(frame_list), desc="...", ncols=0) as pbar:
-        for frame in frame_list:
-            if frame is not None:
-                for ann in frame:
-                    ann_list.append({
-                        'image_id': ann[5],
-                        'category_id': 0,
-                        'bbox': [ann[0], ann[1], ann[2] - ann[0], ann[3] - ann[1]],
-                        'score': ann[4]
-                    })
             pbar.update(1)
-        result = json.dumps(ann_list)
-        with open(os.path.join(out_dir, 'pred.json'), 'w') as f:
-            f.write(result)
+            if (seq.startswith(".")): continue
+            pbar.set_description("Processing " + seq)
+
+            in_seq_dir = os.path.join(in_loc_dir, seq)
+            out_seq_dir = os.path.join(out_loc_dir, seq)
+            os.makedirs(out_seq_dir, exist_ok=True)
 
-def detect(in_dir, out_dir, config, seq_name, model, device, verbose):
+            detect_seq(in_seq_dir, out_seq_dir, config, model, device, verbose)
+
+def detect_seq(in_seq_dir, out_seq_dir, config, model, device, verbose):
+
+    ann_list = []
+    frame_list = detect(in_seq_dir, config, model, device, verbose)
+    for frame in frame_list:
+        if frame is not None:
+            for ann in frame:
+                ann_list.append({
+                    'image_id': ann[5],
+                    'category_id': 0,
+                    'bbox': [ann[0], ann[1], ann[2] - ann[0], ann[3] - ann[1]],
+                    'score': ann[4]
+                })
+    result = json.dumps(ann_list)
+
+    with open(os.path.join(out_seq_dir, 'pred.json'), 'w') as f:
+        f.write(result)
+
+def detect(in_dir, config, model, device, verbose):
     
     #progress_log = lambda p, m: 0
 
     # create dataloader
     dataloader = create_dataloader_frames_only(in_dir)
 
-    try:
-        inference, image_shapes, width, height = do_detection(dataloader, model, device, verbose=verbose)
-    except:
-        print("Error in " + seq_name)
-        with open(os.path.join(out_dir, "ERROR_" + seq_name + ".txt"), 'w') as f:
-            f.write("ERROR")
-        return
+    inference, image_shapes, width, height = do_detection(dataloader, model, device, verbose=verbose)
 
 
     outputs = do_suppression(inference, conf_thres=config['conf_threshold'], iou_thres=config['nms_iou'], verbose=verbose)