# HG changeset patch # User bgruening # Date 1749813815 0 # Node ID 252fd085940d15349a70a0346520e606358a27f2 planemo upload for repository https://github.com/bgruening/galaxytools/tree/master/tools commit 67e0e1d123bcfffb10bab8cc04ae67259caec557 diff -r 000000000000 -r 252fd085940d json2yolosegment.py --- /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() diff -r 000000000000 -r 252fd085940d json2yolosegment.xml --- /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 @@ + + with ultralytics + + macros.xml + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff -r 000000000000 -r 252fd085940d macros.xml --- /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 @@ + + 8.3.0 + 0 + + + + + + + + + operation_3443 + + + + + quay.io/bgruening/docker-ultralytics:@TOOL_VERSION@ + + + + + + @software{ultralytics_yolov8, + author = {Ultralytics team}, + title = {YOLOv8}, + url = {https://github.com/ultralytics/ultralytics} + } + + + @software{yolo11_ultralytics, + author = {Ultralytics team}, + title = {YOLO11}, + url = {https://github.com/ultralytics/ultralytics}, + } + + + @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}, + } + + + + + diff -r 000000000000 -r 252fd085940d preprocessing.py --- /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() diff -r 000000000000 -r 252fd085940d test-data/class_name.txt --- /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 diff -r 000000000000 -r 252fd085940d test-data/class_names.txt --- /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 diff -r 000000000000 -r 252fd085940d test-data/in_json.json --- /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": "" +} diff -r 000000000000 -r 252fd085940d test-data/in_json.txt --- /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 diff -r 000000000000 -r 252fd085940d test-data/in_json1.json --- /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": "" +} diff -r 000000000000 -r 252fd085940d test-data/in_json1.txt --- /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 diff -r 000000000000 -r 252fd085940d test-data/jpegs/0001.jpg Binary file test-data/jpegs/0001.jpg has changed diff -r 000000000000 -r 252fd085940d test-data/jpegs/0003.jpg Binary file test-data/jpegs/0003.jpg has changed diff -r 000000000000 -r 252fd085940d test-data/pred-test01.jpg Binary file test-data/pred-test01.jpg has changed diff -r 000000000000 -r 252fd085940d test-data/pred-test01.png Binary file test-data/pred-test01.png has changed diff -r 000000000000 -r 252fd085940d test-data/results_plot.png Binary file test-data/results_plot.png has changed diff -r 000000000000 -r 252fd085940d yolov8.py --- /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"))