Yolo class (#1790)

* Modify write_yolo_yaml to accept class names

* Refactor YOLOv8 script for improved clarity

* Modify YOLO training XML for input formats and parameters

Updated input image formats and corrected parameter attributes.

* Remove class_name parameter from YOLO prediction

Removed class_name parameter from inputs and tests for YOLO prediction.

* Update YOLO tool version in XML configuration

* Update YOLO training tool version for galaxy6

* Enhance input image parameter for TIFF support

Updated input image parameter to support TIFF format and added a validator for RGB channel count.

* Clean up yolo_prediction.xml by removing comments

Removed commented-out test cases for segment and detect modes.

* Enhance input_images parameter with help and validation

Updated input_images parameter to include a help message and validation for RGB channels.

* Update yolo_prediction.xml

* Fix help text for input images in yolo_training.xml

* Revise input images parameter help and validation

Updated help text and validation for input images in YOLO training configuration.

* Fix validator expression for input image formats

* Modify input_images parameter for TIFF support

Updated input_images parameter to support TIFF files with 1 or 3 channels and improved validation message.

* Apply suggestion from @bgruening

---------

Co-authored-by: Björn Grüning <bjoern@gruenings.eu>

Author: TuturBaba

tools/image_processing/yolo/preprocessing.py

@@ -38,23 +38,27 @@ def copy_pairs(pairs, image_src, label_src, image_dst, label_dst):
         copy_file(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 write_yolo_yaml(output_dir, classes):
+    yaml_path = os.path.join(output_dir, "yolo.yml")
+    with open(yaml_path, "w") as f:
+        f.write("train: train/images\n")
+        f.write("val: valid/images\n")
+        f.write("test: test/images\n\n")
+        f.write(f"nc: {len(classes)}\n\n")
+        f.write("names:\n")
+        for i, name in enumerate(classes):
+            f.write(f"  {i}: {name}\n")
+    print("\n--- YAML CONTENT ---")
+    with open(yaml_path, "r") as f:
+        print(f.read())
 
 
 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("-c", "--classes", required=False, default="")
     parser.add_argument("-p", "--train_percent", type=int, default=70)
     args = parser.parse_args()
 
@@ -68,8 +72,8 @@ def main():
     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)
+    classes_list = [c.strip() for c in args.classes.split(",")]
+    write_yolo_yaml(args.output, classes_list)
 
 
 if __name__ == "__main__":

tools/image_processing/yolo/yolo_prediction.xml

@@ -1,4 +1,4 @@
-<tool id="yolo_predict" name="Perform YOLO image labeling" version="@TOOL_VERSION@+galaxy@VERSION_SUFFIX@" profile="24.2">
+<tool id="yolo_predict" name="Perform YOLO image labeling" version="@TOOL_VERSION@+galaxy5" profile="24.2">
     <description>with ultralytics</description>
     <macros>
         <import>macros.xml</import>
@@ -16,8 +16,6 @@
           ln -s '$filename' './input_images/${filename.element_identifier}.${filename.ext}' &&
       #end for
 
-      cp '$class_name' './models/class_name.txt' &&
-
       cp '$model' './models/model.pt' &&
 	    
       python '$__tool_directory__/yolov8.py' 
@@ -29,26 +27,25 @@
 	    --image_size='$image_size'
 	    --mode='$mode'
 	    --foldername='overlaid_images'
-	    --class_names_file='$class_name'
-	    --num_classes=`wc -l < ./models/class_name.txt`
 	    --confidence='$confidence'
 	    --iou='$iou'
 	    --max_det='$max_det' 
 	    --headless
     ]]>
     </command>
     <inputs>
-        <param name="input_images" type="data" format="jpg,png" multiple="true"  label="Input images"/>
-	<param name="class_name" type="data" format="txt" label="YOLO class name file" />
-	<param name="model" type="data" format="data" label="Model file" />
-	<param name="mode" type="select" label="Prediction mode">
-          <option value="segment">segment</option>
-          <option value="detect">detect</option>
-	</param>
-	<param name="image_size" type="integer" value="1000" min="16" label="Image size" help="All input images will be re-sized to squares with sides of this length (in pixels). This value governs the trade-offs of speed (smaller values) vs accuracy (larger values)." />
- 	<param name="confidence" type="float" value="0.5" min="0.0" max="1.0" label="Confidence" help="Confidence value (0-1) for each detected bounding box." />
-	<param name="iou" type="float" value="0.7" min="0.1" max="1.0" label="IoU" help="Intersection over Union threshold for non-maximum suppression." />
-	<param name="max_det" type="integer" value="300" min="100" max="1000" label="Max. number of detections" help="Maximum number of detections allowed per image. Limits the total number of objects the model can detect in a single inference, preventing excessive outputs in dense scenes." />
+        <param name="input_images" type="data" format="jpg,png,tiff,tif" multiple="true"  label="Input images">
+            <validator type="expression" message="PNG/JPG acceptés, TIFF uniquement si 3 canaux (RGB)">value is not None and (value.ext in ('png', 'jpg', 'jpeg') or (value.ext in ('tiff', 'tif') and getattr(value.metadata, 'channels', 0) == 3))</validator>
+        </param>
+        <param name="model" type="data" format="data" label="Model file" />
+        <param name="mode" type="select" label="Prediction mode">
+            <option value="segment">segment</option>
+            <option value="detect">detect</option>
+        </param>
+        <param name="image_size" type="integer" value="1000" min="16" label="Image size" help="All input images will be re-sized to squares with sides of this length (in pixels). This value governs the trade-offs of speed (smaller values) vs accuracy (larger values)." />
+        <param name="confidence" type="float" value="0.5" min="0.0" max="1.0" label="Confidence" help="Confidence value (0-1) for each detected bounding box." />
+        <param name="iou" type="float" value="0.7" min="0.1" max="1.0" label="IoU" help="Intersection over Union threshold for non-maximum suppression." />
+        <param name="max_det" type="integer" value="300" min="100" max="1000" label="Max. number of detections" help="Maximum number of detections allowed per image. Limits the total number of objects the model can detect in a single inference, preventing excessive outputs in dense scenes." />
     </inputs>
     <outputs>
         <collection name="txt_results" format="txt" type="list" label="YOLO bounding box and annotation (text)">
@@ -71,16 +68,13 @@
             <param name="iou" value="0.7" />
             <param name="max_det" value="100" />
             <param name="mode" value="segment" />
-            <param name="class_name" value="class_name.txt" />
-
             <output_collection name="txt_results" type="list" count="1">
                 <element name="pred-test01.jpg">
                     <assert_contents>
                         <has_n_lines n="100"/>
                     </assert_contents>
                 </element>
             </output_collection>
-
             <output_collection name="mask_results" type="list" count="1">
                 <element name="pred-test01.jpg">
                     <assert_contents>
@@ -99,14 +93,9 @@
                     </assert_contents>
 		</element>
             </output_collection>
-
         </test>    
-        <!--new version's test-->
-	<!-- SEGMENT MODE TEST -->
-
         <test expect_num_outputs="3">
             <param name="input_images" value="bus.jpg" />
-            <param name="class_name" value="yolo-test-classes.txt" />
             <param name="model" location="https://github.com/ultralytics/assets/releases/download/v8.3.0/yolov8n-seg.pt" />
             <param name="mode" value="segment" />
             <param name="image_size" value="640" />
@@ -123,12 +112,8 @@
                 <element name="bus.jpg" file="bus_seg.jpg" compare="sim_size" />
             </output_collection>
         </test>
-
-	<!-- DETECT MODE TEST -->
-	
         <test expect_num_outputs="2">
             <param name="input_images" value="bus.jpg" />
-            <param name="class_name" value="yolo-test-classes.txt" />
             <param name="model" location="https://github.com/ultralytics/assets/releases/download/v8.3.0/yolov8n-seg.pt" />
             <param name="mode" value="detect" />
             <param name="image_size" value="640" />

tools/image_processing/yolo/yolo_training.xml

@@ -32,6 +32,7 @@
         -i './input_images' 
         -y './input_yolo' 
         -o ./training 
+        -c '$class_names'
         -p '$num_train' &&
 
       #if $model_choice.pretrained == 'personal'
@@ -43,7 +44,7 @@
       python '$__tool_directory__/yolov8.py' 
 		    --train --yaml_path='./training/yolo.yml' 
 		    --model_path='./models' 
-     		    --model_name='yolo_model'
+     		--model_name='yolo_model'
 		    --run_dir='./runs' 
 		    --image_size='$training_params.image_size'
 		    --epochs='$training_params.epochs'
@@ -66,7 +67,10 @@
     ]]>
     </command>
     <inputs>
-        <param name="input_images" type="data" format="jpg" multiple="true"  label="Input images"/>
+        <param name="input_images" type="data" format="jpg,png,tiff" multiple="true" label="Input Images" 
+            help="Upload the images to use for training (JPG, PNG, or TIFF format). or TIFF files, only 3-channel RGB images are supported. Each image must have a corresponding YOLO annotation file with the same base filename." >
+            <validator type="expression" message="PNG/JPG acceptés, TIFF uniquement si 3 canaux (RGB)">value is not None and (value.ext in ('png', 'jpg') or (value.ext in ('tiff') and getattr(value.metadata, 'channels', 0) == 3))</validator>
+        </param>
 	    <param name="input_yolo" type="data" format="txt" multiple="true"  label="Input YOLO txt files" help="The YOLO text files, each text file must correspond to one input image (same name, different extension)." />
         <conditional name="model_choice">
             <param name="pretrained" type="select" label="Choose pretrained YOLO model or your own">
@@ -82,6 +86,7 @@
             <when value="https://github.com/ultralytics/assets/releases/download/v8.3.0/yolo11n.pt"/>
             <when value="https://github.com/ultralytics/assets/releases/download/v8.3.0/yolov8n.pt"/>
         </conditional>
+        <param name="class_names" type="text" optional="true" label="List class name comma separated"/>
 		<section name="training_params" title="Training Parameters">
 	        <param name="num_train" type="integer" value="70" min="50" max="90" label="How do you want to split your images for training?" help="Proportion of training images wrt valid and test images (in percents)." />
 	        <param name="epochs" type="integer" value="100" min="10" max="1000" label="Number of epochs for taining." />
@@ -141,7 +146,7 @@
                         <has_size min="10000"/>
                         <has_image_channels channels="4"/>
                         <has_image_height height="1200"/>
-						<has_image_width height="3600"/>						
+						<has_image_width width="3600"/>						
                     </assert_contents>
                 </output>
         </test>
@@ -181,7 +186,7 @@
                         <has_size min="10000"/>
                         <has_image_channels channels="4"/>
                         <has_image_height height="1200"/>
-						<has_image_width height="3600"/>	
+						<has_image_width width="3600"/>	
                     </assert_contents>
                 </output>
         </test>
@@ -221,7 +226,7 @@
                         <has_size min="10000"/>
                         <has_image_channels channels="4"/>
                         <has_image_height height="1200"/>
-						<has_image_width height="3600"/>	
+						<has_image_width width="3600"/>	
                     </assert_contents>
                 </output>
         </test>

tools/image_processing/yolo/yolov8.py

@@ -18,9 +18,9 @@
 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\
+        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 \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)
+        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 """, formatter_class=RawTextHelpFormatter)
 parser.add_argument("--dir_path",
                     help=(
                         "Path to the training data directory."
@@ -76,9 +76,6 @@
 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=(
@@ -299,7 +296,7 @@ def predict(model, source_datapath, **kwargs):
         iou_value = 0.5
 
     if "num_classes" in kwargs:
-        class_array = list(range(kwargs['num_classes']))
+        class_array = list(model.names.keys())
     else:
         class_array = [0, 1]
 
@@ -373,14 +370,7 @@ def save_yolo_bounding_boxes_to_txt(predictions, save_dir):
         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 = 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'))
@@ -487,12 +477,6 @@ def save_yolo_bounding_boxes_to_txt(predictions, save_dir):
                 print(colored(f"TYX mask stack saved as : '{mask_save_as}'", 'magenta'))
             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