|
|
|
@ -60,9 +60,9 @@ class Heatmap: |
|
|
|
|
self.heatmap_alpha = heatmap_alpha |
|
|
|
|
|
|
|
|
|
# Predict/track information |
|
|
|
|
self.boxes = None |
|
|
|
|
self.track_ids = None |
|
|
|
|
self.clss = None |
|
|
|
|
self.boxes = [] |
|
|
|
|
self.track_ids = [] |
|
|
|
|
self.clss = [] |
|
|
|
|
self.track_history = defaultdict(list) |
|
|
|
|
|
|
|
|
|
# Region & Line Information |
|
|
|
@ -107,16 +107,17 @@ class Heatmap: |
|
|
|
|
print("Using Circular shape now") |
|
|
|
|
self.shape = "circle" |
|
|
|
|
|
|
|
|
|
def extract_results(self, tracks, _intialized=False): |
|
|
|
|
def extract_results(self, tracks): |
|
|
|
|
""" |
|
|
|
|
Extracts results from the provided data. |
|
|
|
|
|
|
|
|
|
Args: |
|
|
|
|
tracks (list): List of tracks obtained from the object tracking process. |
|
|
|
|
""" |
|
|
|
|
self.boxes = tracks[0].boxes.xyxy.cpu() |
|
|
|
|
self.clss = tracks[0].boxes.cls.cpu().tolist() |
|
|
|
|
self.track_ids = tracks[0].boxes.id.int().cpu().tolist() |
|
|
|
|
if tracks[0].boxes.id is not None: |
|
|
|
|
self.boxes = tracks[0].boxes.xyxy.cpu() |
|
|
|
|
self.clss = tracks[0].boxes.cls.tolist() |
|
|
|
|
self.track_ids = tracks[0].boxes.id.int().tolist() |
|
|
|
|
|
|
|
|
|
def generate_heatmap(self, im0, tracks): |
|
|
|
|
""" |
|
|
|
@ -138,7 +139,7 @@ class Heatmap: |
|
|
|
|
self.extract_results(tracks) |
|
|
|
|
self.annotator = Annotator(self.im0, self.tf, None) |
|
|
|
|
|
|
|
|
|
if self.track_ids is not None: |
|
|
|
|
if self.track_ids: |
|
|
|
|
# Draw counting region |
|
|
|
|
if self.count_reg_pts is not None: |
|
|
|
|
self.annotator.draw_region( |
|
|
|
|