Mercurial > repos > bgruening > json2yolosegment
changeset 0:252fd085940d draft
planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557
| author | bgruening |
|---|---|
| date | Fri, 13 Jun 2025 11:23:35 +0000 |
| parents | |
| children | dfda27273ead |
| files | json2yolosegment.py json2yolosegment.xml macros.xml preprocessing.py test-data/class_name.txt test-data/class_names.txt test-data/in_json.json test-data/in_json.txt test-data/in_json1.json test-data/in_json1.txt test-data/jpegs/0001.jpg test-data/jpegs/0003.jpg test-data/pred-test01.jpg test-data/pred-test01.png test-data/results_plot.png yolov8.py |
| diffstat | 16 files changed, 1396 insertions(+), 0 deletions(-) [+] |
line wrap: on
line diff
--- /dev/null Thu Jan 01 00:00:00 1970 +0000 +++ b/json2yolosegment.py Fri Jun 13 11:23:35 2025 +0000 @@ -0,0 +1,64 @@ +import argparse +import json +import os + + +def convert_json_to_yolo(input_dir, save_dir, class_names_file): + + with open(class_names_file, 'r') as f: + class_names = [line.strip() for line in f.readlines()] + + class_to_index = {class_name: i for i, class_name in enumerate(class_names)} + + for filename in os.listdir(input_dir): + json_filepath = os.path.join(input_dir, filename) + + with open(json_filepath, 'r') as f: + data = json.load(f) + + image_width = data.get('imageWidth') + image_height = data.get('imageHeight') + if image_width is None or image_height is None: + print(f"Skipping {filename}: missing image dimensions.") + return + + annotations = data.get('shapes', []) + + base, _ = os.path.splitext(filename) + output_file = f"{base}.txt" + output_filepath = os.path.join(save_dir, output_file) + + with open(output_filepath, 'w') as f: + for annotation in annotations: + label = annotation.get('label') + class_index = class_to_index[label] + + points = annotation.get('points', []) + + if not points: + print(f"No points found for annotation '{label}', skipping.") + continue + + x = [point[0] / (image_width - 1) for point in points] + y = [point[1] / (image_height - 1) for point in points] + + segmentation_points = ['{} {}'.format(x[i], y[i]) for i in range(len(x))] + segmentation_points_string = ' '.join(segmentation_points) + line = '{} {}\n'.format(class_index, segmentation_points_string) + f.write(line) + + print(f"Converted annotations saved to: {output_filepath}") + + +def main(): + parser = argparse.ArgumentParser(description="Convert JSON annotations to YOLO segment format.") + parser.add_argument('-i', '--input_dir', type=str, help='Full path of the folder containing AnyLabeling JSON files.') + parser.add_argument('-o', '--save_dir', type=str, help='Path to the directory to save converted YOLO files.') + parser.add_argument('-c', '--class_names_file', type=str, help='Path to the text file containing class names, one per line.') + args = parser.parse_args() + + convert_json_to_yolo(args.input_dir, args.save_dir, args.class_names_file) + + +if __name__ == "__main__": + main()
--- /dev/null Thu Jan 01 00:00:00 1970 +0000 +++ b/json2yolosegment.xml Fri Jun 13 11:23:35 2025 +0000 @@ -0,0 +1,56 @@ +<tool id="json2yolosegment" name="Convert AnyLabeling JSON to YOLO text" version="@TOOL_VERSION@+galaxy@VERSION_SUFFIX@" profile="24.2"> + <description>with ultralytics</description> + <macros> + <import>macros.xml</import> + </macros> + <expand macro="creator" /> + <expand macro="edam" /> + <command detect_errors="aggressive"> + <![CDATA[ + mkdir ./input ./output && + + #for $filename in $in_json: + ln -s '$filename' './input/${filename.element_identifier}' && + #end for + + python '$__tool_directory__/json2yolosegment.py' -i ./input/ -o ./output -c '$class_name' + + ]]> + </command> + <inputs> + <param name="in_json" type="data" format="json" multiple="true" label="Input label files" help="The label file is the output of AnyLabeling in JSON format."/> + <param name="class_name" type="data" format="txt" label="Class file" help="Class names text file, contains the names of classes, one class per row."/> + </inputs> + <outputs> + <collection name="output_yolo" type="list" label="YOLO text files"> + <discover_datasets pattern="(?P<name>.+)\.txt" visible="true" directory="./output" format="txt"/> + </collection> + </outputs> + <tests> + <test> + <param name="in_json" value="in_json.json,in_json1.json" /> + <param name="class_name" value="class_names.txt" /> + <output_collection name="output_yolo"> + <element name="in_json"> + <assert_contents> + <has_n_lines n="1" /> + <has_text text="0.8532289628180039" /> + </assert_contents> + </element> + <element name="in_json1"> + <assert_contents> + <has_n_lines n="1" /> + <has_text text="0.7710371819960861" /> + </assert_contents> + </element> + </output_collection> + </test> + </tests> + <help> +<![CDATA[ + **What it does** + This tool converts label files generated by AnyLabeling into YOLO supported text files. +]]> + </help> + <expand macro="citations" /> +</tool>
--- /dev/null Thu Jan 01 00:00:00 1970 +0000 +++ b/macros.xml Fri Jun 13 11:23:35 2025 +0000 @@ -0,0 +1,46 @@ +<macros> + <token name="@TOOL_VERSION@">8.3.0</token> + <token name="@VERSION_SUFFIX@">0</token> + <xml name="creator"> + <creator> + <person name="Yi Sun" email="sunyi000@gmail.com" /> + <organization name="NFDI4Bioimage" url="https://nfdi4bioimage.de/" /> + </creator> + </xml> + <xml name="edam"> + <edam_operations> + <edam_operation>operation_3443</edam_operation> + </edam_operations> + </xml> + <xml name="requirements"> + <requirements> + <container type="docker">quay.io/bgruening/docker-ultralytics:@TOOL_VERSION@</container> + </requirements> + </xml> + <xml name="citations"> + <citations> + <citation type="bibtex"> + @software{ultralytics_yolov8, + author = {Ultralytics team}, + title = {YOLOv8}, + url = {https://github.com/ultralytics/ultralytics} + } + </citation> + <citation type="bibtex"> + @software{yolo11_ultralytics, + author = {Ultralytics team}, + title = {YOLO11}, + url = {https://github.com/ultralytics/ultralytics}, + } + </citation> + <citation type="bibtex"> + @misc{embl2025tail, + author = {EMBL Bioimage Informatics Team}, + title = {Tail Analysis Workflow}, + url = {https://git.embl.de/grp-cba/tail-analysis/-/blob/main/analysis_workflow.md}, + } + </citation> + </citations> + </xml> +</macros> +
--- /dev/null Thu Jan 01 00:00:00 1970 +0000 +++ b/preprocessing.py Fri Jun 13 11:23:35 2025 +0000 @@ -0,0 +1,68 @@ +import argparse +import os +import shutil + +from sklearn.model_selection import train_test_split + + +def get_basename(f): + return os.path.splitext(os.path.basename(f))[0] + + +def pair_files(images_dir, labels_dir): + + img_files = [f for f in os.listdir(images_dir)] + lbl_files = [f for f in os.listdir(labels_dir)] + + image_dict = {get_basename(f): f for f in img_files} + label_dict = {get_basename(f): f for f in lbl_files} + + keys = sorted(set(image_dict) & set(label_dict)) + + return [(image_dict[k], label_dict[k]) for k in keys] + + +def copy_pairs(pairs, image_src, label_src, image_dst, label_dst): + os.makedirs(image_dst, exist_ok=True) + os.makedirs(label_dst, exist_ok=True) + for img, lbl in pairs: + shutil.copy(os.path.join(image_src, img), os.path.join(image_dst, img)) + shutil.copy(os.path.join(label_src, lbl), os.path.join(label_dst, lbl)) + + +def write_yolo_yaml(output_dir): + + yolo_yaml_path = os.path.join(output_dir, "yolo.yml") + with open(yolo_yaml_path, 'w') as f: + f.write(f"path: {output_dir}\n") + f.write("train: train\n") + f.write("val: valid\n") + f.write("test: test\n") + f.write("\n") + f.write("names: ['dataset']\n") + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("-i", "--images", required=True) + parser.add_argument("-y", "--labels", required=True) + parser.add_argument("-o", "--output", required=True) + parser.add_argument("-p", "--train_percent", type=int, default=70) + args = parser.parse_args() + + all_pairs = pair_files(args.images, args.labels) + train_size = args.train_percent / 100.0 + val_test_size = 1.0 - train_size + + train_pairs, val_test_pairs = train_test_split(all_pairs, test_size=val_test_size, random_state=42) + val_pairs, test_pairs = train_test_split(val_test_pairs, test_size=0.5, random_state=42) + + copy_pairs(train_pairs, args.images, args.labels, os.path.join(args.output, "train/images"), os.path.join(args.output, "train/labels")) + copy_pairs(val_pairs, args.images, args.labels, os.path.join(args.output, "valid/images"), os.path.join(args.output, "valid/labels")) + copy_pairs(test_pairs, args.images, args.labels, os.path.join(args.output, "test/images"), os.path.join(args.output, "test/labels")) + + write_yolo_yaml(args.output) + + +if __name__ == "__main__": + main()
--- /dev/null Thu Jan 01 00:00:00 1970 +0000 +++ b/test-data/class_name.txt Fri Jun 13 11:23:35 2025 +0000 @@ -0,0 +1,1 @@ +tail
--- /dev/null Thu Jan 01 00:00:00 1970 +0000 +++ b/test-data/class_names.txt Fri Jun 13 11:23:35 2025 +0000 @@ -0,0 +1,1 @@ +animal
--- /dev/null Thu Jan 01 00:00:00 1970 +0000 +++ b/test-data/in_json.json Fri Jun 13 11:23:35 2025 +0000 @@ -0,0 +1,332 @@ +{ + "version": "0.4.10", + "flags": {}, + "shapes": [ + { + "label": "animal", + "text": "", + "points": [ + [ + 394.0, + 205.0 + ], + [ + 388.0, + 206.0 + ], + [ + 376.0, + 201.0 + ], + [ + 354.0, + 202.0 + ], + [ + 348.0, + 198.0 + ], + [ + 332.0, + 195.0 + ], + [ + 317.0, + 199.0 + ], + [ + 296.0, + 214.0 + ], + [ + 273.0, + 218.0 + ], + [ + 241.0, + 249.0 + ], + [ + 218.0, + 288.0 + ], + [ + 208.0, + 315.0 + ], + [ + 191.0, + 395.0 + ], + [ + 191.0, + 442.0 + ], + [ + 181.0, + 481.0 + ], + [ + 180.0, + 510.0 + ], + [ + 175.0, + 538.0 + ], + [ + 174.0, + 580.0 + ], + [ + 170.0, + 598.0 + ], + [ + 171.0, + 611.0 + ], + [ + 167.0, + 618.0 + ], + [ + 163.0, + 675.0 + ], + [ + 158.0, + 682.0 + ], + [ + 158.0, + 691.0 + ], + [ + 147.0, + 712.0 + ], + [ + 143.0, + 739.0 + ], + [ + 149.0, + 746.0 + ], + [ + 150.0, + 752.0 + ], + [ + 160.0, + 761.0 + ], + [ + 165.0, + 773.0 + ], + [ + 172.0, + 779.0 + ], + [ + 183.0, + 796.0 + ], + [ + 200.0, + 811.0 + ], + [ + 206.0, + 814.0 + ], + [ + 209.0, + 812.0 + ], + [ + 215.0, + 817.0 + ], + [ + 220.0, + 816.0 + ], + [ + 227.0, + 819.0 + ], + [ + 233.0, + 817.0 + ], + [ + 270.0, + 817.0 + ], + [ + 273.0, + 813.0 + ], + [ + 288.0, + 811.0 + ], + [ + 315.0, + 800.0 + ], + [ + 318.0, + 796.0 + ], + [ + 321.0, + 797.0 + ], + [ + 333.0, + 777.0 + ], + [ + 338.0, + 773.0 + ], + [ + 345.0, + 753.0 + ], + [ + 361.0, + 726.0 + ], + [ + 363.0, + 707.0 + ], + [ + 369.0, + 694.0 + ], + [ + 373.0, + 673.0 + ], + [ + 377.0, + 671.0 + ], + [ + 380.0, + 663.0 + ], + [ + 382.0, + 642.0 + ], + [ + 385.0, + 639.0 + ], + [ + 386.0, + 631.0 + ], + [ + 384.0, + 598.0 + ], + [ + 387.0, + 579.0 + ], + [ + 386.0, + 536.0 + ], + [ + 381.0, + 518.0 + ], + [ + 382.0, + 474.0 + ], + [ + 389.0, + 450.0 + ], + [ + 394.0, + 410.0 + ], + [ + 418.0, + 346.0 + ], + [ + 423.0, + 341.0 + ], + [ + 425.0, + 330.0 + ], + [ + 429.0, + 326.0 + ], + [ + 430.0, + 316.0 + ], + [ + 434.0, + 309.0 + ], + [ + 434.0, + 298.0 + ], + [ + 437.0, + 294.0 + ], + [ + 435.0, + 286.0 + ], + [ + 437.0, + 277.0 + ], + [ + 436.0, + 245.0 + ], + [ + 433.0, + 237.0 + ], + [ + 429.0, + 236.0 + ], + [ + 424.0, + 225.0 + ] + ], + "group_id": null, + "shape_type": "polygon", + "flags": {} + } + ], + "imagePath": "C2-MAX_20230629_DE_W0003_P0001-0001.tif", + "imageData": null, + "imageHeight": 1024, + "imageWidth": 512, + "text": "" +}
--- /dev/null Thu Jan 01 00:00:00 1970 +0000 +++ b/test-data/in_json.txt Fri Jun 13 11:23:35 2025 +0000 @@ -0,0 +1,1 @@ +0 0.7710371819960861 0.20039100684261973 0.7592954990215264 0.2013685239491691 0.735812133072407 0.19648093841642228 0.6927592954990215 0.19745845552297164 0.6810176125244618 0.1935483870967742 0.649706457925636 0.1906158357771261 0.6203522504892368 0.19452590420332355 0.5792563600782779 0.20918866080156404 0.5342465753424658 0.2130987292277615 0.47162426614481406 0.2434017595307918 0.42661448140900193 0.28152492668621704 0.4070450097847358 0.30791788856304986 0.37377690802348335 0.386119257086999 0.37377690802348335 0.43206256109481916 0.3542074363992172 0.4701857282502444 0.3522504892367906 0.49853372434017595 0.3424657534246575 0.5259042033235581 0.3405088062622309 0.5669599217986315 0.33268101761252444 0.5845552297165201 0.33463796477495106 0.5972629521016618 0.3268101761252446 0.6041055718475073 0.31898238747553814 0.6598240469208211 0.30919765166340507 0.6666666666666666 0.30919765166340507 0.6754643206256109 0.2876712328767123 0.6959921798631477 0.27984344422700586 0.7223851417399805 0.29158512720156554 0.729227761485826 0.29354207436399216 0.7350928641251222 0.3131115459882583 0.7438905180840665 0.32289628180039137 0.7556207233626588 0.33659491193737767 0.761485826001955 0.35812133072407043 0.7781036168132942 0.3913894324853229 0.7927663734115347 0.40313111545988256 0.7956989247311828 0.4090019569471624 0.793743890518084 0.4207436399217221 0.7986314760508308 0.43052837573385516 0.7976539589442815 0.44422700587084146 0.8005865102639296 0.45596868884540115 0.7986314760508308 0.5283757338551859 0.7986314760508308 0.5342465753424658 0.7947214076246334 0.5636007827788649 0.7927663734115347 0.6164383561643836 0.7820136852394917 0.6223091976516634 0.7781036168132942 0.6281800391389433 0.7790811339198436 0.6516634050880626 0.7595307917888563 0.6614481409001957 0.7556207233626588 0.675146771037182 0.7360703812316716 0.7064579256360078 0.7096774193548387 0.7103718199608611 0.6911045943304008 0.7221135029354208 0.678396871945259 0.7299412915851272 0.6578690127077224 0.7377690802348337 0.6559139784946236 0.7436399217221135 0.6480938416422287 0.7475538160469667 0.6275659824046921 0.7534246575342466 0.624633431085044 0.7553816046966731 0.6168132942326491 0.7514677103718199 0.5845552297165201 0.7573385518590998 0.5659824046920822 0.7553816046966731 0.5239491691104594 0.7455968688845401 0.5063538611925709 0.7475538160469667 0.4633431085043988 0.761252446183953 0.4398826979472141 0.7710371819960861 0.40078201368523947 0.8180039138943248 0.33822091886608013 0.8277886497064579 0.3333333333333333 0.8317025440313112 0.3225806451612903 0.8395303326810176 0.31867057673509286 0.8414872798434442 0.3088954056695992 0.8493150684931506 0.3020527859237537 0.8493150684931506 0.2913000977517107 0.8551859099804305 0.2873900293255132 0.8512720156555773 0.27956989247311825 0.8551859099804305 0.270772238514174 0.8532289628180039 0.23949169110459434 0.8473581213307241 0.2316715542521994 0.8395303326810176 0.23069403714565004 0.8297455968688845 0.21994134897360704
--- /dev/null Thu Jan 01 00:00:00 1970 +0000 +++ b/test-data/in_json1.json Fri Jun 13 11:23:35 2025 +0000 @@ -0,0 +1,332 @@ +{ + "version": "0.4.10", + "flags": {}, + "shapes": [ + { + "label": "animal", + "text": "", + "points": [ + [ + 394.0, + 205.0 + ], + [ + 388.0, + 206.0 + ], + [ + 376.0, + 201.0 + ], + [ + 354.0, + 202.0 + ], + [ + 348.0, + 198.0 + ], + [ + 332.0, + 195.0 + ], + [ + 317.0, + 199.0 + ], + [ + 296.0, + 214.0 + ], + [ + 273.0, + 218.0 + ], + [ + 241.0, + 249.0 + ], + [ + 218.0, + 288.0 + ], + [ + 208.0, + 315.0 + ], + [ + 191.0, + 395.0 + ], + [ + 191.0, + 442.0 + ], + [ + 181.0, + 481.0 + ], + [ + 180.0, + 510.0 + ], + [ + 175.0, + 538.0 + ], + [ + 174.0, + 580.0 + ], + [ + 170.0, + 598.0 + ], + [ + 171.0, + 611.0 + ], + [ + 167.0, + 618.0 + ], + [ + 163.0, + 675.0 + ], + [ + 158.0, + 682.0 + ], + [ + 158.0, + 691.0 + ], + [ + 147.0, + 712.0 + ], + [ + 143.0, + 739.0 + ], + [ + 149.0, + 746.0 + ], + [ + 150.0, + 752.0 + ], + [ + 160.0, + 761.0 + ], + [ + 165.0, + 773.0 + ], + [ + 172.0, + 779.0 + ], + [ + 183.0, + 796.0 + ], + [ + 200.0, + 811.0 + ], + [ + 206.0, + 814.0 + ], + [ + 209.0, + 812.0 + ], + [ + 215.0, + 817.0 + ], + [ + 220.0, + 816.0 + ], + [ + 227.0, + 819.0 + ], + [ + 233.0, + 817.0 + ], + [ + 270.0, + 817.0 + ], + [ + 273.0, + 813.0 + ], + [ + 288.0, + 811.0 + ], + [ + 315.0, + 800.0 + ], + [ + 318.0, + 796.0 + ], + [ + 321.0, + 797.0 + ], + [ + 333.0, + 777.0 + ], + [ + 338.0, + 773.0 + ], + [ + 345.0, + 753.0 + ], + [ + 361.0, + 726.0 + ], + [ + 363.0, + 707.0 + ], + [ + 369.0, + 694.0 + ], + [ + 373.0, + 673.0 + ], + [ + 377.0, + 671.0 + ], + [ + 380.0, + 663.0 + ], + [ + 382.0, + 642.0 + ], + [ + 385.0, + 639.0 + ], + [ + 386.0, + 631.0 + ], + [ + 384.0, + 598.0 + ], + [ + 387.0, + 579.0 + ], + [ + 386.0, + 536.0 + ], + [ + 381.0, + 518.0 + ], + [ + 382.0, + 474.0 + ], + [ + 389.0, + 450.0 + ], + [ + 394.0, + 410.0 + ], + [ + 418.0, + 346.0 + ], + [ + 423.0, + 341.0 + ], + [ + 425.0, + 330.0 + ], + [ + 429.0, + 326.0 + ], + [ + 430.0, + 316.0 + ], + [ + 434.0, + 309.0 + ], + [ + 434.0, + 298.0 + ], + [ + 437.0, + 294.0 + ], + [ + 435.0, + 286.0 + ], + [ + 437.0, + 277.0 + ], + [ + 436.0, + 245.0 + ], + [ + 433.0, + 237.0 + ], + [ + 429.0, + 236.0 + ], + [ + 424.0, + 225.0 + ] + ], + "group_id": null, + "shape_type": "polygon", + "flags": {} + } + ], + "imagePath": "C2-MAX_20230629_DE_W0003_P0001-0001.tif", + "imageData": null, + "imageHeight": 1024, + "imageWidth": 512, + "text": "" +}
--- /dev/null Thu Jan 01 00:00:00 1970 +0000 +++ b/test-data/in_json1.txt Fri Jun 13 11:23:35 2025 +0000 @@ -0,0 +1,1 @@ +0 0.7710371819960861 0.20039100684261973 0.7592954990215264 0.2013685239491691 0.735812133072407 0.19648093841642228 0.6927592954990215 0.19745845552297164 0.6810176125244618 0.1935483870967742 0.649706457925636 0.1906158357771261 0.6203522504892368 0.19452590420332355 0.5792563600782779 0.20918866080156404 0.5342465753424658 0.2130987292277615 0.47162426614481406 0.2434017595307918 0.42661448140900193 0.28152492668621704 0.4070450097847358 0.30791788856304986 0.37377690802348335 0.386119257086999 0.37377690802348335 0.43206256109481916 0.3542074363992172 0.4701857282502444 0.3522504892367906 0.49853372434017595 0.3424657534246575 0.5259042033235581 0.3405088062622309 0.5669599217986315 0.33268101761252444 0.5845552297165201 0.33463796477495106 0.5972629521016618 0.3268101761252446 0.6041055718475073 0.31898238747553814 0.6598240469208211 0.30919765166340507 0.6666666666666666 0.30919765166340507 0.6754643206256109 0.2876712328767123 0.6959921798631477 0.27984344422700586 0.7223851417399805 0.29158512720156554 0.729227761485826 0.29354207436399216 0.7350928641251222 0.3131115459882583 0.7438905180840665 0.32289628180039137 0.7556207233626588 0.33659491193737767 0.761485826001955 0.35812133072407043 0.7781036168132942 0.3913894324853229 0.7927663734115347 0.40313111545988256 0.7956989247311828 0.4090019569471624 0.793743890518084 0.4207436399217221 0.7986314760508308 0.43052837573385516 0.7976539589442815 0.44422700587084146 0.8005865102639296 0.45596868884540115 0.7986314760508308 0.5283757338551859 0.7986314760508308 0.5342465753424658 0.7947214076246334 0.5636007827788649 0.7927663734115347 0.6164383561643836 0.7820136852394917 0.6223091976516634 0.7781036168132942 0.6281800391389433 0.7790811339198436 0.6516634050880626 0.7595307917888563 0.6614481409001957 0.7556207233626588 0.675146771037182 0.7360703812316716 0.7064579256360078 0.7096774193548387 0.7103718199608611 0.6911045943304008 0.7221135029354208 0.678396871945259 0.7299412915851272 0.6578690127077224 0.7377690802348337 0.6559139784946236 0.7436399217221135 0.6480938416422287 0.7475538160469667 0.6275659824046921 0.7534246575342466 0.624633431085044 0.7553816046966731 0.6168132942326491 0.7514677103718199 0.5845552297165201 0.7573385518590998 0.5659824046920822 0.7553816046966731 0.5239491691104594 0.7455968688845401 0.5063538611925709 0.7475538160469667 0.4633431085043988 0.761252446183953 0.4398826979472141 0.7710371819960861 0.40078201368523947 0.8180039138943248 0.33822091886608013 0.8277886497064579 0.3333333333333333 0.8317025440313112 0.3225806451612903 0.8395303326810176 0.31867057673509286 0.8414872798434442 0.3088954056695992 0.8493150684931506 0.3020527859237537 0.8493150684931506 0.2913000977517107 0.8551859099804305 0.2873900293255132 0.8512720156555773 0.27956989247311825 0.8551859099804305 0.270772238514174 0.8532289628180039 0.23949169110459434 0.8473581213307241 0.2316715542521994 0.8395303326810176 0.23069403714565004 0.8297455968688845 0.21994134897360704
--- /dev/null Thu Jan 01 00:00:00 1970 +0000 +++ b/yolov8.py Fri Jun 13 11:23:35 2025 +0000 @@ -0,0 +1,494 @@ +import argparse +import os +import pathlib +import shutil +import time +from argparse import RawTextHelpFormatter +from collections import defaultdict + +import cv2 +import numpy as np +from termcolor import colored +from tifffile import imwrite +from ultralytics import YOLO + + +# +# Input arguments +# +parser = argparse.ArgumentParser( + description='train/predict dataset with YOLOv8', + epilog="""USAGE EXAMPLE:\n\n~~~~Prediction~~~~\n\ + python yolov8.py --test_path=/g/group/user/data --model_path=/g/cba/models --model_name=yolov8n --save_dir=/g/group/user/results --iou=0.7 --confidence=0.5 --image_size=320 --run_dir=/g/group/user/runs --foldername=batch --headless --num_classes=1 max_det=1 --class_names_file=/g/group/user/class_names.txt\n\ + \n~~~~Training~~~~ \n\ + python yolov8.py --train --yaml_path=/g/group/user/example.yaml --model_path=/g/cba/models --model_name=yolov8n --run_dir=/g/group/user/runs/ --image_size=320 --epochs=150 --scale=0.3 --hsv_v=0.5 --model_format=pt --degrees=180 --class_names_file=/g/group/user/class_names.txt""", formatter_class=RawTextHelpFormatter) +parser.add_argument("--dir_path", + help=( + "Path to the training data directory." + ), + type=str) +parser.add_argument("--yaml_path", + help=( + "YAML file with all the data paths" + " i.e. for train, test, valid data." + ), + type=str) +parser.add_argument("--test_path", + help=( + "Path to the prediction folder." + ), + type=str) +parser.add_argument("--save_dir", + help=( + "Path to the directory where bounding boxes text files" + " would be saved." + ), + type=str) +parser.add_argument("--run_dir", + help=( + "Path where overlaid images would be saved." + "For example: `RUN_DIR=projectName/results`." + "This should exist." + ), + type=str) +parser.add_argument("--foldername", + help=("Folder to save overlaid images.\n" + "For example: FOLDERNAME=batch.\n" + "This should not exist as a new folder named `batch`\n" + " will be created in RUN_DIR.\n" + " If it exists already then, a new folder named `batch1`\n" + " will be created automatically as it does not overwrite\n" + ), + type=str) + +# For selecting and loading model +parser.add_argument("--model_name", + help=("Models for task `detect` can be seen here:\n" + "https://docs.ultralytics.com/tasks/detect/#models \n\n" + "Models for task `segment` can be seen here:\n" + "https://docs.ultralytics.com/tasks/segment/#models \n\n" + " . Use `yolov8n` for `detect` tasks. " + "For custom model, use `best`" + ), + default='yolov8n', type=str) +parser.add_argument("--model_path", + help="Full absolute path to the model directory", + type=str) +parser.add_argument("--model_format", + help="Format of the YOLO model i.e pt, yaml etc.", + default='pt', type=str) +parser.add_argument("--class_names_file", + help="Path to the text file containing class names.", + type=str) + +# For training the model and prediction +parser.add_argument("--mode", + help=( + "detection, segmentation, classification, and pose \n. " + " Only detection mode available currently i.e. `detect`" + ), default='detect', type=str) +parser.add_argument('--train', + help="Do training", + action='store_true') +parser.add_argument("--confidence", + help="Confidence value (0-1) for each detected bounding box", + default=0.5, type=float) +parser.add_argument("--epochs", + help="Number of epochs for training. Default: 100", + default=100, type=int) +parser.add_argument("--init_lr", + help="Number of epochs for training. Default: 100", + default=0.01, type=float) +parser.add_argument("--weight_decay", + help="Number of epochs for training. Default: 100", + default=0.0005, type=float) + +parser.add_argument("--num_classes", + help="Number of classes to be predicted. Default: 2", + default=2, type=int) +parser.add_argument("--iou", + help="Intersection over union (IoU) threshold for NMS", + default=0.7, type=float) +parser.add_argument("--image_size", + help=("Size of input image to be used only as integer of w,h. \n" + "For training choose <= 1000. \n\n" + "Prediction will be done on original image size" + ), + default=320, type=int) +parser.add_argument("--max_det", + help=("Maximum number of detections allowed per image. \n" + "Limits the total number of objects the model can detect in a single inference, \n" + "preventing excessive outputs in dense scenes.\n\n" + ), + default=300, type=int) + +# For tracking +parser.add_argument("--tracker_file", + help=("Path to the configuration file of the tracker used. \n"), + default='bytetrack.yaml', type=str) + +# For headless operation +parser.add_argument('--headless', action='store_true') +parser.add_argument('--nextflow', action='store_true') + +# For data augmentation +parser.add_argument("--hsv_h", + help="(float) image HSV-Hue augmentation (fraction)", + default=0.015, type=float) +parser.add_argument("--hsv_s", + help="(float) image HSV-Saturation augmentation (fraction)", + default=0.7, type=float) +parser.add_argument("--hsv_v", + help="(float) image HSV-Value augmentation (fraction)", + default=0.4, type=float) +parser.add_argument("--degrees", + help="(float) image rotation (+/- deg)", + default=0.0, type=float) +parser.add_argument("--translate", + help="(float) image translation (+/- fraction)", + default=0.1, type=float) +parser.add_argument("--scale", + help="(float) image scale (+/- gain)", + default=0.5, type=float) +parser.add_argument("--shear", + help="(float) image shear (+/- deg)", + default=0.0, type=float) +parser.add_argument("--perspective", + help="(float) image perspective (+/- fraction), range 0-0.001", + default=0.0, type=float) +parser.add_argument("--flipud", + help="(float) image flip up-down (probability)", + default=0.0, type=float) +parser.add_argument("--fliplr", + help="(float) image flip left-right (probability)", + default=0.5, type=float) +parser.add_argument("--mosaic", + help="(float) image mosaic (probability)", + default=1.0, type=float) +parser.add_argument("--crop_fraction", + help="(float) crops image to a fraction of its size to " + "emphasize central features and adapt to object scales, " + "reducing background distractions", + default=1.0, type=float) + + +# +# Functions +# +# Train a new model on the dataset mentioned in yaml file +def trainModel(model_path, model_name, yaml_filepath, **kwargs): + if "imgsz" in kwargs: + image_size = kwargs['imgsz'] + else: + image_size = 320 + + if "epochs" in kwargs: + n_epochs = kwargs['epochs'] + else: + n_epochs = 100 + + if "hsv_h" in kwargs: + aug_hsv_h = kwargs['hsv_h'] + else: + aug_hsv_h = 0.015 + + if "hsv_s" in kwargs: + aug_hsv_s = kwargs['hsv_s'] + else: + aug_hsv_s = 0.7 + + if "hsv_v" in kwargs: + aug_hsv_v = kwargs['hsv_v'] + else: + aug_hsv_v = 0.4 + + if "degrees" in kwargs: + aug_degrees = kwargs['degrees'] + else: + aug_degrees = 10.0 + + if "translate" in kwargs: + aug_translate = kwargs['translate'] + else: + aug_translate = 0.1 + + if "scale" in kwargs: + aug_scale = kwargs['scale'] + else: + aug_scale = 0.2 + + if "shear" in kwargs: + aug_shear = kwargs['shear'] + else: + aug_shear = 0.0 + + if "shear" in kwargs: + aug_shear = kwargs['shear'] + else: + aug_shear = 0.0 + + if "perspective" in kwargs: + aug_perspective = kwargs['perspective'] + else: + aug_perspective = 0.0 + + if "fliplr" in kwargs: + aug_fliplr = kwargs['fliplr'] + else: + aug_fliplr = 0.5 + + if "flipud" in kwargs: + aug_flipud = kwargs['flipud'] + else: + aug_flipud = 0.0 + + if "mosaic" in kwargs: + aug_mosaic = kwargs['mosaic'] + else: + aug_mosaic = 1.0 + + if "crop_fraction" in kwargs: + aug_crop_fraction = kwargs['crop_fraction'] + else: + aug_crop_fraction = 1.0 + + if "weight_decay" in kwargs: + weight_decay = kwargs['weight_decay'] + else: + weight_decay = 1.0 + + if "init_lr" in kwargs: + init_lr = kwargs['init_lr'] + else: + init_lr = 1.0 + + train_save_path = os.path.expanduser('~/runs/' + args.mode + '/train/') + if os.path.isdir(train_save_path): + shutil.rmtree(train_save_path) + # Load a pretrained YOLO model (recommended for training) + if args.model_format == 'pt': + model = YOLO(os.path.join(model_path, model_name + "." + args.model_format)) + else: + model = YOLO(model_name + "." + args.model_format) + model.train(data=yaml_filepath, epochs=n_epochs, project=args.run_dir, + imgsz=image_size, verbose=True, hsv_h=aug_hsv_h, + hsv_s=aug_hsv_s, hsv_v=aug_hsv_v, degrees=aug_degrees, + translate=aug_translate, shear=aug_shear, scale=aug_scale, + perspective=aug_perspective, fliplr=aug_fliplr, + flipud=aug_flipud, mosaic=aug_mosaic, crop_fraction=aug_crop_fraction, + weight_decay=weight_decay, lr0=init_lr, seed=42) + return model + + +# Validate the trained model +def validateModel(model): + # Remove prediction save path if already exists + val_save_path = os.path.expanduser('~/runs/' + args.mode + '/val/') + if os.path.isdir(val_save_path): + shutil.rmtree(val_save_path) + # Validate the model + metrics = model.val() # no args needed, dataset & settings remembered + metrics.box.map # map50-95 + metrics.box.map50 # map50 + metrics.box.map75 # map75 + metrics.box.maps # a list contains map50-95 of each category + + +# Do predictions on images/videos using trained/loaded model +def predict(model, source_datapath, **kwargs): + if "imgsz" in kwargs: + image_size = kwargs['imgsz'] + else: + image_size = 320 + + if "conf" in kwargs: + confidence = kwargs['conf'] + else: + confidence = 0.5 + + if "iou" in kwargs: + iou_value = kwargs['iou'] + else: + iou_value = 0.5 + + if "num_classes" in kwargs: + class_array = list(range(kwargs['num_classes'])) + else: + class_array = [0, 1] + + if "max_det" in kwargs: + maximum_detections = args.max_det + else: + maximum_detections = 300 + + if "run_dir" in kwargs: + run_save_dir = kwargs['run_dir'] + else: + # Remove prediction save path if already exists + pred_save_path = os.path.expanduser('~/runs/' + args.mode + '/predict/') + if os.path.isdir(pred_save_path): + shutil.rmtree(pred_save_path) + if "foldername" in kwargs: + save_folder_name = kwargs['foldername'] + # infer on a local image or directory containing images/videos + prediction = model.predict(source=source_datapath, save=True, stream=True, + conf=confidence, imgsz=image_size, + save_conf=True, iou=iou_value, max_det=maximum_detections, + classes=class_array, save_txt=False, + project=run_save_dir, name=save_folder_name, verbose=True) + return prediction + + +# Save bounding boxes +def save_yolo_bounding_boxes_to_txt(predictions, save_dir): + """ + Function to save YOLO bounding boxes to text files. + Parameters: + - predictions: List of results from YOLO model inference. + - save_dir: Directory where the text files will be saved. + """ + for result in predictions: + result = result.to("cpu").numpy() + # Using bounding_boxes, confidence_scores, and class_num which are defined in the list + bounding_boxes = result.boxes.xyxy # Bounding boxes in xyxy format + confidence_scores = result.boxes.conf # Confidence scores + class_nums = result.boxes.cls # Class numbers + # Create save directory if it doesn't exist + save_path = pathlib.Path(save_dir).absolute() + save_path.mkdir(parents=True, exist_ok=True) + # Construct filename for the text file + image_filename = pathlib.Path(result.path).stem + text_filename = save_path / f"{image_filename}.txt" + # Write bounding boxes info into the text file + with open(text_filename, 'w') as f: + for i in range(bounding_boxes.shape[0]): + x1, y1, x2, y2 = bounding_boxes[i] + confidence = confidence_scores[i] + class_num = int(class_nums[i]) + f.write(f'{class_num:01} {x1:06.2f} {y1:06.2f} {x2:06.2f} {y2:06.2f} {confidence:0.02} \n') + print(colored(f"Bounding boxes saved in: {text_filename}", 'green')) + + +if __name__ == '__main__': + args = parser.parse_args() + os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE" + # Train/load model + if (args.train): + model = trainModel(args.model_path, args.model_name, args.yaml_path, + imgsz=args.image_size, epochs=args.epochs, + hsv_h=args.hsv_h, hsv_s=args.hsv_s, hsv_v=args.hsv_v, + degrees=args.degrees, translate=args.translate, + shear=args.shear, scale=args.scale, + perspective=args.perspective, fliplr=args.fliplr, + flipud=args.flipud, mosaic=args.mosaic) + validateModel(model) + else: + t = time.time() + train_save_path = os.path.expanduser('~/runs/' + args.mode + '/') + if os.path.isfile(os.path.join(train_save_path, + "train", "weights", "best.pt")) and (args.model_name == 'sam'): + model = YOLO(os.path.join(train_save_path, + "train", "weights", "best.pt")) + else: + model = YOLO(os.path.join(args.model_path, + args.model_name + ".pt")) + model.info(verbose=True) + elapsed = time.time() - t + print(colored(f"\nYOLO model loaded in : '{elapsed}' sec \n", 'white', 'on_yellow')) + + if (args.save_dir): + # Do predictions (optionally show image results with bounding boxes) + t = time.time() + datapath_for_prediction = args.test_path + # Extracting class names from the model + class_names = model.names + predictions = predict(model, datapath_for_prediction, + imgsz=args.image_size, conf=args.confidence, + iou=args.iou, run_dir=args.run_dir, + foldername=args.foldername, num_classes=args.num_classes, max_det=args.max_det) + elapsed = time.time() - t + print(colored(f"\nYOLO prediction done in : '{elapsed}' sec \n", 'white', 'on_cyan')) + + if (args.mode == "detect"): + # Save bounding boxes + save_yolo_bounding_boxes_to_txt(predictions, args.save_dir) + elif (args.mode == "track"): + results = model.track(source=datapath_for_prediction, + tracker=args.tracker_file, + conf=args.confidence, + iou=args.iou, + persist=False, + show=True, + save=True, + project=args.run_dir, + name=args.foldername) + # Store the track history + track_history = defaultdict(lambda: []) + + for result in results: + # Get the boxes and track IDs + if result.boxes and result.boxes.is_track: + boxes = result.boxes.xywh.cpu() + track_ids = result.boxes.id.int().cpu().tolist() + # Visualize the result on the frame + frame = result.plot() + # Plot the tracks + for box, track_id in zip(boxes, track_ids): + x, y, w, h = box + track = track_history[track_id] + track.append((float(x), float(y))) # x, y center point + if len(track) > 30: # retain 30 tracks for 30 frames + track.pop(0) + + # Draw the tracking lines + points = np.hstack(track).astype(np.int32).reshape((-1, 1, 2)) + cv2.polylines(frame, [points], isClosed=False, color=(230, 230, 230), thickness=2) + + # Display the annotated frame + cv2.imshow("YOLO11 Tracking", frame) + print(colored(f"Tracking results saved in : '{args.save_dir}' \n", 'green')) + elif (args.mode == "segment"): + # Read class names from the file + with open(args.class_names_file, 'r') as f: + class_names = [line.strip() for line in f.readlines()] + # Create a mapping from class names to indices + class_to_index = {class_name: i for i, class_name in enumerate(class_names)} + + # Save polygon coordinates + for result in predictions: + # Create binary mask + img = np.copy(result.orig_img) + filename = pathlib.Path(result.path).stem + b_mask = np.zeros(img.shape[:2], np.uint8) + mask_save_as = str(pathlib.Path(os.path.join(args.save_dir, filename + "_mask.tiff")).absolute()) + # Define output file path for text file + output_filename = os.path.splitext(filename)[0] + ".txt" + txt_save_as = str(pathlib.Path(os.path.join(args.save_dir, filename + ".txt")).absolute()) + + for c, ci in enumerate(result): + # Extract contour result + contour = ci.masks.xy.pop() + # Changing the type + contour = contour.astype(np.int32) + # Reshaping + contour = contour.reshape(-1, 1, 2) + # Draw contour onto mask + _ = cv2.drawContours(b_mask, [contour], -1, (255, 255, 255), cv2.FILLED) + + # Normalized polygon points + points = ci.masks.xyn.pop() + obj_class = int(ci.boxes.cls.to("cpu").numpy().item()) + confidence = result.boxes.conf.to("cpu").numpy()[c] + + with open(txt_save_as, 'a') as f: + segmentation_points = ['{} {}'.format(points[i][0], points[i][1]) for i in range(len(points))] + segmentation_points_string = ' '.join(segmentation_points) + line = '{} {} {}\n'.format(obj_class, segmentation_points_string, confidence) + f.write(line) + + imwrite(mask_save_as, b_mask, imagej=True) # save image + print(colored(f"Saved cropped image as : \n '{mask_save_as}' \n", 'magenta')) + print(colored(f"Polygon coordinates saved as : \n '{txt_save_as}' \n", 'cyan')) + + else: + raise Exception(("Currently only 'detect' and 'segment' modes are available"))