prediction functions with ms

PatBall1 · Dec 4, 2024 · 7dae05f · 7dae05f
1 parent cbdc5cb
commit 7dae05f
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 1 deletion.
diff --git 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
@@ -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