diff --git a/sahi/predict.py b/sahi/predict.py index 8863eb367..0cca03dd2 100644 --- a/sahi/predict.py +++ b/sahi/predict.py @@ -360,6 +360,8 @@ def predict( visual_bbox_thickness: int = None, visual_text_size: float = None, visual_text_thickness: int = None, + visual_hide_labels: bool = False, + visual_hide_conf: bool = False, visual_export_format: str = "png", verbose: int = 1, return_dict: bool = False, @@ -437,6 +439,8 @@ def predict( visual_bbox_thickness: int visual_text_size: float visual_text_thickness: int + visual_hide_labels: bool + visual_hide_conf: bool visual_export_format: str Can be specified as 'jpg' or 'png' verbose: int @@ -607,6 +611,8 @@ def predict( text_size=visual_text_size, text_th=visual_text_thickness, color=color, + hide_labels=visual_hide_labels, + hide_conf=visual_hide_conf, output_dir=None, file_name=None, export_format=None, @@ -619,6 +625,8 @@ def predict( text_size=visual_text_size, text_th=visual_text_thickness, color=color, + hide_labels=visual_hide_labels, + hide_conf=visual_hide_conf, output_dir=output_dir, file_name=filename_without_extension, export_format=visual_export_format, @@ -649,6 +657,8 @@ def predict( rect_th=visual_bbox_thickness, text_size=visual_text_size, text_th=visual_text_thickness, + hide_labels=visual_hide_labels, + hide_conf=visual_hide_conf, output_dir=output_dir if not source_is_video else None, file_name=filename_without_extension, export_format=visual_export_format, diff --git a/sahi/prediction.py b/sahi/prediction.py index d3faf46eb..7475d15f7 100644 --- a/sahi/prediction.py +++ b/sahi/prediction.py @@ -165,7 +165,13 @@ def __init__( self.durations_in_seconds = durations_in_seconds def export_visuals( - self, export_dir: str, text_size: float = None, rect_th: int = None, file_name: str = "prediction_visual" + self, + export_dir: str, + text_size: float = None, + rect_th: int = None, + hide_labels: bool = False, + hide_conf: bool = False, + file_name: str = "prediction_visual", ): """ @@ -173,6 +179,8 @@ def export_visuals( export_dir: directory for resulting visualization to be exported text_size: size of the category name over box rect_th: rectangle thickness + hide_labels: hide labels + hide_conf: hide confidence file_name: saving name Returns: @@ -185,6 +193,8 @@ def export_visuals( text_size=text_size, text_th=None, color=None, + hide_labels=hide_labels, + hide_conf=hide_conf, output_dir=export_dir, file_name=file_name, export_format="png", diff --git a/sahi/utils/cv.py b/sahi/utils/cv.py index a804b0b2c..0204df498 100644 --- a/sahi/utils/cv.py +++ b/sahi/utils/cv.py @@ -320,6 +320,7 @@ def visualize_prediction( text_size: float = None, text_th: float = None, color: tuple = None, + hide_labels: bool = False, output_dir: Optional[str] = None, file_name: Optional[str] = "prediction_visual", ): @@ -370,22 +371,24 @@ def visualize_prediction( color=color, thickness=rect_th, ) - # arange bounding box text location - label = f"{class_}" - w, h = cv2.getTextSize(label, 0, fontScale=text_size, thickness=text_th)[0] # label width, height - outside = p1[1] - h - 3 >= 0 # label fits outside box - p2 = p1[0] + w, p1[1] - h - 3 if outside else p1[1] + h + 3 - # add bounding box text - cv2.rectangle(image, p1, p2, color, -1, cv2.LINE_AA) # filled - cv2.putText( - image, - label, - (p1[0], p1[1] - 2 if outside else p1[1] + h + 2), - 0, - text_size, - (255, 255, 255), - thickness=text_th, - ) + + if not hide_labels: + # arange bounding box text location + label = f"{class_}" + w, h = cv2.getTextSize(label, 0, fontScale=text_size, thickness=text_th)[0] # label width, height + outside = p1[1] - h - 3 >= 0 # label fits outside box + p2 = p1[0] + w, p1[1] - h - 3 if outside else p1[1] + h + 3 + # add bounding box text + cv2.rectangle(image, p1, p2, color, -1, cv2.LINE_AA) # filled + cv2.putText( + image, + label, + (p1[0], p1[1] - 2 if outside else p1[1] + h + 2), + 0, + text_size, + (255, 255, 255), + thickness=text_th, + ) if output_dir: # create output folder if not present Path(output_dir).mkdir(parents=True, exist_ok=True) @@ -404,6 +407,8 @@ def visualize_object_predictions( text_size: float = None, text_th: float = None, color: tuple = None, + hide_labels: bool = False, + hide_conf: bool = False, output_dir: Optional[str] = None, file_name: str = "prediction_visual", export_format: str = "png", @@ -417,6 +422,8 @@ def visualize_object_predictions( text_size: size of the category name over box text_th: text thickness color: annotation color in the form: (0, 255, 0) + hide_labels: hide labels + hide_conf: hide confidence output_dir: directory for resulting visualization to be exported file_name: exported file will be saved as: output_dir+file_name+".png" export_format: can be specified as 'jpg' or 'png' @@ -473,22 +480,28 @@ def visualize_object_predictions( color=color, thickness=rect_th, ) - # arange bounding box text location - label = f"{category_name} {score:.2f}" - w, h = cv2.getTextSize(label, 0, fontScale=text_size, thickness=text_th)[0] # label width, height - outside = p1[1] - h - 3 >= 0 # label fits outside box - p2 = p1[0] + w, p1[1] - h - 3 if outside else p1[1] + h + 3 - # add bounding box text - cv2.rectangle(image, p1, p2, color, -1, cv2.LINE_AA) # filled - cv2.putText( - image, - label, - (p1[0], p1[1] - 2 if outside else p1[1] + h + 2), - 0, - text_size, - (255, 255, 255), - thickness=text_th, - ) + + if not hide_labels: + # arange bounding box text location + label = f"{category_name}" + + if not hide_conf: + label += f" {score:.2f}" + + w, h = cv2.getTextSize(label, 0, fontScale=text_size, thickness=text_th)[0] # label width, height + outside = p1[1] - h - 3 >= 0 # label fits outside box + p2 = p1[0] + w, p1[1] - h - 3 if outside else p1[1] + h + 3 + # add bounding box text + cv2.rectangle(image, p1, p2, color, -1, cv2.LINE_AA) # filled + cv2.putText( + image, + label, + (p1[0], p1[1] - 2 if outside else p1[1] + h + 2), + 0, + text_size, + (255, 255, 255), + thickness=text_th, + ) # export if output_dir is present if output_dir is not None: