thinh-huynh-re commited on
Commit
e5b9cea
1 Parent(s): 70b3e07
Files changed (2) hide show
  1. run_opencv.py +11 -4
  2. utils/img_container.py +2 -2
run_opencv.py CHANGED
@@ -35,6 +35,10 @@ class ArgParser(Tap):
35
 
36
  id2label: Optional[str] = "labels/kinetics_400.json"
37
 
 
 
 
 
38
 
39
  class ActivityModel:
40
  def __init__(self, args: ArgParser):
@@ -77,9 +81,12 @@ class ActivityModel:
77
  max_index = logits.argmax(-1).item()
78
  predicted_label = self.model.config.id2label[max_index]
79
 
80
- img_container.frame_rate.label = (
81
- f"{predicted_label}_{logits[0][max_index]:.2f}%"
82
- )
 
 
 
83
 
84
  # logits = np.squeeze(logits)
85
  logits = logits.squeeze().numpy()
@@ -105,7 +112,7 @@ class ActivityModel:
105
 
106
  def main(args: ArgParser):
107
  activity_model = ActivityModel(args)
108
- img_container = ImgContainer(activity_model.frames_per_video)
109
 
110
  num_skips = 0
111
 
 
35
 
36
  id2label: Optional[str] = "labels/kinetics_400.json"
37
 
38
+ threshold: Optional[float] = 10.0
39
+
40
+ max_confidence: Optional[float] = 20.0 # Set None if not scale
41
+
42
 
43
  class ActivityModel:
44
  def __init__(self, args: ArgParser):
 
81
  max_index = logits.argmax(-1).item()
82
  predicted_label = self.model.config.id2label[max_index]
83
 
84
+ confidence = logits[0][max_index]
85
+
86
+ if (self.args.threshold is None) or (
87
+ self.args.threshold is not None and confidence >= self.args.threshold
88
+ ):
89
+ img_container.frame_rate.label = f"{predicted_label}_{confidence:.2f}%"
90
 
91
  # logits = np.squeeze(logits)
92
  logits = logits.squeeze().numpy()
 
112
 
113
  def main(args: ArgParser):
114
  activity_model = ActivityModel(args)
115
+ img_container = ImgContainer(activity_model.frames_per_video, args.is_recording)
116
 
117
  num_skips = 0
118
 
utils/img_container.py CHANGED
@@ -7,13 +7,13 @@ from .frame_rate import FrameRate
7
 
8
 
9
  class ImgContainer:
10
- def __init__(self, frames_per_video: int = 8) -> None:
11
  self.img: Optional[np.ndarray] = None # raw image
12
  self.frame_rate: FrameRate = FrameRate()
13
  self.imgs: List[np.ndarray] = []
14
  self.frames_per_video = frames_per_video
15
  self.rs: Optional[DataFrame] = None
16
- self.is_recording: bool = False
17
 
18
  def add_frame(self, frame: np.ndarray) -> None:
19
  if len(self.imgs) >= self.frames_per_video:
 
7
 
8
 
9
  class ImgContainer:
10
+ def __init__(self, frames_per_video: int = 8, is_recording: bool = False) -> None:
11
  self.img: Optional[np.ndarray] = None # raw image
12
  self.frame_rate: FrameRate = FrameRate()
13
  self.imgs: List[np.ndarray] = []
14
  self.frames_per_video = frames_per_video
15
  self.rs: Optional[DataFrame] = None
16
+ self.is_recording = is_recording
17
 
18
  def add_frame(self, frame: np.ndarray) -> None:
19
  if len(self.imgs) >= self.frames_per_video: