Examples

Python Examples

Baseline Boxes Example

Download baseline_boxes.py

import os
import argparse
import warnings
from typing import List, Tuple

import numpy as np
import pandas as pd
import torch

from deepforest import main as df_main
from deepforest.utilities import read_file, format_geometry
import matplotlib.pyplot as plt
import matplotlib.patches as patches

from milliontrees import get_dataset
from milliontrees.common.data_loaders import get_eval_loader


def format_deepforest_predictions(
    images: np.ndarray,
    metadata: torch.Tensor,
    targets: List[dict],
    model: "df_main.deepforest",
    dataset,
    batch_index: int,
) -> Tuple[List[dict], List[pd.DataFrame]]:
    """Run DeepForest on a batch and convert to MillionTrees format.

    Returns a tuple of:
    - list of y_pred dicts for MillionTrees evaluation
    - list of original prediction DataFrames (for plotting)
    """
    warnings.filterwarnings("ignore")

    images_tensor = torch.tensor(images)
    predictions = model.predict_step(images_tensor, batch_index)

    batch_y_pred: List[dict] = []
    formatted_predictions: List[pd.DataFrame] = []

    for image_metadata, pred, image_targets, image in zip(
            metadata, predictions, targets, images_tensor):
        basename = dataset._filename_id_to_code[int(image_metadata[0])]

        if pred is None or len(pred["boxes"]) == 0:
            y_pred = {
                "y": torch.zeros((0, 4), dtype=torch.float32),
                "labels": torch.zeros((0,), dtype=torch.int64),
                "scores": torch.zeros((0,), dtype=torch.float32),
            }
            formatted_pred = pd.DataFrame(
                columns=["xmin", "ymin", "xmax", "ymax", "score", "label"])  # empty
            formatted_pred.root_dir = os.path.join(dataset._data_dir._str,
                                                   "images")
            formatted_pred["image_path"] = basename
        else:
            formatted_pred = format_geometry(pred)
            formatted_pred.root_dir = os.path.join(dataset._data_dir._str,
                                                   "images")
            formatted_pred["image_path"] = basename

            y_pred = {
                "y": torch.tensor(
                    formatted_pred[["xmin", "ymin", "xmax",
                                     "ymax"]].values.astype("float32")),
                "labels": torch.tensor(formatted_pred.label.values.astype(
                    np.int64)),
                "scores": torch.tensor(
                    formatted_pred.score.values.astype("float32")),
            }

        batch_y_pred.append(y_pred)
        formatted_predictions.append(formatted_pred)

    return batch_y_pred, formatted_predictions


def _draw_boxes(ax, boxes: np.ndarray, color: str):
    for box in boxes:
        if len(box) != 4:
            continue
        xmin, ymin, xmax, ymax = box
        width = max(0.0, float(xmax) - float(xmin))
        height = max(0.0, float(ymax) - float(ymin))
        rect = patches.Rectangle((float(xmin), float(ymin)), width, height,
                                 linewidth=0.8,
                                 edgecolor=color,
                                 facecolor='none',
                                 alpha=0.8)
        ax.add_patch(rect)


def save_gallery(thumbnails: List[dict], rows: int, cols: int, dpi: int,
                 output_path: str):
    if len(thumbnails) == 0:
        return
    n = min(len(thumbnails), rows * cols)
    fig, axes = plt.subplots(rows,
                             cols,
                             figsize=(cols * 3, rows * 3),
                             dpi=dpi)
    axes = np.atleast_2d(axes)

    for idx in range(rows * cols):
        r = idx // cols
        c = idx % cols
        ax = axes[r, c]
        if idx < n:
            thumb = thumbnails[idx]
            ax.imshow(thumb["image"], interpolation='nearest')
            _draw_boxes(ax, thumb.get("gt_boxes", np.zeros((0, 4))),
                        color='orange')
            _draw_boxes(ax, thumb.get("pred_boxes", np.zeros((0, 4))),
                        color='royalblue')
            ax.set_title(thumb.get("title", ""), fontsize=8)
        ax.axis('off')

    plt.tight_layout(pad=0.2)
    os.makedirs(os.path.dirname(output_path), exist_ok=True)
    fig.savefig(output_path, bbox_inches='tight')
    plt.close(fig)


