diff --git a/segmentation_mAP.py b/segmentation_mAP.py new file mode 100644 index 0000000..e4987b2 --- /dev/null +++ b/segmentation_mAP.py @@ -0,0 +1,256 @@ +''' +URL : https://github.com/cosmiq/solaris +Code based on: https://github.com/CosmiQ/solaris/blob/main/solaris/eval/vector.py +''' + +import os +import glob +from tqdm import tqdm +import numpy as np +import geopandas as gpd + + +def average_score_by_class(ious, threshold=0.5): + """ for a list of object ious by class, test if they are a counted as a + positive or a negative. + Arguments + --------- + ious : list of lists + A list containing individual lists of ious for eachobject class. + threshold : float + A value between 0.0 and 1.0 that determines the threshold for a true positve. + Returns + --------- + average_by_class : list + A list containing the ratio of true positives for each class + """ + binary_scoring_lists = [] + for x in ious: + items = [] + for i in x: + if i >= threshold: + items.append(1) + else: + items.append(0) + binary_scoring_lists.append(items) + average_by_class = [] + for l in binary_scoring_lists: + average_by_class.append(np.nanmean(l)) + return average_by_class + + + + +def calculate_iou(pred_poly, test_data_GDF): + """Get the best intersection over union for a predicted polygon. + + Arguments + --------- + pred_poly : :py:class:`shapely.Polygon` + Prediction polygon to test. + test_data_GDF : :py:class:`geopandas.GeoDataFrame` + GeoDataFrame of ground truth polygons to test ``pred_poly`` against. + + Returns + ------- + iou_GDF : :py:class:`geopandas.GeoDataFrame` + A subset of ``test_data_GDF`` that overlaps ``pred_poly`` with an added + column ``iou_score`` which indicates the intersection over union value. + + """ + + # Fix bowties and self-intersections + if not pred_poly.is_valid: + pred_poly = pred_poly.buffer(0.0) + + precise_matches = test_data_GDF[test_data_GDF.intersects(pred_poly)] + + iou_row_list = [] + for _, row in precise_matches.iterrows(): + # Load ground truth polygon and check exact iou + test_poly = row.geometry + # Ignore invalid polygons for now + if pred_poly.is_valid and test_poly.is_valid: + intersection = pred_poly.intersection(test_poly).area + union = pred_poly.union(test_poly).area + # Calculate iou + iou_score = intersection / float(union) + else: + iou_score = 0 + row['iou_score'] = iou_score + iou_row_list.append(row) + + iou_GDF = gpd.GeoDataFrame(iou_row_list) + return iou_GDF + +def get_all_objects(proposal_polygons_dir, gt_polygons_dir, + prediction_cat_attrib="class", gt_cat_attrib='make', + file_format="geojson"): + """ Using the proposal and ground truth polygons, calculate the total. + Filenames of predictions and ground-truth must be identical. + unique classes present in each + Arguments + --------- + proposal_polygons_dir : str + The path that contains any model proposal polygons + gt_polygons_dir : str + The path that contains the ground truth polygons + prediction_cat_attrib : str + The column or attribute within the predictions that specifies + unique classes + gt_cat_attrib : str + The column or attribute within the ground truth that + specifies unique classes + file_format : str + The extension or file format for predictions + Returns + --------- + prop_objs : list + All unique objects that exist in the proposals + gt_obj : list + All unique objects that exist in the ground truth + all_objs : list + A union of the prop_objs and gt_objs lists + """ + objs = [] + os.chdir(proposal_polygons_dir) + search = "*" + file_format + proposal_geojsons = glob.glob(search) + for geojson in tqdm(proposal_geojsons): + ground_truth_poly = os.path.join(gt_polygons_dir, geojson) + if os.path.exists(ground_truth_poly): + ground_truth_gdf = gpd.read_file(ground_truth_poly) + proposal_gdf = gpd.read_file(geojson) + for index, row in (proposal_gdf.iterrows()): + objs.append(row[prediction_cat_attrib]) + prop_objs = list(set(objs)) + os.chdir(gt_polygons_dir) + search = "*" + file_format + objs = [] + gt_geojsons = glob.glob(search) + for geojson in tqdm(gt_geojsons): + proposal_poly = os.path.join(proposal_polygons_dir, geojson) + if os.path.exists(proposal_poly): + proposal_gdf = gpd.read_file(proposal_poly) + ground_truth_gdf = gpd.read_file(geojson) + for index, row in (ground_truth_gdf.iterrows()): + objs.append(row[gt_cat_attrib]) + gt_objs = list(set(objs)) + all_objs = gt_objs + prop_objs + all_objs = list(set(all_objs)) + return prop_objs, gt_objs, all_objs + + +def precision_calc(proposal_polygons_dir, gt_polygons_dir, + prediction_cat_attrib="class", gt_cat_attrib='make', confidence_attrib=None, + object_subset=[], threshold=0.5, file_format="geojson"): + """ Using the proposal and ground truth polygons, calculate precision metrics. + Filenames of predictions and ground-truth must be identical. Will only + calculate metric for classes that exist in the ground truth. + Arguments + --------- + proposal_polygons_dir : str + The path that contains any model proposal polygons + gt_polygons_dir : str + The path that contains the ground truth polygons + prediction_cat_attrib : str + The column or attribute within the predictions that specifies + unique classes + gt_cat_attrib : str + The column or attribute within the ground truth that + specifies unique classes + confidence_attrib : str + The column or attribute within the proposal polygons that + specifies model confidence for each prediction + object_subset : list + A list or subset of the unique objects that are contained within the + ground truth polygons. If empty, this will be + auto-created using all classes that appear ground truth polygons. + threshold : float + A value between 0.0 and 1.0 that determines the IOU threshold for a + true positve. + file_format : str + The extension or file format for predictions + Returns + --------- + iou_holder : list of lists + An iou score for each object per class (precision specific) + precision_by_class : list + A list containing the precision score for each class + mPrecision : float + The mean precision score of precision_by_class + confidences : list of lists + All confidences for each object for each class + """ + ious = [] + dir = os.chdir(proposal_polygons_dir) + search = "*" + file_format + proposal_geojsons = glob.glob(search) + iou_holder = [] + confidences = [] + if len(object_subset) == 0: + prop_objs, object_subset, all_objs = get_all_objects( + proposal_polygons_dir, gt_polygons_dir, + prediction_cat_attrib=prediction_cat_attrib, + gt_cat_attrib=gt_cat_attrib, file_format=file_format) + for i in range(len(object_subset)): + iou_holder.append([]) + confidences.append([]) + + for geojson in tqdm(proposal_geojsons): + ground_truth_poly = os.path.join(gt_polygons_dir, geojson) + if os.path.exists(ground_truth_poly): + ground_truth_gdf = gpd.read_file(ground_truth_poly) + proposal_gdf = gpd.read_file(geojson) + i = 0 + for obj in object_subset: + conf_holder = [] + proposal_gdf2 = proposal_gdf[proposal_gdf[prediction_cat_attrib] == obj] + for index, row in (proposal_gdf2.iterrows()): + if confidence_attrib is not None: + conf_holder.append(row[confidence_attrib]) + iou_GDF = calculate_iou(row.geometry, ground_truth_gdf) + if 'iou_score' in iou_GDF.columns: + iou = iou_GDF.iou_score.max() + max_iou_row = iou_GDF.loc[iou_GDF['iou_score'].idxmax(axis=0, skipna=True)] + id_1 = row[prediction_cat_attrib] + id_2 = ground_truth_gdf.loc[max_iou_row.name][gt_cat_attrib] + if id_1 == id_2: + ious.append(iou) + ground_truth_gdf.drop(max_iou_row.name, axis=0, inplace=True) + else: + iou = 0 + ious.append(iou) + else: + iou = 0 + ious.append(iou) + for item in ious: + iou_holder[i].append(item) + if confidence_attrib is not None: + for conf in conf_holder: + confidences[i].append(conf) + ious = [] + i += 1 + else: + print("Warning- No ground truth for:", geojson) + proposal_gdf = gpd.read_file(geojson) + i = 0 + + for obj in object_subset: + proposal_gdf2 = proposal_gdf[proposal_gdf[gt_cat_attrib] == obj] + for z in range(len(proposal_gdf2)): + ious.append(0) + for item in ious: + iou_holder[i].append(item) + if confidence_attrib is not None: + for conf in conf_holder: + confidences[i].append(conf) + i += 1 + ious = [] + precision_by_class = average_score_by_class(iou_holder, threshold=threshold) + precision_by_class = list(np.nan_to_num(precision_by_class)) + mPrecision = np.nanmean(precision_by_class) + print("mPrecision:", mPrecision) + + return iou_holder, precision_by_class, mPrecision, confidences +