diff --git a/docs/en/guides/heatmaps.md b/docs/en/guides/heatmaps.md index 3b8b91116b40..1cb25a161994 100644 --- a/docs/en/guides/heatmaps.md +++ b/docs/en/guides/heatmaps.md @@ -65,7 +65,8 @@ A heatmap generated with [Ultralytics YOLOv8](https://github.com/ultralytics/ult imw=w, imh=h, view_img=True, - shape="circle") + shape="circle", + classes_names=model.names) while cap.isOpened(): success, im0 = cap.read() @@ -110,7 +111,8 @@ A heatmap generated with [Ultralytics YOLOv8](https://github.com/ultralytics/ult imh=h, view_img=True, shape="circle", - count_reg_pts=line_points) + count_reg_pts=line_points, + classes_names=model.names) while cap.isOpened(): success, im0 = cap.read() @@ -126,6 +128,51 @@ A heatmap generated with [Ultralytics YOLOv8](https://github.com/ultralytics/ult video_writer.release() cv2.destroyAllWindows() ``` + + === "Polygon Counting" + ```python + from ultralytics import YOLO + import heatmap + import cv2 + + model = YOLO("yolov8n.pt") + cap = cv2.VideoCapture("path/to/video/file.mp4") + assert cap.isOpened(), "Error reading video file" + w, h, fps = (int(cap.get(x)) for x in (cv2.CAP_PROP_FRAME_WIDTH, cv2.CAP_PROP_FRAME_HEIGHT, cv2.CAP_PROP_FPS)) + + # Video writer + video_writer = cv2.VideoWriter("heatmap_output.avi", + cv2.VideoWriter_fourcc(*'mp4v'), + fps, + (w, h)) + + # Define polygon points + region_points = [(20, 400), (1080, 404), (1080, 360), (20, 360), (20, 400)] + + # Init heatmap + heatmap_obj = heatmap.Heatmap() + heatmap_obj.set_args(colormap=cv2.COLORMAP_PARULA, + imw=w, + imh=h, + view_img=True, + shape="circle", + count_reg_pts=region_points, + classes_names=model.names) + + while cap.isOpened(): + success, im0 = cap.read() + if not success: + print("Video frame is empty or video processing has been successfully completed.") + break + tracks = model.track(im0, persist=True, show=False) + + im0 = heatmap_obj.generate_heatmap(im0, tracks) + video_writer.write(im0) + + cap.release() + video_writer.release() + cv2.destroyAllWindows() + ``` === "Region Counting" @@ -155,7 +202,8 @@ A heatmap generated with [Ultralytics YOLOv8](https://github.com/ultralytics/ult imh=h, view_img=True, shape="circle", - count_reg_pts=region_points) + count_reg_pts=region_points, + classes_names=model.names) while cap.isOpened(): success, im0 = cap.read() @@ -190,7 +238,8 @@ A heatmap generated with [Ultralytics YOLOv8](https://github.com/ultralytics/ult imw=w, imh=h, view_img=True, - shape="circle") + shape="circle", + classes_names=model.names) results = model.track(im0, persist=True) im0 = heatmap_obj.generate_heatmap(im0, tracks=results) @@ -223,7 +272,8 @@ A heatmap generated with [Ultralytics YOLOv8](https://github.com/ultralytics/ult imw=w, imh=h, view_img=True, - shape="circle") + shape="circle", + classes_names=model.names) while cap.isOpened(): success, im0 = cap.read() @@ -243,22 +293,27 @@ A heatmap generated with [Ultralytics YOLOv8](https://github.com/ultralytics/ult ### Arguments `set_args` -| Name | Type | Default | Description | -|-----------------------|----------------|-------------------|-----------------------------------------------------------| -| `view_img` | `bool` | `False` | Display the frame with heatmap | -| `colormap` | `cv2.COLORMAP` | `None` | cv2.COLORMAP for heatmap | -| `imw` | `int` | `None` | Width of Heatmap | -| `imh` | `int` | `None` | Height of Heatmap | -| `heatmap_alpha` | `float` | `0.5` | Heatmap alpha value | -| `count_reg_pts` | `list` | `None` | Object counting region points | -| `count_txt_thickness` | `int` | `2` | Count values text size | -| `count_txt_color` | `RGB Color` | `(0, 0, 0)` | Foreground color for Object counts text | -| `count_color` | `RGB Color` | `(255, 255, 255)` | Background color for Object counts text | -| `count_reg_color` | `RGB Color` | `(255, 0, 255)` | Counting region color | -| `region_thickness` | `int` | `5` | Counting region thickness value | -| `decay_factor` | `float` | `0.99` | Decay factor for heatmap area removal after specific time | -| `shape` | `str` | `circle` | Heatmap shape for display "rect" or "circle" supported | -| `line_dist_thresh` | `int` | `15` | Euclidean Distance threshold for line counter | +| Name | Type | Default | Description | +|-----------------------|----------------|---------------------|-----------------------------------------------------------| +| `view_img` | `bool` | `False` | Display the frame with heatmap | +| `colormap` | `cv2.COLORMAP` | `None` | cv2.COLORMAP for heatmap | +| `imw` | `int` | `None` | Width of Heatmap | +| `imh` | `int` | `None` | Height of Heatmap | +| `view_in_counts` | `bool` | `True` | Display in-counts only on video frame | +| `view_out_counts` | `bool` | `True` | Display out-counts only on video frame | +| `classes_names` | `dict` | `model.model.names` | Dictionary of Class Names | +| `heatmap_alpha` | `float` | `0.5` | Heatmap alpha value | +| `count_reg_pts` | `list` | `None` | Object counting region points | +| `count_txt_thickness` | `int` | `2` | Count values text size | +| `count_txt_color` | `RGB Color` | `(0, 0, 0)` | Foreground color for Object counts text | +| `count_reg_color` | `RGB Color` | `(255, 0, 255)` | Counting region color | +| `region_thickness` | `int` | `5` | Counting region thickness value | +| `decay_factor` | `float` | `0.99` | Decay factor for heatmap area removal after specific time | +| `shape` | `str` | `circle` | Heatmap shape for display "rect" or "circle" supported | +| `line_dist_thresh` | `int` | `15` | Euclidean Distance threshold for line counter | +| `fontsize` | `float` | `0.6` | Font size of counting text | +| `line_color` | `RGB Color` | `(255, 255, 255)` | Count highlighter color | +| `cls_txtdisplay_gap` | `int` | `50` | Display gap between each class count | ### Arguments `model.track` diff --git a/docs/en/guides/object-counting.md b/docs/en/guides/object-counting.md index 0c3ce72a3d46..8b5f148b5cbc 100644 --- a/docs/en/guides/object-counting.md +++ b/docs/en/guides/object-counting.md @@ -229,9 +229,11 @@ Object counting with [Ultralytics YOLOv8](https://github.com/ultralytics/ultraly | `track_color` | `RGB Color` | `(0, 255, 0)` | Color for each track line | | `line_dist_thresh` | `int` | `15` | Euclidean Distance threshold for line counter | | `count_txt_thickness` | `int` | `2` | Thickness of Object counts text | -| `count_txt_color` | `RGB Color` | `(0, 0, 0)` | Foreground color for Object counts text | -| `count_color` | `RGB Color` | `(255, 255, 255)` | Background color for Object counts text | +| `count_txt_color` | `RGB Color` | `(255, 255, 255)` | Foreground color for Object counts text | | `region_thickness` | `int` | `5` | Thickness for object counter region or line | +| `fontsize` | `float` | `0.6` | Font size of counting text | +| `line_color` | `RGB Color` | `(255, 255, 255)` | Count highlighter color | +| `cls_txtdisplay_gap` | `int` | `50` | Display gap between each class count | ### Arguments `model.track` diff --git a/ultralytics/solutions/heatmap.py b/ultralytics/solutions/heatmap.py index f70e62bb7288..6524abb06ddd 100644 --- a/ultralytics/solutions/heatmap.py +++ b/ultralytics/solutions/heatmap.py @@ -24,6 +24,8 @@ def __init__(self): self.view_img = False self.shape = "circle" + self.names = None # Classes names + # Image information self.imw = None self.imh = None @@ -52,10 +54,13 @@ def __init__(self): # Object Counting Information self.in_counts = 0 self.out_counts = 0 - self.counting_list = [] + self.count_ids = [] + self.class_wise_count = {} self.count_txt_thickness = 0 - self.count_txt_color = (0, 0, 0) - self.count_color = (255, 255, 255) + self.count_txt_color = (255, 255, 255) + self.line_color = (255, 255, 255) + self.cls_txtdisplay_gap = 50 + self.fontsize = 0.6 # Decay factor self.decay_factor = 0.99 @@ -67,6 +72,7 @@ def set_args( self, imw, imh, + classes_names=None, colormap=cv2.COLORMAP_JET, heatmap_alpha=0.5, view_img=False, @@ -74,13 +80,15 @@ def set_args( view_out_counts=True, count_reg_pts=None, count_txt_thickness=2, - count_txt_color=(0, 0, 0), - count_color=(255, 255, 255), + count_txt_color=(255, 255, 255), + fontsize=0.8, + line_color=(255, 255, 255), count_reg_color=(255, 0, 255), region_thickness=5, line_dist_thresh=15, decay_factor=0.99, shape="circle", + cls_txtdisplay_gap=50, ): """ Configures the heatmap colormap, width, height and display parameters. @@ -89,6 +97,7 @@ def set_args( colormap (cv2.COLORMAP): The colormap to be set. imw (int): The width of the frame. imh (int): The height of the frame. + classes_names (dict): Classes names heatmap_alpha (float): alpha value for heatmap display view_img (bool): Flag indicating frame display view_in_counts (bool): Flag to control whether to display the incounts on video stream. @@ -96,13 +105,16 @@ def set_args( count_reg_pts (list): Object counting region points count_txt_thickness (int): Text thickness for object counting display count_txt_color (RGB color): count text color value - count_color (RGB color): count text background color value + fontsize (float): Text display font size + line_color (RGB color): count highlighter line color count_reg_color (RGB color): Color of object counting region region_thickness (int): Object counting Region thickness line_dist_thresh (int): Euclidean Distance threshold for line counter decay_factor (float): value for removing heatmap area after object passed shape (str): Heatmap shape, rect or circle shape supported + cls_txtdisplay_gap (int): Display gap between each class count """ + self.names = classes_names self.imw = imw self.imh = imh self.heatmap_alpha = heatmap_alpha @@ -116,29 +128,29 @@ def set_args( if len(count_reg_pts) == 2: print("Line Counter Initiated.") self.count_reg_pts = count_reg_pts - self.counting_region = LineString(count_reg_pts) - - elif len(count_reg_pts) == 4: - print("Region Counter Initiated.") + self.counting_region = LineString(self.count_reg_pts) + elif len(count_reg_pts) >= 3: + print("Polygon Counter Initiated.") self.count_reg_pts = count_reg_pts self.counting_region = Polygon(self.count_reg_pts) - else: - print("Region or line points Invalid, 2 or 4 points supported") + print("Invalid Region points provided, region_points must be 2 for lines or >= 3 for polygons.") print("Using Line Counter Now") - self.counting_region = Polygon([(20, 400), (1260, 400)]) # dummy points + self.counting_region = LineString(self.count_reg_pts) # Heatmap new frame self.heatmap = np.zeros((int(self.imh), int(self.imw)), dtype=np.float32) self.count_txt_thickness = count_txt_thickness self.count_txt_color = count_txt_color - self.count_color = count_color + self.fontsize = fontsize + self.line_color = line_color self.region_color = count_reg_color self.region_thickness = region_thickness self.decay_factor = decay_factor self.line_dist_thresh = line_dist_thresh self.shape = shape + self.cls_txtdisplay_gap = cls_txtdisplay_gap # shape of heatmap, if not selected if self.shape not in ["circle", "rect"]: @@ -183,6 +195,12 @@ def generate_heatmap(self, im0, tracks): ) for box, cls, track_id in zip(self.boxes, self.clss, self.track_ids): + # Store class info + if self.names[cls] not in self.class_wise_count: + if len(self.names[cls]) > 5: + self.names[cls] = self.names[cls][:5] + self.class_wise_count[self.names[cls]] = {"in": 0, "out": 0} + if self.shape == "circle": center = (int((box[0] + box[2]) // 2), int((box[1] + box[3]) // 2)) radius = min(int(box[2]) - int(box[0]), int(box[3]) - int(box[1])) // 2 @@ -203,23 +221,39 @@ def generate_heatmap(self, im0, tracks): if len(track_line) > 30: track_line.pop(0) - # Count objects - if len(self.count_reg_pts) == 4: - if self.counting_region.contains(Point(track_line[-1])) and track_id not in self.counting_list: - self.counting_list.append(track_id) - if box[0] < self.counting_region.centroid.x: - self.out_counts += 1 - else: + prev_position = self.track_history[track_id][-2] if len(self.track_history[track_id]) > 1 else None + + # Count objects in any polygon + if len(self.count_reg_pts) >= 3: + is_inside = self.counting_region.contains(Point(track_line[-1])) + + if prev_position is not None and is_inside and track_id not in self.count_ids: + self.count_ids.append(track_id) + + if (box[0] - prev_position[0]) * (self.counting_region.centroid.x - prev_position[0]) > 0: self.in_counts += 1 + self.class_wise_count[self.names[cls]]["in"] += 1 + else: + self.out_counts += 1 + self.class_wise_count[self.names[cls]]["out"] += 1 + # Count objects using line elif len(self.count_reg_pts) == 2: - distance = Point(track_line[-1]).distance(self.counting_region) - if distance < self.line_dist_thresh and track_id not in self.counting_list: - self.counting_list.append(track_id) - if box[0] < self.counting_region.centroid.x: - self.out_counts += 1 - else: - self.in_counts += 1 + is_inside = (box[0] - prev_position[0]) * (self.counting_region.centroid.x - prev_position[0]) > 0 + + if prev_position is not None and is_inside and track_id not in self.count_ids: + distance = Point(track_line[-1]).distance(self.counting_region) + + if distance < self.line_dist_thresh and track_id not in self.count_ids: + self.count_ids.append(track_id) + + if (box[0] - prev_position[0]) * (self.counting_region.centroid.x - prev_position[0]) > 0: + self.in_counts += 1 + self.class_wise_count[self.names[cls]]["in"] += 1 + else: + self.out_counts += 1 + self.class_wise_count[self.names[cls]]["out"] += 1 + else: for box, cls in zip(self.boxes, self.clss): if self.shape == "circle": @@ -240,26 +274,30 @@ def generate_heatmap(self, im0, tracks): heatmap_normalized = cv2.normalize(self.heatmap, None, 0, 255, cv2.NORM_MINMAX) heatmap_colored = cv2.applyColorMap(heatmap_normalized.astype(np.uint8), self.colormap) - incount_label = f"In Count : {self.in_counts}" - outcount_label = f"OutCount : {self.out_counts}" - - # Display counts based on user choice - counts_label = None - if not self.view_in_counts and not self.view_out_counts: - counts_label = None - elif not self.view_in_counts: - counts_label = outcount_label - elif not self.view_out_counts: - counts_label = incount_label - else: - counts_label = f"{incount_label} {outcount_label}" + label = "Ultralytics Analytics \t" + + for key, value in self.class_wise_count.items(): + if value["in"] != 0 or value["out"] != 0: + if not self.view_in_counts and not self.view_out_counts: + label = None + elif not self.view_in_counts: + label += f"{str.capitalize(key)}: IN {value['in']} \t" + elif not self.view_out_counts: + label += f"{str.capitalize(key)}: OUT {value['out']} \t" + else: + label += f"{str.capitalize(key)}: IN {value['in']} OUT {value['out']} \t" + + label = label.rstrip() + label = label.split("\t") - if self.count_reg_pts is not None and counts_label is not None: - self.annotator.count_labels( - counts=counts_label, - count_txt_size=self.count_txt_thickness, + if self.count_reg_pts is not None and label is not None: + self.annotator.display_counts( + counts=label, + tf=self.count_txt_thickness, + fontScale=self.fontsize, txt_color=self.count_txt_color, - color=self.count_color, + line_color=self.line_color, + classwise_txtgap=self.cls_txtdisplay_gap, ) self.im0 = cv2.addWeighted(self.im0, 1 - self.heatmap_alpha, heatmap_colored, self.heatmap_alpha, 0) diff --git a/ultralytics/solutions/object_counter.py b/ultralytics/solutions/object_counter.py index 18f42c624936..a3fd0906f80c 100644 --- a/ultralytics/solutions/object_counter.py +++ b/ultralytics/solutions/object_counter.py @@ -43,16 +43,19 @@ def __init__(self): # Object counting Information self.in_counts = 0 self.out_counts = 0 - self.counting_dict = {} + self.count_ids = [] + self.class_wise_count = {} self.count_txt_thickness = 0 - self.count_txt_color = (0, 0, 0) - self.count_color = (255, 255, 255) + self.count_txt_color = (255, 255, 255) + self.line_color = (255, 255, 255) + self.cls_txtdisplay_gap = 50 + self.fontsize = 0.6 # Tracks info self.track_history = defaultdict(list) self.track_thickness = 2 self.draw_tracks = False - self.track_color = (0, 255, 0) + self.track_color = None # Check if environment support imshow self.env_check = check_imshow(warn=True) @@ -68,12 +71,14 @@ def set_args( view_in_counts=True, view_out_counts=True, draw_tracks=False, - count_txt_thickness=2, - count_txt_color=(0, 0, 0), - count_color=(255, 255, 255), - track_color=(0, 255, 0), + count_txt_thickness=3, + count_txt_color=(255, 255, 255), + fontsize=0.8, + line_color=(255, 255, 255), + track_color=None, region_thickness=5, line_dist_thresh=15, + cls_txtdisplay_gap=50, ): """ Configures the Counter's image, bounding box line thickness, and counting region points. @@ -89,11 +94,13 @@ def set_args( draw_tracks (Bool): draw tracks count_txt_thickness (int): Text thickness for object counting display count_txt_color (RGB color): count text color value - count_color (RGB color): count text background color value + fontsize (float): Text display font size + line_color (RGB color): count highlighter line color count_reg_color (RGB color): Color of object counting region track_color (RGB color): color for tracks region_thickness (int): Object counting Region thickness line_dist_thresh (int): Euclidean Distance threshold for line counter + cls_txtdisplay_gap (int): Display gap between each class count """ self.tf = line_thickness self.view_img = view_img @@ -108,7 +115,7 @@ def set_args( self.reg_pts = reg_pts self.counting_region = LineString(self.reg_pts) elif len(reg_pts) >= 3: - print("Region Counter Initiated.") + print("Polygon Counter Initiated.") self.reg_pts = reg_pts self.counting_region = Polygon(self.reg_pts) else: @@ -120,10 +127,12 @@ def set_args( self.track_color = track_color self.count_txt_thickness = count_txt_thickness self.count_txt_color = count_txt_color - self.count_color = count_color + self.fontsize = fontsize + self.line_color = line_color self.region_color = count_reg_color self.region_thickness = region_thickness self.line_dist_thresh = line_dist_thresh + self.cls_txtdisplay_gap = cls_txtdisplay_gap def mouse_event_for_region(self, event, x, y, flags, params): """ @@ -171,7 +180,13 @@ def extract_and_process_tracks(self, tracks): # Extract tracks for box, track_id, cls in zip(boxes, track_ids, clss): # Draw bounding box - self.annotator.box_label(box, label=f"{track_id}:{self.names[cls]}", color=colors(int(track_id), True)) + self.annotator.box_label(box, label=f"{self.names[cls]}#{track_id}", color=colors(int(track_id), True)) + + # Store class info + if self.names[cls] not in self.class_wise_count: + if len(self.names[cls]) > 5: + self.names[cls] = self.names[cls][:5] + self.class_wise_count[self.names[cls]] = {"in": 0, "out": 0} # Draw Tracks track_line = self.track_history[track_id] @@ -182,68 +197,68 @@ def extract_and_process_tracks(self, tracks): # Draw track trails if self.draw_tracks: self.annotator.draw_centroid_and_tracks( - track_line, color=self.track_color, track_thickness=self.track_thickness + track_line, + color=self.track_color if self.track_color else colors(int(track_id), True), + track_thickness=self.track_thickness, ) prev_position = self.track_history[track_id][-2] if len(self.track_history[track_id]) > 1 else None - centroid = Point((box[:2] + box[2:]) / 2) - - # Count objects - if len(self.reg_pts) >= 3: # any polygon - is_inside = self.counting_region.contains(centroid) - current_position = "in" if is_inside else "out" - if prev_position is not None: - if self.counting_dict[track_id] != current_position and is_inside: - self.in_counts += 1 - self.counting_dict[track_id] = "in" - elif self.counting_dict[track_id] != current_position and not is_inside: - self.out_counts += 1 - self.counting_dict[track_id] = "out" - else: - self.counting_dict[track_id] = current_position + # Count objects in any polygon + if len(self.reg_pts) >= 3: + is_inside = self.counting_region.contains(Point(track_line[-1])) - else: - self.counting_dict[track_id] = current_position + if prev_position is not None and is_inside and track_id not in self.count_ids: + self.count_ids.append(track_id) - elif len(self.reg_pts) == 2: - if prev_position is not None: - is_inside = (box[0] - prev_position[0]) * ( - self.counting_region.centroid.x - prev_position[0] - ) > 0 - current_position = "in" if is_inside else "out" - - if self.counting_dict[track_id] != current_position and is_inside: + if (box[0] - prev_position[0]) * (self.counting_region.centroid.x - prev_position[0]) > 0: self.in_counts += 1 - self.counting_dict[track_id] = "in" - elif self.counting_dict[track_id] != current_position and not is_inside: - self.out_counts += 1 - self.counting_dict[track_id] = "out" + self.class_wise_count[self.names[cls]]["in"] += 1 else: - self.counting_dict[track_id] = current_position - else: - self.counting_dict[track_id] = None - - incount_label = f"In Count : {self.in_counts}" - outcount_label = f"OutCount : {self.out_counts}" - - # Display counts based on user choice - counts_label = None - if not self.view_in_counts and not self.view_out_counts: - counts_label = None - elif not self.view_in_counts: - counts_label = outcount_label - elif not self.view_out_counts: - counts_label = incount_label - else: - counts_label = f"{incount_label} {outcount_label}" + self.out_counts += 1 + self.class_wise_count[self.names[cls]]["out"] += 1 - if counts_label is not None: - self.annotator.count_labels( - counts=counts_label, - count_txt_size=self.count_txt_thickness, + # Count objects using line + elif len(self.reg_pts) == 2: + is_inside = (box[0] - prev_position[0]) * (self.counting_region.centroid.x - prev_position[0]) > 0 + + if prev_position is not None and is_inside and track_id not in self.count_ids: + distance = Point(track_line[-1]).distance(self.counting_region) + + if distance < self.line_dist_thresh and track_id not in self.count_ids: + self.count_ids.append(track_id) + + if (box[0] - prev_position[0]) * (self.counting_region.centroid.x - prev_position[0]) > 0: + self.in_counts += 1 + self.class_wise_count[self.names[cls]]["in"] += 1 + else: + self.out_counts += 1 + self.class_wise_count[self.names[cls]]["out"] += 1 + + label = "Ultralytics Analytics \t" + + for key, value in self.class_wise_count.items(): + if value["in"] != 0 or value["out"] != 0: + if not self.view_in_counts and not self.view_out_counts: + label = None + elif not self.view_in_counts: + label += f"{str.capitalize(key)}: IN {value['in']} \t" + elif not self.view_out_counts: + label += f"{str.capitalize(key)}: OUT {value['out']} \t" + else: + label += f"{str.capitalize(key)}: IN {value['in']} OUT {value['out']} \t" + + label = label.rstrip() + label = label.split("\t") + + if label is not None: + self.annotator.display_counts( + counts=label, + tf=self.count_txt_thickness, + fontScale=self.fontsize, txt_color=self.count_txt_color, - color=self.count_color, + line_color=self.line_color, + classwise_txtgap=self.cls_txtdisplay_gap, ) def display_frames(self): diff --git a/ultralytics/utils/plotting.py b/ultralytics/utils/plotting.py index d0215ba5e3ea..04d305bcef55 100644 --- a/ultralytics/utils/plotting.py +++ b/ultralytics/utils/plotting.py @@ -363,35 +363,49 @@ def draw_centroid_and_tracks(self, track, color=(255, 0, 255), track_thickness=2 cv2.polylines(self.im, [points], isClosed=False, color=color, thickness=track_thickness) cv2.circle(self.im, (int(track[-1][0]), int(track[-1][1])), track_thickness * 2, color, -1) - def count_labels(self, counts=0, count_txt_size=2, color=(255, 255, 255), txt_color=(0, 0, 0)): + def display_counts( + self, counts=None, tf=2, fontScale=0.6, line_color=(0, 0, 0), txt_color=(255, 255, 255), classwise_txtgap=55 + ): """ - Plot counts for object counter. + Display counts on im0. Args: - counts (int): objects counts value - count_txt_size (int): text size for counts display - color (tuple): background color of counts display - txt_color (tuple): text color of counts display + counts (str): objects count data + tf (int): text thickness for display + fontScale (float): text fontsize for display + line_color (RGB Color): counts highlighter color + txt_color (RGB Color): counts display color + classwise_txtgap (int): Gap between each class count data """ - self.tf = count_txt_size - tl = self.tf or round(0.002 * (self.im.shape[0] + self.im.shape[1]) / 2) + 1 + + tl = tf or round(0.002 * (self.im.shape[0] + self.im.shape[1]) / 2) + 1 tf = max(tl - 1, 1) - # Get text size for in_count and out_count - t_size_in = cv2.getTextSize(str(counts), 0, fontScale=tl / 2, thickness=tf)[0] + t_sizes = [cv2.getTextSize(str(count), 0, fontScale=0.8, thickness=tf)[0] for count in counts] - # Calculate positions for counts label - text_width = t_size_in[0] - text_x = (self.im.shape[1] - text_width) // 2 # Center x-coordinate - text_y = t_size_in[1] + max_text_width = max([size[0] for size in t_sizes]) + max_text_height = max([size[1] for size in t_sizes]) - # Create a rounded rectangle for in_count - cv2.rectangle( - self.im, (text_x - 5, text_y - 5), (text_x + text_width + 7, text_y + t_size_in[1] + 7), color, -1 - ) - cv2.putText( - self.im, str(counts), (text_x, text_y + t_size_in[1]), 0, tl / 2, txt_color, self.tf, lineType=cv2.LINE_AA - ) + text_x = self.im.shape[1] - max_text_width - 20 + text_y = classwise_txtgap + + for i, count in enumerate(counts): + text_x_pos = text_x + text_y_pos = text_y + i * classwise_txtgap + + cv2.putText( + self.im, + str(count), + (text_x_pos, text_y_pos), + cv2.FONT_HERSHEY_SIMPLEX, + fontScale=fontScale, + color=txt_color, + thickness=tf, + lineType=cv2.LINE_AA, + ) + + line_y_pos = text_y_pos + max_text_height + 5 + cv2.line(self.im, (text_x_pos, line_y_pos), (text_x_pos + max_text_width, line_y_pos), line_color, tf) @staticmethod def estimate_pose_angle(a, b, c):