def main():
    parser = argparse.ArgumentParser(
        description="Run baseline DeepForest evaluation on TreeBoxes.")
    parser.add_argument(
        "--root-dir",
        type=str,
        default=os.environ.get("MT_ROOT",
                               "/orange/ewhite/web/public/MillionTrees"),
        help="Dataset root directory",
    )
    parser.add_argument("--batch-size", type=int, default=12)
    parser.add_argument(
        "--plot-interval",
        type=int,
        default=250,
        help="Plot every Nth image; set 0 to disable plotting",
    )
    parser.add_argument("--gallery-rows", type=int, default=2)
    parser.add_argument("--gallery-cols", type=int, default=3)
    parser.add_argument("--gallery-dpi", type=int, default=72)
    parser.add_argument("--output-dir", type=str, default=None)
    parser.add_argument("--max-batches", type=int, default=None)
    parser.add_argument("--mini", action="store_true", help="Use mini datasets for fast dev")
    parser.add_argument("--download", action="store_true", help="Download dataset if missing")
    parser.add_argument("--split-scheme",
                        type=str,
                        default="random",
                        choices=["random", "zeroshot", "crossgeometry"],
                        help="Dataset split scheme")
    args = parser.parse_args()

    # Load model
    model = df_main.deepforest()
    model.load_model("weecology/deepforest-tree")
    model.eval()

    # Load dataset
    box_dataset = get_dataset("TreeBoxes",
                              download=args.download,
                              mini=args.mini,
                              root_dir=args.root_dir,
                              split_scheme=args.split_scheme)
    test_subset = box_dataset.get_subset("test")
    test_loader = get_eval_loader("standard",
                                  test_subset,
                                  batch_size=args.batch_size)

    print(f"There are {len(test_loader)} batches in the test loader")

    all_y_pred: List[dict] = []
    all_y_true: List[dict] = []

    batch_index = 0
    thumbnails: List[dict] = []
    max_thumbnails = args.gallery_rows * args.gallery_cols if args.output_dir else 0
    for batch in test_loader:
        metadata, images, targets = batch
        mt_preds, df_preds = format_deepforest_predictions(images, metadata,
                                                           targets, model,
                                                           box_dataset,
                                                           batch_index)

        for image_metadata, y_pred, pred, image_targets, image in zip(
                metadata, mt_preds, df_preds, targets, images):
            
            # Collect thumbnails for gallery
            if args.output_dir and len(thumbnails) < max_thumbnails:
                basename = (pred.image_path.unique()[0]
                            if isinstance(pred, pd.DataFrame)
                            and len(pred) > 0 else "")
                image_np = (image.permute(1, 2, 0).numpy() * 255).clip(
                    0, 255).astype("uint8")
                pred_boxes = y_pred.get("y", torch.zeros((0, 4))).detach().cpu().numpy()
                gt_boxes = image_targets["y"].detach().cpu().numpy()
                recall = box_dataset.metrics["recall"]._recall(
                    image_targets["y"], y_pred.get("y", torch.zeros((0, 4))),
                    iou_threshold=0.4)
                title = f"{basename} R@0.4={float(recall):.2f}"
                thumbnails.append({
                    "image": image_np,
                    "pred_boxes": pred_boxes,
                    "gt_boxes": gt_boxes,
                    "title": title
                })

            all_y_pred.append(y_pred)
            all_y_true.append(image_targets)
            batch_index += 1

        if args.max_batches is not None and batch_index >= args.max_batches:
            break

    results, results_str = box_dataset.eval(all_y_pred, all_y_true,
                                            test_subset.metadata_array[:len(all_y_true)])
    print(results_str)

    if args.output_dir:
        os.makedirs(args.output_dir, exist_ok=True)
        with open(os.path.join(args.output_dir, "results_boxes.txt"), "w",
                  encoding="utf-8") as f:
            f.write(results_str)
        # Save a compact gallery image
        gallery_path = os.path.join(args.output_dir, "gallery_boxes.png")
        save_gallery(thumbnails, args.gallery_rows, args.gallery_cols,
                     args.gallery_dpi, gallery_path)


if __name__ == "__main__":
    main()

Baseline Points Example

Download baseline_points.py

import os
import argparse
import warnings
from typing import List, Tuple

import numpy as np
import pandas as pd
import torch

from deepforest import main as df_main
from deepforest.utilities import read_file, format_geometry
from deepforest.visualize import plot_results
import geopandas as gpd

from milliontrees import get_dataset
from milliontrees.common.data_loaders import get_eval_loader


def _map_labels_to_int(values: pd.Series, model: "df_main.deepforest") -> np.ndarray:
    if values.dtype == object:
        return values.apply(lambda x: model.label_dict.get(x, 0)).values.astype(np.int64)
    return values.values.astype(np.int64)


def format_deepforest_predictions(
    images: np.ndarray,
    metadata: torch.Tensor,
    targets: List[dict],
    model: "df_main.deepforest",
    dataset,
    batch_index: int,
) -> Tuple[List[dict], List[pd.DataFrame]]:
    """Run DeepForest on a batch and convert to MillionTrees point format."""
    warnings.filterwarnings("ignore")

    images_tensor = torch.tensor(images)
    predictions = model.predict_step(images_tensor, batch_index)

    batch_y_pred: List[dict] = []
    formatted_predictions: List[pd.DataFrame] = []

    for image_metadata, image_pred, image_targets, image in zip(
            metadata, predictions, targets, images_tensor):
        basename = dataset._filename_id_to_code[int(image_metadata[0])]

        if len(image_pred["boxes"]) == 0:
            y_pred = {
                "y": torch.zeros((0, 2), dtype=torch.float32),
                "labels": torch.zeros((0,), dtype=torch.int64),
                "scores": torch.zeros((0,), dtype=torch.float32),
            }
            formatted_pred = pd.DataFrame(columns=["x", "y", "score", "label"])  # empty
        else:
            formatted_pred = format_geometry(image_pred)
            formatted_pred["image_path"] = basename
            formatted_pred = read_file(
                formatted_pred,
                root_dir=os.path.join(dataset._data_dir._str, "images"),
                image_path=basename,
            )
            
            # Convert boxes to centroids
            formatted_pred["geometry"] = gpd.GeoSeries(formatted_pred["geometry"]).centroid
            formatted_pred[["x", "y"]] = formatted_pred["geometry"].apply(lambda g: pd.Series([g.x, g.y]))
            y_pred = {
                "y": torch.tensor(formatted_pred[["x", "y"]].values.astype("float32")),
                "labels": torch.tensor(_map_labels_to_int(formatted_pred.label, model)),
                "scores": torch.tensor(formatted_pred.score.values.astype("float32")),
            }

        batch_y_pred.append(y_pred)
        formatted_predictions.append(formatted_pred)

    return batch_y_pred, formatted_predictions


