From a8d74ff7500b7926c1a5844960b0d959514d101e Mon Sep 17 00:00:00 2001
From: Michal Kindl <kindl.michal@gmail.com>
Date: Wed, 16 Nov 2022 18:49:18 +0100
Subject: [PATCH] Added frame number to infer text result

---
 yolov6/core/inferer.py | 11 ++++++++---
 1 file changed, 8 insertions(+), 3 deletions(-)

diff --git a/yolov6/core/inferer.py b/yolov6/core/inferer.py
index b3601a7a..46328ea9 100644
--- a/yolov6/core/inferer.py
+++ b/yolov6/core/inferer.py
@@ -68,7 +68,8 @@ def infer(self, conf_thres, iou_thres, classes, agnostic_nms, max_det, save_dir,
         ''' Model Inference and results visualization '''
         vid_path, vid_writer, windows = None, None, []
         fps_calculator = CalcFPS()
-        for img_src, img_path, vid_cap in tqdm(self.files):
+        frame_num = 0
+        for img_src, img_path, vid_cap in tqdm(self.files):            
             img, img_src = self.precess_image(img_src, self.img_size, self.stride, self.half)
             img = img.to(self.device)
             if len(img.shape) == 3:
@@ -97,7 +98,7 @@ def infer(self, conf_thres, iou_thres, classes, agnostic_nms, max_det, save_dir,
                 for *xyxy, conf, cls in reversed(det):
                     if save_txt:  # Write to file
                         xywh = (self.box_convert(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist()  # normalized xywh
-                        line = (cls, *xywh, conf)
+                        line = (cls, *xywh, conf, frame_num)
                         with open(txt_path + '.txt', 'a') as f:
                             f.write(('%g ' * len(line)).rstrip() % line + '\n')
 
@@ -107,7 +108,7 @@ def infer(self, conf_thres, iou_thres, classes, agnostic_nms, max_det, save_dir,
 
                         self.plot_box_and_label(img_ori, max(round(sum(img_ori.shape) / 2 * 0.003), 2), xyxy, label, color=self.generate_colors(class_num, True))
 
-                img_src = np.asarray(img_ori)
+                img_src = np.asarray(img_ori)            
 
             # FPS counter
             fps_calculator.update(1.0 / (t2 - t1))
@@ -150,6 +151,10 @@ def infer(self, conf_thres, iou_thres, classes, agnostic_nms, max_det, save_dir,
                         save_path = str(Path(save_path).with_suffix('.mp4'))  # force *.mp4 suffix on results videos
                         vid_writer = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))
                     vid_writer.write(img_src)
+            
+            # increment frame number
+            frame_num += 1
+
 
     @staticmethod
     def precess_image(img_src, img_size, stride, half):