diff --git a/hot_fair_utilities/training/yolo_v8_v1/train.py b/hot_fair_utilities/training/yolo_v8_v1/train.py index edc41da..e9ff763 100644 --- a/hot_fair_utilities/training/yolo_v8_v1/train.py +++ b/hot_fair_utilities/training/yolo_v8_v1/train.py @@ -9,7 +9,7 @@ # Reader imports from hot_fair_utilities.model.yolo import YOLOSegWithPosWeight - +from ...utils import compute_iou_chart_from_yolo_results, get_yolo_iou_metrics # Get environment variables with fallbacks ROOT = Path(os.getenv("YOLO_ROOT", Path(__file__).parent.absolute())) DATA_ROOT = str(Path(os.getenv("YOLO_DATA_ROOT", ROOT / "yolo-training"))) @@ -121,7 +121,14 @@ def train( device=[int(i) for i in gpu.split(",")] if "," in gpu else gpu, **kwargs, ) - return weights + compute_iou_chart_from_yolo_results(results_csv_path=os.path.join(output_path,"checkpoints", name,'results.csv'),results_output_chart_path=os.path.join(output_path,"checkpoints", name,'iou_chart.png')) + + output_model_path=os.path.join(os.path.join(output_path,"checkpoints"), name, "weights", "best.pt") + + iou_model_accuracy=get_yolo_iou_metrics(output_model_path) + + return output_model_path,iou_model_accuracy + def check4checkpoint(name, weights,output_path): diff --git a/hot_fair_utilities/training/yolo_v8_v2/train.py b/hot_fair_utilities/training/yolo_v8_v2/train.py index c473810..9afd1c7 100644 --- a/hot_fair_utilities/training/yolo_v8_v2/train.py +++ b/hot_fair_utilities/training/yolo_v8_v2/train.py @@ -6,7 +6,7 @@ # Third party imports import torch import ultralytics - +from ...utils import get_yolo_iou_metrics,compute_iou_chart_from_yolo_results # Reader imports from hot_fair_utilities.model.yolo import YOLOSegWithPosWeight @@ -86,7 +86,13 @@ def train(data, weights, epochs, batch_size, pc, output_path, dataset_yaml_path, # metrics = model.val(save_json=True, plots=True) # print(model.val()) - return os.path.join(os.path.join(output_path,"checkpoints"), name, "weights", "best.pt") + compute_iou_chart_from_yolo_results(results_csv_path=os.path.join(output_path,"checkpoints", name,'results.csv'),results_output_chart_path=os.path.join(output_path,"checkpoints", name,'iou_chart.png')) + + output_model_path=os.path.join(os.path.join(output_path,"checkpoints"), name, "weights", "best.pt") + + iou_model_accuracy=get_yolo_iou_metrics(output_model_path) + + return output_model_path,iou_model_accuracy def check4checkpoint(name, weights,output_path): diff --git a/hot_fair_utilities/utils.py b/hot_fair_utilities/utils.py index f0883c6..c30d358 100644 --- a/hot_fair_utilities/utils.py +++ b/hot_fair_utilities/utils.py @@ -5,12 +5,15 @@ import math import os import re +import ultralytics + import time import urllib.request import zipfile from glob import glob from typing import Tuple - +import pandas as pd +import matplotlib.pyplot as plt # Third party imports # Third-party imports import geopandas @@ -251,3 +254,45 @@ def fetch_osm_data(payload: json, API_URL="https://api-prod.raw-data.hotosm.org/ with zip_ref.open("Export.geojson") as file: my_export_geojson = json.loads(file.read()) return my_export_geojson + + +import pandas as pd +import matplotlib.pyplot as plt + + +def compute_iou_chart_from_yolo_results(results_csv_path,results_output_chart_path): + + data = pd.read_csv(results_csv_path) + + + data['IoU(M)'] = 1 / ( + 1 / data['metrics/precision(M)'] + 1 / data['metrics/recall(M)'] - 1 + ) + + plt.figure(figsize=(10, 5)) + plt.plot(data['epoch'], data['IoU(M)'], label='IoU (Mask)') + plt.xlabel('Epoch') + plt.xticks(data['epoch'].astype(int)) + plt.ylabel('IoU') + plt.title('IoU over Epochs') + plt.legend() + plt.grid() + + plt.savefig(results_output_chart_path) + return results_output_chart_path + + +def get_yolo_iou_metrics(model_path): + + model_val = ultralytics.YOLO(model_path) + model_val_metrics = ( + model_val.val().results_dict + ) ### B and M denotes bounding box and mask respectively + # print(metrics) + iou_accuracy = 1 / ( + 1 / model_val_metrics["metrics/precision(M)"] + + 1 / model_val_metrics["metrics/recall(M)"] + - 1 + ) # ref here https://github.com/ultralytics/ultralytics/issues/9984#issuecomment-2422551315 + final_accuracy = iou_accuracy * 100 + return final_accuracy \ No newline at end of file diff --git a/test_yolo_v1.py b/test_yolo_v1.py index 9cf4109..b24f044 100644 --- a/test_yolo_v1.py +++ b/test_yolo_v1.py @@ -53,10 +53,10 @@ def __exit__(self, type, value, traceback): p_val=0.05, ) -output_model_path = train_yolo( +output_model_path,output_model_iou_accuracy = train_yolo( data=f"{base_path}", weights=f"{os.getcwd()}/yolov8s_v1-seg-best.pt", - gpu="cpu", + # gpu="cpu", epochs=2, batch_size=16, pc=2.0, @@ -64,6 +64,7 @@ def __exit__(self, type, value, traceback): dataset_yaml_path=os.path.join(yolo_data_dir,'yolo_dataset.yaml') ) +print(output_model_iou_accuracy) prediction_output = f"{base_path}/prediction/output" # model_path = f"{output_path}/weights/best.pt" diff --git a/test_yolo_v2.py b/test_yolo_v2.py index c0c7b3d..555bced 100644 --- a/test_yolo_v2.py +++ b/test_yolo_v2.py @@ -2,6 +2,7 @@ import os import time import warnings +import ultralytics os.environ.update(os.environ) os.environ["RAMP_HOME"] = os.getcwd() @@ -50,18 +51,17 @@ def __exit__(self, type, value, traceback): output_path=yolo_data_dir, ) -output_model_path = train_yolo( +output_model_path,output_model_iou_accuracy = train_yolo( data=f"{base_path}", - weights=f"{os.getcwd()}/yolov8s_v2-seg.pt", ## Todo : replace with finetuned ramp model checkpoint - gpu="cpu", + weights=f"{os.getcwd()}/yolov8s_v2-seg.pt", + # gpu="cpu", epochs=2, batch_size=16, pc=2.0, output_path=yolo_data_dir, dataset_yaml_path=os.path.join(yolo_data_dir,'yolo_dataset.yaml') ) - -print(output_model_path) +print(output_model_iou_accuracy) prediction_output = f"{base_path}/prediction/output" # model_path = f"{output_path}/weights/best.pt"