Run on GPU

Dockerfile

@@ -49,4 +49,4 @@ EXPOSE 7860
 
 HEALTHCHECK --interval=30s --timeout=30s --start-period=5s --retries=3 CMD [ "curl", "localhost:7860" ]
 
-CMD ["python", "main.py"]
+CMD ["python", "-u", "main.py"]

detectree2model/predictions/predict.py

@@ -55,7 +55,7 @@ def predict(tile_path, overlap_threshold, confidence_threshold, simplify_value,
     cfg = setup_cfg(update_model=trained_model, out_dir=store_path)
 
     # hash the following line if you have gpu support
-    cfg.MODEL.DEVICE = "cpu"
+    # cfg.MODEL.DEVICE = "cpu"
     predict_on_data(tile_path, predictor=DefaultPredictor(cfg))
 
     project_to_geojson(tile_path, tile_path + "predictions/", tile_path + "predictions_geo/")