update inference

stazizov · Dec 12, 2022 · 2da8e26 · 2da8e26
1 parent e49ceae
commit 2da8e26
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 4 deletions.
diff --git a/.gitignore b/.gitignore
@@ -7,4 +7,5 @@ WSS1*
 **.pth 
 **.pt 
 **.png
+**.jpg 
 **.pyc
diff --git a/faster_rcnn/inference.py b/faster_rcnn/inference.py
@@ -24,16 +24,18 @@
     parser = argparse.ArgumentParser()
     parser.add_argument("--path2image", type=str)
     parser.add_argument("--path2save", type=str)
+    parser.add_argument("--threshold", type=float, default=0.25)
     args = parser.parse_args()
 
     cfg = get_cfg()
     cfg.merge_from_file(model_zoo.get_config_file("COCO-Detection/faster_rcnn_R_50_FPN_1x.yaml"))
     cfg.MODEL.WEIGHTS = "model_final.pth"  # path to the model we just trained
-    cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.001  # set a custom testing threshold
+    cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = args.threshold  # set a custom testing threshold
     cfg.MODEL.DEVICE='cpu'
     cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 128   
     cfg.MODEL.ROI_HEADS.NUM_CLASSES = 1  
-    cfg.SOLVER.BASE_LR = 0.00025  
+    cfg.MODEL.IOU_THRESHOLD = 0.6
+    cfg.SOLVER.BASE_LR = args.threshold
 
     predictor = DefaultPredictor(cfg)
 

diff --git a/notebooks/yolov5.ipynb b/notebooks/yolov5.ipynb
@@ -20,7 +20,7 @@
    "outputs": [],
    "source": [
     "from google.colab import drive\n",
-    "drive.mount('/content/drive')" 
+    "drive.mount('/content/drive')"
    ]
   },
   {
@@ -2037,7 +2037,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.6"
+   "version": "3.10.6 | packaged by conda-forge | (main, Aug 22 2022, 20:38:29) [Clang 13.0.1 ]"
   },
   "vscode": {
    "interpreter": {
-Original file line number
+Diff line change
@@ Expand Up / @@ -7,4 +7,5 @@ WSS1* @@
     **.pth
     **.pt
     **.png
+    **.jpg
     **.pyc