From 7dae05fb0865f89dac35136f6233844a65b25e84 Mon Sep 17 00:00:00 2001 From: James Ball Date: Wed, 4 Dec 2024 12:31:49 +0000 Subject: [PATCH] prediction functions with ms --- detectree2/models/predict.py | 2 ++ detectree2/models/train.py | 2 +- 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/detectree2/models/predict.py b/detectree2/models/predict.py index 0478aaa..247b0c3 100644 --- a/detectree2/models/predict.py +++ b/detectree2/models/predict.py @@ -7,6 +7,8 @@ from pathlib import Path import cv2 +import numpy as np +import rasterio from detectron2.engine import DefaultPredictor from detectron2.evaluation.coco_evaluation import instances_to_coco_json diff --git a/detectree2/models/train.py b/detectree2/models/train.py index 527bf4f..7553d9c 100644 --- a/detectree2/models/train.py +++ b/detectree2/models/train.py @@ -32,7 +32,7 @@ build_detection_train_loader, ) from detectron2.data import detection_utils as utils -from detectron2.engine import DefaultTrainer +from detectron2.engine import DefaultPredictor, DefaultTrainer from detectron2.engine.hooks import HookBase from detectron2.evaluation import COCOEvaluator, verify_results from detectron2.evaluation.coco_evaluation import instances_to_coco_json