added prediction stage to dvc pipeline

dvc.lock

@@ -8,14 +8,14 @@ stages:
  size: 28867246
  nfiles: 1482
  - path: params.yaml
- md5: f2fc712e78c50f898902306181630ead
+ md5: a0c123b909925d3386149e8b024832a0
  size: 215
  - path: src/train.py
- md5: 8402b07d9f68ed886e86f8b644b2a843
- size: 2002
+ md5: 6746a939a17bbf4e2e2bd23a226b69ec
+ size: 1774
  - path: src/utils.py
- md5: 4d04addab005a6475665898aae662ace
- size: 1815
+ md5: 5440f0cb53d7c318e2233a2414c32147
+ size: 1228
  outs:
  - path: models/model.pt
  md5: 7fe4c6173c334d2e926e71236366b8e0
@@ -26,3 +26,16 @@ stages:
  - path: reports/train_params.yaml
  md5: 2ea9f04ec479dd1db02f44263b2c8328
  size: 1501
+ test:
+ cmd: python src/predict.py
+ deps:
+ - path: data/raw/wildfire-raw-yolov8
+ md5: ffe9083905cf64817ca61146b3356bb6.dir
+ size: 28867246
+ nfiles: 1482
+ - path: models/model.pt
+ md5: 7fe4c6173c334d2e926e71236366b8e0
+ size: 6226104
+ - path: src/predict.py
+ md5: fa6b9e2cfac6cbe9a40e8d182a0a2f4c
+ size: 594

dvc.yaml

@@ -13,3 +13,9 @@ stages:
  metrics:
  - reports/train_metrics.csv:
  cache: false
+ test:
+ cmd: python src/predict.py
+ deps:
+ - data/raw/wildfire-raw-yolov8
+ - models/model.pt
+ - src/predict.py

params.yaml

@@ -3,7 +3,7 @@ pretrained: True
 seed: 0
 imgsz: 640 
 batch: 8 
-epochs: 2
+epochs: 1
 optimizer: SGD # other choices=['SGD', 'Adam', 'AdamW', 'RMSProp']
 lr0: 0.01 # learning rate 
 name: 'yolov8n_exp_v0' # experiment name

src/predict.py

@@ -0,0 +1,23 @@
+import os
+import glob
+from pathlib import Path
+
+from ultralytics import YOLO
+
+
+root_dir = Path(__file__).resolve().parents[1] # root directory
+path_model_weights = os.path.join(root_dir, 'models/model.pt')
+data_dir = os.path.join(root_dir, "data/raw/wildfire-raw-yolov8")
+test_data_path = os.path.join(data_dir, 'test/images')
+
+
+if __name__ == '__main__':
+
+ #load trained model 
+ model = YOLO(path_model_weights)
+
+ for images in glob.glob(f"{test_data_path}/*.jpg"):
+ predictions = model.predict(images, save=True)
+ print(predictions)
+ print(type(predictions)) 
+

src/train.py

@@ -6,7 +6,7 @@
 from dotenv import load_dotenv
 from ultralytics import YOLO
 
-from utils import convert_metrics_csv_to_json, save_metrics_and_params, save_model
+from utils import save_metrics_and_params, save_model
 
 
 load_dotenv()
@@ -53,9 +53,7 @@
  mlflow.log_param('learning_rate', params['lr0'])
 
  # save model
- model_path = save_model(experiment_name=params['name']) 
- # log model artifact to mlflow 
- mlflow.log_artifact(model_path)
+ save_model(experiment_name=params['name']) 
 
  # save metrics csv file and training params 
  save_metrics_and_params(experiment_name=params['name'])