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