def plot_eval_result(
    y_pred: dict,
    pred_df: pd.DataFrame,
    image_targets: dict,
    image_tensor: torch.Tensor,
    dataset,
    batch_index: int,
    output_dir: str = None,
):
    basename = (
        pred_df.image_path.unique()[0]
        if isinstance(pred_df, pd.DataFrame) and len(pred_df) > 0
        else "empty"
    )

    # Ground truth
    gt_df = pd.DataFrame(image_targets["y"].numpy(), columns=["x", "y"])
    gt_df["image_path"] = basename
    gt_df = read_file(
        gt_df,
        root_dir=os.path.join(dataset._data_dir._str, "images"),
        image_path=basename,
        label="Tree",
    )
    gt_df["label"] = "Tree"
    gt_df["score"] = 1

    # Predictions
    if isinstance(pred_df, pd.DataFrame) and len(pred_df) > 0:
        pred_vis_df = read_file(
            pred_df,
            root_dir=os.path.join(dataset._data_dir._str, "images"),
            image_path=basename,
            label="Tree",
        )
        pred_vis_df["geometry"] = gpd.GeoSeries(pred_vis_df["geometry"]).centroid
        pred_vis_df[["x", "y"]] = pred_vis_df["geometry"].apply(lambda g: pd.Series([g.x, g.y]))
        if "label" not in pred_vis_df.columns:
            pred_vis_df["label"] = "Tree"
    else:
        pred_vis_df = pred_df

    # Image channel-last, 0-255
    image = image_tensor.permute(1, 2, 0).numpy() * 255

    pred_vis_df.root_dir = os.path.join(dataset._data_dir._str, "images")
    fig = plot_results(pred_vis_df, gt_df, image=image.astype("int32"))
    if output_dir:
        os.makedirs(output_dir, exist_ok=True)
        out_path = os.path.join(output_dir, f"{batch_index:06d}_{basename}.png")
        fig.savefig(out_path, dpi=150, bbox_inches="tight")

def main():
    parser = argparse.ArgumentParser(description="Run baseline DeepForest on TreePoints using centroid conversion.")
    parser.add_argument(
        "--root-dir",
        type=str,
        default=os.environ.get("MT_ROOT", "/orange/ewhite/web/public/MillionTrees/"),
        help="Dataset root directory",
    )
    parser.add_argument("--batch-size", type=int, default=32)
    parser.add_argument("--plot-interval", type=int, default=1000, help="Plot every Nth image; set 0 to disable plotting")
    parser.add_argument("--output-dir", type=str, default=None)
    parser.add_argument("--max-batches", type=int, default=None)
    parser.add_argument("--mini", action="store_true", help="Use mini datasets for fast dev")
    parser.add_argument("--download", action="store_true", help="Download dataset if missing")
    parser.add_argument("--split-scheme",
                        type=str,
                        default="random",
                        choices=["random", "zeroshot", "crossgeometry"],
                        help="Dataset split scheme")
    args = parser.parse_args()

    # Load model
    model = df_main.deepforest()
    model.load_model("weecology/deepforest-tree")
    model.eval()

    # Load dataset
    point_dataset = get_dataset("TreePoints",
                                root_dir=args.root_dir,
                                mini=args.mini,
                                download=args.download,
                                split_scheme=args.split_scheme)
    test_subset = point_dataset.get_subset("test")
    test_loader = get_eval_loader("standard", test_subset, batch_size=args.batch_size)

    print(f"There are {len(test_loader)} batches in the test loader")

    all_y_pred: List[dict] = []
    all_y_true: List[dict] = []

    batch_index = 0
    for batch in test_loader:
        metadata, images, targets = batch
        mt_preds, df_preds = format_deepforest_predictions(images, metadata, targets, model, point_dataset, batch_index)

        for image_metadata, y_pred, pred, image_targets, image in zip(metadata, mt_preds, df_preds, targets, images):
            if args.plot_interval and args.plot_interval > 0 and (batch_index % args.plot_interval == 0):
                plot_eval_result(y_pred, pred, image_targets, image, point_dataset, batch_index, args.output_dir)

            all_y_pred.append(y_pred)
            all_y_true.append(image_targets)
            batch_index += 1

        if args.max_batches is not None and batch_index >= args.max_batches:
            break

    # Evaluate using dataset's metric implementation
    results, results_str = point_dataset.eval(all_y_pred, all_y_true, test_subset.metadata_array[:len(all_y_true)])
    print(results_str)

    if args.output_dir:
        os.makedirs(args.output_dir, exist_ok=True)
        with open(os.path.join(args.output_dir, "results_points.txt"), "w", encoding="utf-8") as f:
            f.write(results_str)


