(images_path, output_path, onnx_file, labels_file, sigmoid_threshold=0.8)
| 79 | return label_json |
| 80 | |
| 81 | def segmentImages(images_path, output_path, onnx_file, labels_file, sigmoid_threshold=0.8): |
| 82 | # check if the onnx network exists |
| 83 | if(not os.path.exists(onnx_file)): |
| 84 | # download the onnx network |
| 85 | import urllib.request |
| 86 | url = "https://github.com/eokeeffe/UAV_Aerial_Segmentation_cpp_onnx/raw/refs/heads/main/networks/aerial_segmentation.onnx" |
| 87 | if not os.path.isabs(onnx_file): |
| 88 | onnx_file = os.path.join(os.path.dirname(os.path.abspath(__file__)), onnx_file) |
| 89 | print(f"Downloading segmentation model to {onnx_file}...") |
| 90 | urllib.request.urlretrieve(url, onnx_file) |
| 91 | |
| 92 | # load the onnx network |
| 93 | ort_session = ort.InferenceSession(onnx_file) |
| 94 | |
| 95 | # get the image locations |
| 96 | all_images = os.listdir(images_path) |
| 97 | |
| 98 | # create the output folder if it doesn't exist |
| 99 | Path(output_path).mkdir(parents=True, exist_ok=True) |
| 100 | |
| 101 | # segment each image |
| 102 | print("Starting segmentation ...") |
| 103 | for image in tqdm(all_images): |
| 104 | input_image = os.path.join(images_path, image) |
| 105 | output_image = os.path.join(output_path, os.path.splitext(image)[0] + '.mask.png') |
| 106 | |
| 107 | if(not os.path.exists(input_image)): |
| 108 | print(input_image," doesn't exist") |
| 109 | continue |
| 110 | if(os.path.exists(output_image)): |
| 111 | print(output_image," already exists") |
| 112 | continue |
| 113 | |
| 114 | # format the image to the correct dimensions |
| 115 | preprocessed_image,h,w = loadImage(input_image) |
| 116 | # run the inference |
| 117 | outputs = ort_session.run(["sigmoid"], {'image': preprocessed_image})[0] |
| 118 | # process the output to classified pixels |
| 119 | classified_image = extractSegmentedImage(outputs, h, w, sigmoid_threshold=sigmoid_threshold) |
| 120 | # save the segmented image |
| 121 | cv2.imwrite(output_image, classified_image) |
| 122 | |
| 123 | # save a json file with the pixel value to label relationship |
| 124 | if labels_file is not None: |
| 125 | if not os.path.isabs(labels_file): |
| 126 | labels_file = os.path.join(output_path, labels_file) |
| 127 | with open(labels_file, "w") as outfile: |
| 128 | json.dump(createPxielLabels(), outfile) |
| 129 | |
| 130 | ort_session = None |
| 131 | print("... segmentation completed!") |
| 132 | |
| 133 | if __name__=="__main__": |
| 134 | parser = argparse.ArgumentParser() |
no test coverage detected