henry000 commited on
Commit
a0976c9
·
1 Parent(s): 70a7f92

:speech_balloon: [Add] FPS when using stream data

Browse files
Files changed (1) hide show
  1. yolo/tools/solver.py +6 -0
yolo/tools/solver.py CHANGED
@@ -1,4 +1,5 @@
1
  import os
 
2
 
3
  import torch
4
  from loguru import logger
@@ -120,6 +121,8 @@ class ModelTester:
120
  if dataloader.is_stream:
121
  import cv2
122
  import numpy as np
 
 
123
  try:
124
  for idx, images in enumerate(dataloader):
125
  images = images.to(self.device)
@@ -131,6 +134,9 @@ class ModelTester:
131
 
132
  if dataloader.is_stream:
133
  img = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)
 
 
 
134
  cv2.imshow("Prediction", img)
135
  if cv2.waitKey(1) & 0xFF == ord("q"):
136
  break
 
1
  import os
2
+ import time
3
 
4
  import torch
5
  from loguru import logger
 
121
  if dataloader.is_stream:
122
  import cv2
123
  import numpy as np
124
+
125
+ last_time = time.time()
126
  try:
127
  for idx, images in enumerate(dataloader):
128
  images = images.to(self.device)
 
134
 
135
  if dataloader.is_stream:
136
  img = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)
137
+ fps = 1 / (time.time() - last_time)
138
+ cv2.putText(img, f"FPS: {fps:.2f}", (0, 15), 0, 0.5, (100, 255, 0), 1, cv2.LINE_AA)
139
+ last_time = time.time()
140
  cv2.imshow("Prediction", img)
141
  if cv2.waitKey(1) & 0xFF == ord("q"):
142
  break