if __name__ == "__main__":
    main()

Baseline Polygons Example

Download baseline_polygons.py

import os
import argparse
import warnings
from typing import List, Tuple

import numpy as np
import pandas as pd
import torch

from deepforest import main as df_main
from deepforest.utilities import read_file, format_geometry
from deepforest.visualize import plot_results

from milliontrees import get_dataset
from milliontrees.common.data_loaders import get_eval_loader


def format_deepforest_predictions(
    images: np.ndarray,
    metadata: torch.Tensor,
    targets: List[dict],
    model,  # Remove the problematic type hint
    dataset,
    batch_index: int,
) -> Tuple[List[dict], List[pd.DataFrame]]:
    """
    Run DeepForest on a batch and convert to MillionTrees format.
    
    Args:
        images: Input images as numpy array
        metadata: Tensor containing metadata information  
        targets: List of target dictionaries
        model: DeepForest model instance
        dataset: Dataset instance
        batch_index: Index of the current batch
        
    Returns:
        Tuple containing predictions in MillionTrees format and formatted DataFrames
    """
    warnings.filterwarnings("ignore")

    images_tensor = torch.tensor(images)
    predictions = model.predict_step(images_tensor, batch_index)

    batch_y_pred: List[dict] = []
    formatted_predictions: List[pd.DataFrame] = []

    for image_metadata, pred, image_targets, image in zip(
        metadata, predictions, targets, images_tensor
    ):
        basename = dataset._filename_id_to_code[int(image_metadata[0])]

        if pred is None or len(pred["boxes"]) == 0:
            y_pred = {
                "y": torch.zeros((0, 4), dtype=torch.float32),
                "labels": torch.zeros((0,), dtype=torch.int64),
                "scores": torch.zeros((0,), dtype=torch.float32),
            }
            formatted_pred = pd.DataFrame(
                columns=["xmin", "ymin", "xmax", "ymax", "score", "label"]
            )
            formatted_pred.root_dir = os.path.join(dataset._data_dir._str, "images")
            formatted_pred["image_path"] = basename
        else:
            formatted_pred = format_geometry(pred)
            formatted_pred.root_dir = os.path.join(dataset._data_dir._str, "images")
            formatted_pred["image_path"] = basename

            y_pred = {
                "y": torch.tensor(
                    formatted_pred[["xmin", "ymin", "xmax", "ymax"]].values.astype("float32")
                ),
                "labels": torch.tensor(formatted_pred.label.values.astype(np.int64)),
                "scores": torch.tensor(formatted_pred.score.values.astype("float32")),
            }

        batch_y_pred.append(y_pred)
        formatted_predictions.append(formatted_pred)

    return batch_y_pred, formatted_predictions


def plot_eval_result(
    y_pred: dict,
    pred_df: pd.DataFrame,
    image_targets: dict,
    image_tensor: torch.Tensor,
    dataset,
    batch_index: int,
    output_dir: str = None,
):
    basename = pred_df.image_path.unique()[0] if len(pred_df) > 0 else "empty"

    # Ground truth
    gt_df = pd.DataFrame(
        image_targets["bboxes"],
        columns=["xmin", "ymin", "xmax", "ymax"],
    )
    gt_df["image_path"] = basename
    gt_df = read_file(
        gt_df,
        root_dir=os.path.join(dataset._data_dir._str, "images"),
        image_path=basename,
        label="Tree",
    )
    gt_df["label"] = "Tree"

    # Predictions
    pred_vis_df = read_file(
        pred_df,
        root_dir=os.path.join(dataset._data_dir._str, "images"),
        image_path=basename,
    ) if len(pred_df) > 0 else pred_df
    if len(pred_df) > 0 and "label" not in pred_vis_df.columns:
        pred_vis_df["label"] = "Tree"

    # Image channel-last, 0-255
    image = image_tensor.permute(1, 2, 0).numpy() * 255

    # Simple recall example for logging
    recall = dataset.metrics["recall"]._recall(
        image_targets["bboxes"],
        y_pred.get("bboxes", torch.zeros((0, 4))),
        iou_threshold=0.4
    )

    # Plot
    try:
        fig = plot_results(pred_vis_df, gt_df, image=image.astype("int32"))
        if output_dir:
            os.makedirs(output_dir, exist_ok=True)
            out_path = os.path.join(output_dir, f"{batch_index:06d}_{basename}.png")
            fig.savefig(out_path, dpi=150, bbox_inches="tight")
    except Exception:
        pass

    print(f"Image: {basename}, idx {batch_index}, Recall@0.4: {float(recall):.2f}")


