a single commit for general viz function to implement plot_results, r…

…emove plotting from predict functions, closes #761 and update docs

deepforest/main.py

@@ -341,13 +341,13 @@ def predict_image(self,
         Args:
             image: a float32 numpy array of a RGB with channels last format
             path: optional path to read image from disk instead of passing image arg
-            return_plot: Return image with plotted detections
             color: color of the bounding box as a tuple of BGR color, e.g. orange annotations is (0, 165, 255)
             thickness: thickness of the rectangle border line in px
         Returns:
             result: A pandas dataframe of predictions (Default)
             img: The input with predictions overlaid (Optional)
         """
+
         # Ensure we are in eval mode
         self.model.eval()
 
@@ -373,24 +373,20 @@ def predict_image(self,
         result = predict._predict_image_(model=self.model,
                                          image=image,
                                          path=path,
-                                         nms_thresh=self.config["nms_thresh"],
-                                         return_plot=return_plot,
-                                         thickness=thickness,
-                                         color=color)
+                                         nms_thresh=self.config["nms_thresh"])
 
-        if return_plot:
-            return result
+        #If there were no predictions, return None
+        if result is None:
+            return None
         else:
-            #If there were no predictions, return None
-            if result is None:
-                return None
-            else:
-                result["label"] = result.label.apply(
-                    lambda x: self.numeric_to_label_dict[x])
+            result["label"] = result.label.apply(
+                lambda x: self.numeric_to_label_dict[x])
+
+        result = utilities.read_file(result)
 
         return result
 
-    def predict_file(self, csv_file, root_dir, savedir=None, color=None, thickness=1):
+    def predict_file(self, csv_file, root_dir):
         """Create a dataset and predict entire annotation file Csv file format
         is .csv file with the columns "image_path", "xmin","ymin","xmax","ymax"
         for the image name and bounding box position. Image_path is the
@@ -400,12 +396,10 @@ def predict_file(self, csv_file, root_dir, savedir=None, color=None, thickness=1
         Args:
             csv_file: path to csv file
             root_dir: directory of images. If none, uses "image_dir" in config
-            savedir: Optional. Directory to save image plots.
-            color: color of the bounding box as a tuple of BGR color, e.g. orange annotations is (0, 165, 255)
-            thickness: thickness of the rectangle border line in px
         Returns:
             df: pandas dataframe with bounding boxes, label and scores for each image in the csv file
         """
+
         df = utilities.read_file(csv_file)
         ds = dataset.TreeDataset(csv_file=csv_file,
                                  root_dir=root_dir,
@@ -418,10 +412,7 @@ def predict_file(self, csv_file, root_dir, savedir=None, color=None, thickness=1
                                                annotations=df,
                                                dataloader=dataloader,
                                                root_dir=root_dir,
-                                               nms_thresh=self.config["nms_thresh"],
-                                               savedir=savedir,
-                                               color=color,
-                                               thickness=thickness)
+                                               nms_thresh=self.config["nms_thresh"])
 
         return results
 
@@ -431,12 +422,9 @@ def predict_tile(self,
                      patch_size=400,
                      patch_overlap=0.05,
                      iou_threshold=0.15,
-                     return_plot=False,
                      mosaic=True,
                      sigma=0.5,
                      thresh=0.001,
-                     color=None,
-                     thickness=1,
                      crop_model=None,
                      crop_transform=None,
                      crop_augment=False):
@@ -453,12 +441,9 @@ def predict_tile(self,
             iou_threshold: Minimum iou overlap among predictions between
                 windows to be suppressed.
                 Lower values suppress more boxes at edges.
-            return_plot: Should the image be returned with the predictions drawn?
             mosaic: Return a single prediction dataframe (True) or a tuple of image crops and predictions (False)
             sigma: variance of Gaussian function used in Gaussian Soft NMS
             thresh: the score thresh used to filter bboxes after soft-nms performed
-            color: color of the bounding box as a tuple of BGR color, e.g. orange annotations is (0, 165, 255)
-            thickness: thickness of the rectangle border line in px
             cropModel: a deepforest.model.CropModel object to predict on crops
             crop_transform: a torchvision.transforms object to apply to crops
             crop_augment: a boolean to apply augmentations to crops
@@ -510,18 +495,6 @@ def predict_tile(self,
                 lambda x: self.numeric_to_label_dict[x])
             if raster_path:
                 results["image_path"] = os.path.basename(raster_path)
-            if return_plot:
-                # Draw predictions on BGR
-                if raster_path:
-                    tile = rio.open(raster_path).read()
-                else:
-                    tile = self.image
-                drawn_plot = tile[:, :, ::-1]
-                drawn_plot = visualize.plot_predictions(tile,
-                                                        results,
-                                                        color=color,
-                                                        thickness=thickness)
-                return drawn_plot
         else:
             for df in results:
                 df["label"] = df.label.apply(lambda x: self.numeric_to_label_dict[x])
@@ -543,6 +516,8 @@ def predict_tile(self,
                                                    transform=crop_transform,
                                                    augment=crop_augment)
 
+        results = utilities.read_file(results)
+
         return results
 
     def training_step(self, batch, batch_idx):
@@ -732,30 +707,27 @@ def configure_optimizers(self):
         else:
             return optimizer
 
-    def evaluate(self, csv_file, root_dir, iou_threshold=None, savedir=None):
+    def evaluate(self, csv_file, root_dir, iou_threshold=None):
         """Compute intersection-over-union and precision/recall for a given
         iou_threshold.
 
         Args:
             csv_file: location of a csv file with columns "name","xmin","ymin","xmax","ymax","label", each box in a row
             root_dir: location of files in the dataframe 'name' column.
             iou_threshold: float [0,1] intersection-over-union union between annotation and prediction to be scored true positive
-            savedir: optional path dir to save evaluation images
         Returns:
             results: dict of ("results", "precision", "recall") for a given threshold
         """
         ground_df = utilities.read_file(csv_file)
         ground_df["label"] = ground_df.label.apply(lambda x: self.label_dict[x])
         predictions = self.predict_file(csv_file=csv_file,
-                                        root_dir=root_dir,
-                                        savedir=savedir)
+                                        root_dir=root_dir)
 
         results = evaluate_iou.__evaluate_wrapper__(
             predictions=predictions,
             ground_df=ground_df,
             root_dir=root_dir,
             iou_threshold=iou_threshold,
-            numeric_to_label_dict=self.numeric_to_label_dict,
-            savedir=savedir)
+            numeric_to_label_dict=self.numeric_to_label_dict)
 
         return results

deepforest/predict.py

@@ -186,13 +186,6 @@ def _dataloader_wrapper_(model,
 
     results = read_file(results, root_dir)
 
-    if savedir:
-        visualize.plot_prediction_dataframe(results,
-                                            root_dir=root_dir,
-                                            savedir=savedir,
-                                            color=color,
-                                            thickness=thickness)
-
     return results
 
 

deepforest/visualize.py

@@ -256,38 +256,200 @@ def label_to_color(label):
 
     return color_dict[label]
 
-
-def convert_to_sv_format(df):
+def convert_to_sv_format(df, width=None, height=None):
     """Convert DeepForest prediction results to a supervision Detections
     object.
 
     Args:
         df (pd.DataFrame): The results from `predict_image` or `predict_tile`.
                            Expected columns: ['xmin', 'ymin', 'xmax', 'ymax', 'label', 'score', 'image_path'].
+        width (int): The width of the image in pixels. Required if the geometry type is 'polygon'.
+        height (int): The height of the image in pixels. Required if the geometry type is 'polygon'.
 
     Returns:
         sv.Detections: A supervision Detections object containing bounding boxes, class IDs,
                        confidence scores, and class names object mapping classes ids to corresponding
                        class names inside of data dictionary.
 
+
     Example:
         detections = convert_to_sv_format(result)
     """
-    # Extract bounding boxes as a 2D numpy array with shape (_, 4)
-    boxes = df[['xmin', 'ymin', 'xmax', 'ymax']].values.astype(np.float32)
+    geom_type = determine_geometry_type(df)
+
+    if geom_type == "box":
+        # Extract bounding boxes as a 2D numpy array with shape (_, 4)
+        boxes = df.geometry.apply(lambda x: (x.bounds[0], x.bounds[1], x.bounds[2], x.bounds[3])).values
+        boxes = np.stack(boxes)
+
+        label_mapping = {label: idx for idx, label in enumerate(df['label'].unique())}
+
+        # Extract labels as a numpy array
+        labels = df['label'].map(label_mapping).values.astype(int)
+
+        # Extract scores as a numpy array
+        try:
+            scores = np.array(df['score'].tolist())
+        except KeyError:
+            scores = np.ones(len(labels))
+
+        # Create a reverse mapping from integer to string labels
+        class_name = {v: k for k, v in label_mapping.items()}
+
+        detections = sv.Detections(
+            xyxy=boxes,
+            class_id=labels,
+            confidence=scores,
+            data={"class_name": [class_name[class_id] for class_id in labels]})
+
+    elif geom_type == "polygon":
+        # Extract bounding boxes as a 2D numpy array with shape (_, 4)
+        boxes = df.geometry.apply(lambda x: (x.bounds[0], x.bounds[1], x.bounds[2], x.bounds[3])).values
+        boxes = np.stack(boxes)
+
+        label_mapping = {label: idx for idx, label in enumerate(df['label'].unique())}
+
+        # Extract labels as a numpy array
+        labels = df['label'].map(label_mapping).values.astype(int)
+
+        # Extract scores as a numpy array
+        scores = np.array(df['score'].tolist())
+        # Create a reverse mapping from integer to string labels
+        class_name = {v: k for k, v in label_mapping.items()}
+
+        # Create masks
+        if height is None or width is None:
+            raise ValueError("height and width of the mask must be provided for polygon predictions")
+
+        polygons = df.geometry.apply(lambda x: np.array(x.exterior.coords)).values
+        # as integers
+        polygons = [np.array(p).round().astype(np.int32) for p in polygons]
+        masks = [ sv.polygon_to_mask(p,(width,height)) for p in polygons ]
+        masks = np.stack(masks)
+
+        detections = sv.Detections(
+            xyxy=boxes,
+            mask=masks,
+            class_id=labels,
+            confidence=scores,
+            data={"class_name": [class_name[class_id] for class_id in labels]})
+
+    elif geom_type == "point":
+        points = df.geometry.apply(lambda x: (x.x, x.y)).values
+        points = np.stack(points)
+        points = np.expand_dims(points, axis=1)
+
+        label_mapping = {label: idx for idx, label in enumerate(df['label'].unique())}
+
+        # Extract labels as a numpy array
+        labels = df['label'].map(label_mapping).values.astype(int)
+
+        # Extract scores as a numpy array
+        scores = np.array(df['score'].tolist())
+        scores = np.expand_dims(np.stack(scores), 1)
+
+        # Create a reverse mapping from integer to string labels
+        class_name = {v: k for k, v in label_mapping.items()}
+
+        detections = sv.KeyPoints(
+            xy=points,
+            class_id=labels,
+            confidence=scores,
+            data={"class_name": [class_name[class_id] for class_id in labels
+            ]})
+
+    return detections
+
+def plot_results(results, root_dir, ground_truth=None, savedir=None, height=None, width=None, results_color=None, ground_truth_color=None, thickness=2, radius=3):
+    """Plot the prediction results.
+
+    Args:
+        df: a pandas dataframe with prediction results
+        root_dir: the root directory where the images are stored
+        ground_truth: an optional pandas dataframe with ground truth annotations
+        savedir: optional path to save the figure. If None (default), the figure will be interactively plotted.
+        height: height of the image in pixels. Required if the geometry type is 'polygon'.
+        width: width of the image in pixels. Required if the geometry type is 'polygon'.
+        results_color (list): color of the results annotations as a tuple of RGB color, e.g. orange annotations is [245, 135, 66]
+    Returns:
+        None
+    """
+    # Convert colors, check for multi-class labels
+    if results_color is None:
+        sv_color = sv.Color(245, 135, 66)
+    elif type(results_color) is list:
+        sv_color = sv.Color(results_color[0], results_color[1], results_color[2])
+    else:
+        sv_color = results_color
+
+    num_labels = len(results.label.unique())
+    if num_labels > 1 and results_color is not None:
+        warnings.warn("Multiple labels detected in the results. Each label will be plotted with a different color using a color ramp, results color argument is ignored.")
+    if num_labels > 1:
+        sv_color = sv.ColorPalette.from_matplotlib('viridis', 5)
+
+    # Read images
+    image_path = os.path.join(root_dir, results.image_path.unique()[0])
+    image = np.array(Image.open(image_path))
+
+    # Plot the results following https://supervision.roboflow.com/annotators/
+    fig, ax = plt.subplots()
+    annotated_scene = _plot_image_with_results(df=results,image=image, sv_color=sv_color, height=height, width=width, thickness=thickness, radius=radius)
+
+    if ground_truth is not None:
+        if ground_truth_color is None:
+            sv_color = sv.Color(0, 165, 255)
+        elif type(ground_truth_color) is list:
+            sv_color = sv.Color(ground_truth_color[0], ground_truth_color[1], ground_truth_color[2])
+        else:
+            sv_color = ground_truth_color
+        # Plot the ground truth annotations
+        annotated_scene = _plot_image_with_results(df=ground_truth,image=annotated_scene, sv_color=sv_color, height=height, width=width, thickness=thickness, radius=radius)
+
+    if savedir:
+        basename = os.path.splitext(os.path.basename(results.image_path.unique()[0]))[0]
+        image_name = "{}.png".format(basename)
+        plt.savefig(os.path.join(savedir, image_name), bbox_inches='tight', pad_inches=0)
+    else:
+        # Display the image using Matplotlib
+        plt.imshow(annotated_scene)
+        plt.axis('off')  # Hide axes for a cleaner look
+        plt.show()
+
+def _plot_image_with_results(df, image, sv_color, thickness=1, radius=3, height=None, width=None):  
+    """
+    Annotates an image with the given results.
 
-    label_mapping = {label: idx for idx, label in enumerate(df['label'].unique())}
+    Args:
+        df (pandas.DataFrame): The DataFrame containing the results.
+        image (numpy.ndarray): The image to annotate.
+        sv_color (str): The color of the annotations.
+        thickness (int): The thickness of the annotations.
 
-    # Extract labels as a numpy array
-    labels = df['label'].map(label_mapping).values.astype(int)
+    Returns:
+        numpy.ndarray: The annotated image.
+    """
+    # Determine the geometry type
+    geom_type = determine_geometry_type(df)
+    detections = convert_to_sv_format(df, height=height, width=width)
 
-    # Extract scores as a numpy array
-    scores = np.array(df['score'].tolist())
-    # Create a reverse mapping from integer to string labels
-    class_name = {v: k for k, v in label_mapping.items()}
+    if geom_type == "box":
+        bounding_box_annotator = sv.BoxAnnotator(color=sv_color, thickness=thickness)
+        annotated_frame = bounding_box_annotator.annotate(
+            scene=image.copy(),
+            detections=detections,
+        )
+    elif geom_type == "polygon":
 
-    return sv.Detections(
-        xyxy=boxes,
-        class_id=labels,
-        confidence=scores,
-        data={"class_name": [class_name[class_id] for class_id in labels]})
+        polygon_annotator = sv.PolygonAnnotator(color=sv_color, thickness=thickness)
+        annotated_frame = polygon_annotator.annotate(
+            scene=image.copy(),
+            detections=detections,
+        )
+    elif geom_type == "point":
+        point_annotator = sv.VertexAnnotator(color=sv_color, radius=radius)
+        annotated_frame = point_annotator.annotate(
+            scene=image.copy(),
+                    key_points=detections
+                    )
+    return annotated_frame