Skip to content

Commit

Permalink
add a cropmodel sanity check for predicted boxes with less than one p…
Browse files Browse the repository at this point in the history
…ixel width
  • Loading branch information
bw4sz committed Jan 18, 2025
1 parent 7da0b53 commit f1602fb
Showing 1 changed file with 6 additions and 0 deletions.
6 changes: 6 additions & 0 deletions src/deepforest/predict.py
Original file line number Diff line number Diff line change
Expand Up @@ -216,13 +216,19 @@ def _predict_crop_model_(crop_model,
print("No predictions to run crop model on, returning empty dataframe")
return results

# Remove any boxes where xmin = xmax or ymin = ymax
results = results[results.xmin != results.xmax]
results = results[results.ymin != results.ymax]

bounding_box_dataset = dataset.BoundingBoxDataset(
results,
root_dir=os.path.dirname(raster_path),
transform=transform,
augment=augment)

crop_dataloader = crop_model.predict_dataloader(bounding_box_dataset)
crop_results = trainer.predict(crop_model, crop_dataloader)

stacked_outputs = np.vstack(np.concatenate(crop_results))
label = np.argmax(stacked_outputs, 1)
score = np.max(stacked_outputs, 1)
Expand Down

0 comments on commit f1602fb

Please sign in to comment.