--- a/yolo_prediction.xml	Mon Jul 14 18:28:13 2025 +0000
+++ b/yolo_prediction.xml	Mon Jul 21 15:50:37 2025 +0000
@@ -13,7 +13,7 @@
       mkdir -p ./input_images ./runs ./models ./results && 
       
       #for $filename in $input_images:
-          ln -s '$filename' './input_images/${filename.element_identifier}' &&
+          ln -s '$filename' './input_images/${filename.element_identifier}.${filename.ext}' &&
       #end for
 
       cp '$class_name' './models/class_name.txt' &&
@@ -52,46 +52,18 @@
     </inputs>
     <outputs>
         <collection name="txt_results" format="txt" type="list" label="YOLO bounding box and annotation (text)">
-            <discover_datasets pattern="(?P&lt;name&gt;.*)\.txt$" directory="results" visible="true"/>
+            <discover_datasets pattern="(?P&lt;designation&gt;.+)\.txt$" directory="results" />
         </collection>
         <collection name="mask_results" format="tiff" type="list" label="YOLO segmentation masks (TIFF)">
-            <discover_datasets pattern="(?P&lt;name&gt;.*)_mask\.tiff$" directory="results" visible="true"/>
+            <filter>mode == 'segment'</filter>
+            <discover_datasets pattern="(?P&lt;designation&gt;.+)_mask\.tiff$" directory="results" />
         </collection>
         <collection name="overlay_results" format="jpg" type="list" label="YOLO annotated images">
-            <discover_datasets pattern="(?P&lt;name&gt;.*)\.jpg$" directory="runs/overlaid_images" visible="true"/>
+            <discover_datasets pattern="(?P&lt;designation&gt;.+)\.jpg$" directory="runs/overlaid_images" />
         </collection>
     </outputs>
     <tests>
-    <!--tests from first version-->
-        <test>
-        <param name="input_images" value="pred-test01.jpg" />
-        <param name="model" location="https://zenodo.org/records/15611468/files/best.pt" />
-        <param name="image_size" value="512" />
-        <param name="confidence" value="0.003" />
-        <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">
-                <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">
-                <assert_contents>
-                    <has_image_width width="512"/>
-                    <has_image_height height="1024"/>
-                    <has_image_channels channels="1"/>
-                </assert_contents>
-            </element>
-        </output_collection>
-    </test>
-    <test>
+    <test expect_num_outputs="3">
             <param name="input_images" value="pred-test01.jpg" />
             <param name="model" location="https://zenodo.org/records/15611468/files/best.pt" />
             <param name="image_size" value="512" />
@@ -102,7 +74,7 @@
             <param name="class_name" value="class_name.txt" />
 
             <output_collection name="txt_results" type="list" count="1">
-                <element name="pred-test01">
+                <element name="pred-test01.jpg">
                     <assert_contents>
                         <has_n_lines n="100"/>
                     </assert_contents>
@@ -110,18 +82,29 @@
             </output_collection>
 
             <output_collection name="mask_results" type="list" count="1">
-                <element name="pred-test01">
+                <element name="pred-test01.jpg">
                     <assert_contents>
                         <has_image_width width="512"/>
                         <has_image_height height="1024"/>
                         <has_image_channels channels="1"/>
                     </assert_contents>
                 </element>
+	    </output_collection>
+            <output_collection name="overlay_results" type="list" count="1">
+	        <element name="pred-test01.jpg">
+		    <assert_contents>
+                        <has_image_width width="512"/>
+                        <has_image_height height="1024"/>
+                        <has_image_channels channels="3"/>
+                    </assert_contents>
+		</element>
             </output_collection>
+
         </test>    
         <!--new version's test-->
-        <!-- SEGMENT MODE TEST -->
-        <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" />
@@ -131,18 +114,19 @@
             <param name="iou" value="0.45" />
             <param name="max_det" value="300" />
             <output_collection name="txt_results" type="list" count="1">
-                <element name="bus" file="bus.txt" compare="sim_size" />
+                <element name="bus.jpg" file="bus.txt" compare="sim_size" />
             </output_collection>
             <output_collection name="mask_results" type="list" count="1">
-                <element name="bus" file="bus_mask.tiff" compare="sim_size" />
+                <element name="bus.jpg" file="bus_mask.tiff" compare="sim_size" />
             </output_collection>
             <output_collection name="overlay_results" type="list" count="1">
-                <element name="bus" file="bus_seg.jpg" compare="sim_size" />
+                <element name="bus.jpg" file="bus_seg.jpg" compare="sim_size" />
             </output_collection>
         </test>
 
-        <!-- DETECT MODE TEST -->
-        <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" />
@@ -152,12 +136,13 @@
             <param name="iou" value="0.45" />
             <param name="max_det" value="300" />
             <output_collection name="txt_results" type="list" count="1">
-                <element name="bus" file="bus_detect.txt" compare="sim_size" />
+                <element name="bus.jpg" file="bus_detect.txt" compare="sim_size" />
             </output_collection>
             <output_collection name="overlay_results" type="list" count="1">
-                <element name="bus" file="bus_seg.jpg" compare="sim_size" />
+                <element name="bus.jpg" file="bus_seg.jpg" compare="sim_size" />
             </output_collection>
-        </test>
+	</test>
+
     </tests>
 
     <help><![CDATA[
@@ -171,9 +156,9 @@
 
     Choose from two available options using the **mode** parameter:
 
-    - **detect** – Use this if your pre-trained model is a detection model. It outputs bounding boxes, class IDs, and confidence scores for objects detected in each image.
+	    - **detect** – Use this if your pre-trained/trained model is a detection model. It outputs bounding boxes coordinates, class IDs, and confidence scores and overlay images with bounding boxes.
 
-    - **segment** – Use this if your pre-trained model is a segmentation model. It outputs both bounding boxes **and** pixel-level segmentation masks. Additional output files include grayscale TIFF masks.
+	    - **segment** – Use this if your pre-trained/trained model is a segmentation model. It outputs text file containing polygons coordinates, class IDs and confidence score and overlaid images with bounding boxes and mask images.
 
     ⚠️ **Important:** The selected mode must match the type of the model you upload. For example, if you are using a segmentation model (`*-seg.pt`), you must select `segment` mode. Using a mismatched mode and model will lead to errors or empty results.
 
@@ -196,13 +181,13 @@
 
     - **YOLO Annotations (.txt)**:
     - One file per input image.
-    - Contains predicted class ID, confidence score, and bounding box (and segmentation data if applicable).
+    - Contains bounding boxes/polygons coordinates,predicted class ID, confidence scores.
 
     - **Segmentation Masks (.tiff)** (for `segment` mode only): 
-    - Grayscale images showing detected object areas.
+    - Binary images showing detected object areas.
 
     - **Images with annotations(.jpg)**:
-    - Input images with bounding boxes and only for segment mode the overlayer.
+    - Overlaid images with bounding boxes and/or segmentation overlay.
 
     ]]></help>
     <expand macro="citations" />