def main():
    parser = argparse.ArgumentParser(
        description="Run baseline DeepForest evaluation on TreePolygons."
    )
    parser.add_argument(
        "--root-dir",
        type=str,
        default=os.environ.get("MT_ROOT", "/orange/ewhite/web/public/MillionTrees/"),
        help="Dataset root directory",
    )
    parser.add_argument("--batch-size", type=int, default=32)
    parser.add_argument(
        "--plot-interval",
        type=int,
        default=250,
        help="Plot every Nth image; set 0 to disable plotting",
    )
    parser.add_argument("--output-dir", type=str, default=None)
    parser.add_argument("--max-batches", type=int, default=None)
    parser.add_argument(
        "--mini", action="store_true", help="Use mini datasets for fast dev"
    )
    parser.add_argument(
        "--download", action="store_true", help="Download dataset if missing"
    )
    parser.add_argument(
        "--split-scheme",
        type=str,
        default="random",
        choices=["random", "zeroshot", "crossgeometry"],
        help="Dataset split scheme",
    )
    args = parser.parse_args()

    # Load model
    model = df_main.deepforest()
    model.load_model("weecology/deepforest-tree")
    model.eval()

    # Load dataset
    polygon_dataset = get_dataset(
        "TreePolygons",
        root_dir=args.root_dir,
        mini=args.mini,
        download=args.download,
        split_scheme=args.split_scheme,
        image_size=224,
    )
    test_subset = polygon_dataset.get_subset("test")
    test_loader = get_eval_loader(
        "standard", test_subset, batch_size=args.batch_size
    )

    print(f"There are {len(test_loader)} batches in the test loader")

    all_y_pred: List[dict] = []
    all_y_true: List[dict] = []

    batch_index = 0
    for batch in test_loader:
        metadata, images, targets = batch
        mt_preds, df_preds = format_deepforest_predictions(
            images, metadata, targets, model, polygon_dataset, batch_index
        )

        for image_metadata, y_pred, pred, image_targets, image in zip(
            metadata, mt_preds, df_preds, targets, images
        ):
            if args.plot_interval and args.plot_interval > 0 and (
                batch_index % args.plot_interval == 0
            ):
                plot_eval_result(
                    y_pred, pred, image_targets, image,
                    polygon_dataset, batch_index, args.output_dir
                )

            all_y_pred.append(y_pred)
            all_y_true.append(image_targets)
            batch_index += 1

        if args.max_batches is not None and batch_index >= args.max_batches:
            break

    results, results_str = polygon_dataset.eval(
        all_y_pred,
        all_y_true,
        test_subset.metadata_array[:len(all_y_true)],
    )
    print(results_str)

    if args.output_dir:
        os.makedirs(args.output_dir, exist_ok=True)
        output_path = os.path.join(args.output_dir, "results_polygons.txt")
        with open(output_path, "w", encoding="utf-8") as f:
            f.write(results_str)


if __name__ == "__main__":
    main()

View Examples on GitHub

All example files are also available on GitHub:

• baseline_boxes.py

• baseline_points.py

• baseline_polygons.py