def predict_file(model, csv_file, root_dir, savedir, device, iou_threshold=0.1): """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 relative filename, not absolute path, which is in the root_dir directory. One bounding box per line. 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. device: pytorch device of 'cuda' or 'cpu' for gpu prediction. Set internally. Returns: df: pandas dataframe with bounding boxes, label and scores for each image in the csv file """ input_csv = pd.read_csv(csv_file) # Just predict each image once. images = input_csv["image_path"].unique() prediction_list = [] for path in images: image = io.imread("{}/{}".format(root_dir, path)) image = preprocess.preprocess_image(image) # Just predict the images, even though we have the annotations if not device.type == "cpu": image = image.to(device) prediction = model(image) prediction = visualize.format_boxes(prediction[0]) prediction = across_class_nms(prediction, iou_threshold=iou_threshold) prediction["image_path"] = path prediction_list.append(prediction) if savedir: #if on GPU, bring back to cpu for plotting # Just predict the images, even though we have the annotations if not device.type == "cpu": image = image.to("cpu") image = image.squeeze(0).permute(1, 2, 0) plot, ax = visualize.plot_predictions(image, prediction) annotations = input_csv[input_csv.image_path == path] plot = visualize.add_annotations(plot, ax, annotations) plot.savefig("{}/{}.png".format(savedir, os.path.splitext(path)[0]), dpi=300) df = pd.concat(prediction_list, ignore_index=True) return df
def predict_image(model, image, return_plot, device, iou_threshold=0.1): """Predict an image with a deepforest model Args: image: a numpy array of a RGB image ranged from 0-255 path: optional path to read image from disk instead of passing image arg return_plot: Return image with plotted detections device: pytorch device of 'cuda' or 'cpu' for gpu prediction. Set internally. Returns: boxes: A pandas dataframe of predictions (Default) img: The input with predictions overlaid (Optional) """ image = preprocess.preprocess_image(image) image = image.to(device) prediction = model(image) # return None for no predictions if len(prediction[0]["boxes"]) == 0: return None # This function on takes in a single image. df = visualize.format_boxes(prediction[0]) df = across_class_nms(df, iou_threshold=iou_threshold) if return_plot: # Matplotlib likes no batch dim and channels first image = image.squeeze(0).permute(1, 2, 0) plot, ax = visualize.plot_predictions(image, df) return plot else: return df
def predict_image(model, image, return_plot, device, iou_threshold=0.1, color=None, thickness=1): """Predict an image with a deepforest model Args: image: a numpy array of a RGB image ranged from 0-255 path: optional path to read image from disk instead of passing image arg return_plot: Return image with plotted detections device: pytorch device of 'cuda' or 'cpu' for gpu prediction. Set internally. 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: boxes: A pandas dataframe of predictions (Default) img: The input with predictions overlaid (Optional) """ if image.dtype != "float32": warnings.warn( f"Image type is {image.dtype}, transforming to float32. This assumes that the range of pixel values is 0-255, as opposed to 0-1.To suppress this warning, transform image (image.astype('float32')" ) image = image.astype("float32") image = preprocess.preprocess_image(image, device=device) with torch.no_grad(): prediction = model(image) # return None for no predictions if len(prediction[0]["boxes"]) == 0: return None # This function on takes in a single image. df = visualize.format_boxes(prediction[0]) df = across_class_nms(df, iou_threshold=iou_threshold) if return_plot: #Bring to gpu if not device.type == "cpu": image = image.cpu() # Cv2 likes no batch dim, BGR image and channels last, 0-255 image = np.array(image.squeeze(0)) image = np.rollaxis(image, 0, 3) image = image[:, :, ::-1] * 255 image = image.astype("uint8") image = visualize.plot_predictions(image, df, color=color, thickness=thickness) return image else: return df
def test_format_boxes(m): ds = m.val_dataloader() batch = next(iter(ds)) paths, images, targets = batch for path, image, target in zip(paths, images, targets): target_df = visualize.format_boxes(target, scores=False) assert list(target_df.columns.values) == [ "xmin", "ymin", "xmax", "ymax", "label" ] assert not target_df.empty
def test_plot_prediction_dataframe(m, tmpdir): ds = m.val_dataloader() batch = next(iter(ds)) paths, images, targets = batch for path, image, target in zip(paths, images, targets): target_df = visualize.format_boxes(target, scores=False) target_df["image_path"] = path visualize.plot_prediction_dataframe( df=target_df, savedir=tmpdir, root_dir=m.config["validation"]["root_dir"])
def test_plot_predictions(m, tmpdir, label): ds = m.val_dataloader() batch = next(iter(ds)) paths, images, targets = batch for path, image, target in zip(paths, images, targets): target_df = visualize.format_boxes(target, scores=False) target_df["image_path"] = path image = np.array(image)[:, :, ::-1] image = np.rollaxis(image, 0, 3) target_df.label = label image = visualize.plot_predictions(image, target_df) assert image.dtype == "uint8"
def test_confirm_backend(m, tmpdir, show): """When a use call a visualize function we are using matplotlib to save figures, this should not interfere with plotting in global scope""" ds = m.val_dataloader() batch = next(iter(ds)) paths, images, targets = batch original_backend = matplotlib.get_backend() for path, image, target in zip(paths, images, targets): target_df = visualize.format_boxes(target, scores=False) target_df["image_path"] = path visualize.plot_prediction_dataframe( df=target_df, savedir=tmpdir, root_dir=m.config["validation"]["root_dir"], show=show) post_backend = matplotlib.get_backend() assert original_backend == post_backend
def predict_file(model, csv_file, root_dir, savedir, device, iou_threshold=0.1, color=(0, 165, 255), thickness=1): """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 relative filename, not absolute path, which is in the root_dir directory. One bounding box per line. If "label" column is present, these are assumed to be annotations and will be plotted in a different color than predictions 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. device: pytorch device of 'cuda' or 'cpu' for gpu prediction. Set internally. 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 """ model.eval() df = pd.read_csv(csv_file) #Dataloader (when not shuffled) returns a tensor for each image in order paths = df.image_path.unique() ds = dataset.TreeDataset(csv_file=csv_file, root_dir=root_dir, transforms=None, train=False) prediction_list = [] with torch.no_grad(): for i in ds: i = i.to(device) prediction = model(torch.unsqueeze(i, 0)) prediction_list.append(prediction) prediction_list = [item for sublist in prediction_list for item in sublist] results = [] for index, prediction in enumerate(prediction_list): #If there is more than one class, apply NMS Loop through images and apply cross prediction = visualize.format_boxes(prediction) if len(prediction.label.unique()) > 1: prediction = across_class_nms(prediction, iou_threshold=iou_threshold) if savedir: # Just predict the images, even though we have the annotations image = np.array(Image.open("{}/{}".format( root_dir, paths[index])))[:, :, ::-1] image = visualize.plot_predictions(image, prediction) #Plot annotations if they exist annotations = df[df.image_path == paths[index]] image = visualize.plot_predictions(image, annotations, color=color, thickness=thickness) cv2.imwrite( "{}/{}.png".format(savedir, os.path.splitext(paths[index])[0]), image) prediction["image_path"] = paths[index] results.append(prediction) results = pd.concat(results, ignore_index=True) return results