Source code for deepforest.visualize

# Visualize module for plotting and handling predictions
import os
import pandas as pd
import matplotlib
from matplotlib import pyplot as plt
from PIL import Image
import numpy as np
import pandas.api.types as ptypes
import cv2
import random
import warnings
import supervision as sv
import shapely
from deepforest.utilities import determine_geometry_type


[docs] def view_dataset(ds, savedir=None, color=None, thickness=1): """Plot annotations on images for debugging purposes. Args: ds: a deepforest pytorch dataset, see deepforest.dataset or deepforest.load_dataset() to start from a csv file savedir: optional path to save figures. If none (default) images will be interactively plotted 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 """ for i in iter(ds): image_path, image, targets = i df = format_boxes(targets[0], scores=False) image = np.moveaxis(image[0].numpy(), 0, 2) image = plot_predictions(image, df, color=color, thickness=thickness) if savedir: cv2.imwrite("{}/{}".format(savedir, image_path[0]), image) else: plt.imshow(image)
[docs] def format_geometry(predictions, scores=True): """Format a retinanet prediction into a pandas dataframe for a batch of images Args: predictions: a list of dictionaries with keys 'boxes' and 'labels' coming from a retinanet scores: Whether boxes come with scores, during prediction, or without scores, as in during training. Returns: df: a pandas dataframe """ # Detect geometry type geom_type = determine_geometry_type(predictions) if geom_type == "box": df = format_boxes(predictions, scores=scores) df['geometry'] = df.apply( lambda x: shapely.geometry.box(x.xmin, x.ymin, x.xmax, x.ymax), axis=1) elif geom_type == "polygon": raise ValueError("Polygon predictions are not yet supported for formatting") elif geom_type == "point": raise ValueError("Point predictions are not yet supported for formatting") return df
[docs] def format_boxes(prediction, scores=True): """Format a retinanet prediction into a pandas dataframe for a single image. Args: prediction: a dictionary with keys 'boxes' and 'labels' coming from a retinanet scores: Whether boxes come with scores, during prediction, or without scores, as in during training. Returns: df: a pandas dataframe """ df = pd.DataFrame(prediction["boxes"].cpu().detach().numpy(), columns=["xmin", "ymin", "xmax", "ymax"]) df["label"] = prediction["labels"].cpu().detach().numpy() if scores: df["score"] = prediction["scores"].cpu().detach().numpy() return df
[docs] def plot_prediction_and_targets(image, predictions, targets, image_name, savedir): """Plot an image, its predictions, and its ground truth targets for debugging. Args: image: torch tensor, RGB color order targets: torch tensor Returns: figure_path: path on disk with saved figure """ image = np.array(image)[:, :, ::-1].copy() prediction_df = format_boxes(predictions) image = plot_predictions(image, prediction_df) target_df = format_boxes(targets, scores=False) image = plot_predictions(image, target_df) figure_path = "{}/{}.png".format(savedir, image_name) cv2.imwrite(figure_path, image) return figure_path
[docs] def plot_prediction_dataframe(df, root_dir, savedir, color=None, thickness=1, ground_truth=None): """For each row in dataframe, call plot predictions and save plot files to disk. For multi-class labels, boxes will be colored by labels. Ground truth boxes will all be same color, regardless of class. Args: df: a pandas dataframe with image_path, xmin, xmax, ymin, ymax and label columns. The image_path column should be the relative path from root_dir, not the full path. root_dir: relative dir to look for image names from df.image_path ground_truth: an optional pandas dataframe in same format as df holding ground_truth boxes savedir: save the plot to an optional directory path. Returns: written_figures: list of filenames written """ written_figures = [] for name, group in df.groupby("image_path"): image = np.array(Image.open("{}/{}".format(root_dir, name)))[:, :, ::-1].copy() image = plot_predictions(image, group) if ground_truth is not None: annotations = ground_truth[ground_truth.image_path == name] image = plot_predictions(image, annotations, color=color, thickness=thickness) figure_name = "{}/{}.png".format(savedir, os.path.splitext(name)[0]) written_figures.append(figure_name) cv2.imwrite(figure_name, image) return written_figures
[docs] def plot_points(image, points, color=None, radius=5, thickness=1): """Plot points on an image Args: image: a numpy array in *BGR* color order! Channel order is channels first points: a numpy array of shape (N, 2) representing the coordinates of the points color: color of the points as a tuple of BGR color, e.g. orange points is (0, 165, 255) radius: radius of the points in px thickness: thickness of the point border line in px Returns: image: a numpy array with drawn points """ if image.shape[0] == 3: warnings.warn("Input images must be channels last format [h, w, 3] not channels " "first [3, h, w], using np.rollaxis(image, 0, 3) to invert!") image = np.rollaxis(image, 0, 3) if image.dtype == "float32": image = image.astype("uint8") image = image.copy() if not color: color = (0, 165, 255) # Default color is orange for point in points: cv2.circle(image, (int(point[0]), int(point[1])), color=color, radius=radius, thickness=thickness) return image
[docs] def plot_predictions(image, df, color=None, thickness=1): """Plot a set of boxes on an image By default this function does not show, but only plots an axis Label column must be numeric! Image must be BGR color order! Args: image: a numpy array in *BGR* color order! Channel order is channels first df: a pandas dataframe with xmin, xmax, ymin, ymax and label column 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: image: a numpy array with drawn annotations """ if image.shape[0] == 3: warnings.warn("Input images must be channels last format [h, w, 3] not channels " "first [3, h, w], using np.rollaxis(image, 0, 3) to invert!") image = np.rollaxis(image, 0, 3) if image.dtype == "float32": image = image.astype("uint8") image = image.copy() if not color: if not ptypes.is_numeric_dtype(df.label): warnings.warn("No color was provided and the label column is not numeric. " "Using a single default color.") color = (0, 165, 255) for index, row in df.iterrows(): if not color: color = label_to_color(row["label"]) # Plot rectangle, points or polygon if "geometry" in df.columns: # Get geometry type geometry_type = row["geometry"].geom_type if geometry_type == "Polygon": # convert to int32 and numpy array int_coords = lambda x: np.array(x).round().astype(np.int32) polygon = [int_coords(row["geometry"].exterior.coords)] cv2.polylines(image, polygon, True, color, thickness=thickness) elif geometry_type == "Point": int_coords = lambda x: np.array(x).round().astype(np.int32) cv2.circle(image, (int_coords(row["geometry"].x), int_coords(row["geometry"].y)), color=color, radius=5, thickness=thickness) else: raise ValueError("Only polygons and points are supported") elif "xmin" in df.columns: cv2.rectangle(image, (int(row["xmin"]), int(row["ymin"])), (int(row["xmax"]), int(row["ymax"])), color=color, thickness=thickness, lineType=cv2.LINE_AA) elif "x" in df.columns: cv2.circle(image, (row["x"], row["y"]), color=color, radius=5, thickness=thickness) elif "polygon" in df.columns: polygon = np.array(row["polygon"]) cv2.polylines(image, [polygon], True, color, thickness=thickness) return image
[docs] def label_to_color(label): color_dict = {} random.seed(1) colors = [ list((matplotlib.colors.hsv_to_rgb([x, 1.0, 1.0]) * 255).astype(int)) for x in np.arange(0, 1, 1 / 80) ] colors = [tuple([int(y) for y in x]) for x in colors] random.shuffle(colors) for index, color in enumerate(colors): color_dict[index] = color # hand pick the first few colors color_dict[0] = (255, 255, 0) color_dict[1] = (71, 99, 255) color_dict[2] = (255, 0, 0) color_dict[3] = (50, 205, 50) color_dict[4] = (214, 112, 214) color_dict[5] = (60, 20, 220) color_dict[6] = (63, 133, 205) color_dict[7] = (255, 144, 30) color_dict[8] = (0, 215, 255) return color_dict[label]
[docs] 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 includes: ['geometry', 'label', 'score', 'image_path'] for bounding boxes width (int): The width of the image in pixels. Only required if the geometry type is 'polygon'. height (int): The height of the image in pixels. Only required if the geometry type is 'polygon'. Returns: Depending on the geometry type, the function returns either a Detections or a KeyPoints object from the supervision library. """ 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 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()} # 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 __check_color__(color, num_labels): if isinstance(color, list) and len(color) == 3: if num_labels > 1: warnings.warn( """Multiple labels detected in the results and results_color argument provides a single color. Each label will be plotted with a different color using a built-in color ramp. If you want to customize colors with multiple labels pass a supervision.ColorPalette object to results_color with the appropriate number of labels""" ) return sv.ColorPalette.from_matplotlib('viridis', num_labels) else: return sv.Color(color[0], color[1], color[2]) elif isinstance(color, sv.draw.color.ColorPalette): if num_labels > len(color.colors): warnings.warn( """The number of colors provided in results_color does not match the number number of labels. Replacing the provided color palette with a built-in built-in color palette. To use a custom color palette make sure the number of values matches the number of labels""" ) return sv.ColorPalette.from_matplotlib('viridis', num_labels) else: return color elif isinstance(color, list): raise ValueError( "results_color must be either a 3 item list containing RGB values or an sv.ColorPalette instance" ) else: raise TypeError( "results_color must be either a list of RGB values or an sv.ColorPalette instance" )
[docs] def plot_annotations(annotations, savedir=None, height=None, width=None, color=[245, 135, 66], thickness=2, basename=None, root_dir=None, radius=3): """Plot the prediction results. Args: annotations: a pandas dataframe with prediction results 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 or sv.ColorPalette): color of the results annotations as a tuple of RGB color (if a single color), e.g. orange annotations is [245, 135, 66], or an supervision.ColorPalette if multiple labels and specifying colors for each label thickness: thickness of the rectangle border line in px basename: optional basename for the saved figure. If None (default), the basename will be extracted from the image path. root_dir: optional path to the root directory of the images. If None (default), the root directory will be extracted from the annotations dataframe.root_dir attribute. radius: radius of the points in px Returns: None """ # Convert colors, check for multi-class labels num_labels = len(annotations.label.unique()) annotation_color = __check_color__(color, num_labels) # Read images if not hasattr(annotations, 'root_dir'): if root_dir is None: raise ValueError( "The 'annotations.root_dir' attribute does not exist. Please specify the 'root_dir' argument." ) else: root_dir = root_dir else: root_dir = annotations.root_dir image_path = os.path.join(root_dir, annotations.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_geometry(df=annotations, image=image, sv_color=annotation_color, height=height, width=width, thickness=thickness, radius=radius) if savedir: if basename is None: basename = os.path.splitext( os.path.basename(annotations.image_path.unique()[0]))[0] image_name = "{}.png".format(basename) image_path = os.path.join(savedir, image_name) cv2.imwrite(image_path, annotated_scene) else: # Display the image using Matplotlib plt.imshow(annotated_scene) plt.axis('off') # Hide axes for a cleaner look plt.show()
[docs] def plot_results(results, ground_truth=None, savedir=None, height=None, width=None, results_color=[245, 135, 66], ground_truth_color=[0, 165, 255], thickness=2, basename=None, radius=3): """Plot the prediction results. Args: results: a pandas dataframe with prediction results 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 or sv.ColorPalette): color of the results annotations as a tuple of RGB color (if a single color), e.g. orange annotations is [245, 135, 66], or an supervision.ColorPalette if multiple labels and specifying colors for each label ground_truth_color (list): color of the ground truth annotations as a tuple of RGB color, e.g. blue annotations is [0, 165, 255] thickness: thickness of the rectangle border line in px basename: optional basename for the saved figure. If None (default), the basename will be extracted from the image path. radius: radius of the points in px Returns: None """ # Convert colors, check for multi-class labels num_labels = len(results.label.unique()) results_color_sv = __check_color__(results_color, num_labels) ground_truth_color_sv = __check_color__(ground_truth_color, num_labels) # Read images root_dir = results.root_dir 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_geometry(df=results, image=image, sv_color=results_color_sv, height=height, width=width, thickness=thickness, radius=radius) if ground_truth is not None: # Plot the ground truth annotations annotated_scene = _plot_image_with_geometry(df=ground_truth, image=annotated_scene, sv_color=ground_truth_color_sv, height=height, width=width, thickness=thickness, radius=radius) if savedir: if basename is None: basename = os.path.splitext(os.path.basename( results.image_path.unique()[0]))[0] image_name = "{}.png".format(basename) image_path = os.path.join(savedir, image_name) cv2.imwrite(image_path, annotated_scene) 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_geometry(df, image, sv_color, thickness=1, radius=3, height=None, width=None): """Annotates an image with the given results. 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. 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) 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": 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