--- a/yolo_prediction.xml	Fri Jun 13 11:23:57 2025 +0000
+++ b/yolo_prediction.xml	Mon Jul 07 06:47:44 2025 +0000
@@ -8,7 +8,7 @@
     <expand macro="requirements" />
     <command detect_errors="aggressive">
     <![CDATA[
-      export YOLO_CONFIG_DIR='\$HOME/.config/ultralytics' &&
+      export YOLO_CONFIG_DIR="\$HOME/.config/ultralytics" &&
 
       mkdir -p ./input_images ./runs ./models ./results && 
       
@@ -19,22 +19,22 @@
       cp '$class_name' './models/class_name.txt' &&
 
       cp '$model' './models/model.pt' &&
-		    
+	    
       python '$__tool_directory__/yolov8.py' 
-		    --test_path='./input_images' 
-		    --model_path='./models' 
-		    --model_name='model' 
-		    --run_dir='./runs' 
-		    --save_dir='./results'
-		    --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'  
-
+	    --test_path='./input_images' 
+	    --model_path='./models' 
+	    --model_name='model' 
+	    --run_dir='./runs' 
+	    --save_dir='./results'
+	    --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>
@@ -42,8 +42,8 @@
 	<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>
+          <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." />
@@ -51,25 +51,57 @@
 	<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 prediction results (text)">
+        <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"/>
-	</collection>
-	<collection name="mask_results" format="tiff" type="list" label="YOLO prediction masks">
-            <discover_datasets pattern="(?P&lt;name&gt;.*)\.tiff$" directory="results" visible="true"/>
+        </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"/>
         </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"/>
+        </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>
             <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="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">
+            <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"/>
@@ -78,56 +110,100 @@
             </output_collection>
 
             <output_collection name="mask_results" type="list" count="1">
-                <element name="pred-test01_mask">
-                    <assert_contents>
-			<has_image_width width="512"/>
-			<has_image_height height="1024"/>
-                        <has_image_channels channels="1"/>
-                   </assert_contents>
-                </element>
-            </output_collection>
-        </test>
-        <test>
-            <param name="input_images" value="pred-test01.png" />
-            <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_mask">
-                    <assert_contents>
                         <has_image_width width="512"/>
                         <has_image_height height="1024"/>
                         <has_image_channels channels="1"/>
-                   </assert_contents>
+                    </assert_contents>
                 </element>
             </output_collection>
+        </test>    
+        <!--new version's test-->
+        <!-- SEGMENT MODE TEST -->
+        <test>
+            <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" />
+            <param name="confidence" value="0.25" />
+            <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" />
+            </output_collection>
+            <output_collection name="mask_results" type="list" count="1">
+                <element name="bus" 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" />
+            </output_collection>
+        </test>
+
+        <!-- DETECT MODE TEST -->
+        <test>
+            <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" />
+            <param name="confidence" value="0.25" />
+            <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" />
+            </output_collection>
+            <output_collection name="overlay_results" type="list" count="1">
+                <element name="bus" file="bus_seg.jpg" compare="sim_size" />
+            </output_collection>
         </test>
     </tests>
+
     <help><![CDATA[
-        **What it does**
-        This tool performs predictions on input images using yolo compatible model. It supports two modes of operation: segment, detect.
-        
-	**Outputs:**
+    **Overview**
+
+    This Galaxy tool performs object detection or instance segmentation on input images using a **pre-trained YOLOv8 model** from the **Ultralytics** framework.
+
+    ---
+
+    **Prediction Mode**
+
+    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.
+
+    - **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.
+
+    ⚠️ **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.
+
+    ---
+
+    **Inputs**
 
-	 - **Coordinates TXT File** (`.txt`)
-            - A plain text file containing predictions for each detected object.
-         
-	 OR
-         - **Segmentation Mask** (`.tiff`)
-         - An image mask where each pixel value represents a class label.
+    - **Input Images** (`jpg`, `png`): One or more images to analyze.
+    - **YOLO Class Name File** (`.txt`): A plain text file listing all class names used during training (one per line).
+    - **Model File** (`.pt`): A YOLOv8-compatible pre-trained model file.
+    - **Prediction Mode** (`detect` or `segment`): Choose based on the model type as described above.
+    - **Image Size** (default: `1000`): Input images will be resized to square dimensions. Larger sizes yield better accuracy but require more memory.
+    - **Confidence Threshold** (default: `0.5`): Minimum confidence to include a detection.
+    - **IoU Threshold** (default: `0.7`): Intersection-over-Union threshold for non-maximum suppression.
+    - **Max Detections** (default: `300`): Maximum number of objects detected per image.
+
+    ---
 
-    ]]>
-    </help>
+    **Outputs**
+
+    - **YOLO Annotations (.txt)**:
+    - One file per input image.
+    - Contains predicted class ID, confidence score, and bounding box (and segmentation data if applicable).
+
+    - **Segmentation Masks (.tiff)** (for `segment` mode only): 
+    - Grayscale images showing detected object areas.
+
+    - **Images with annotations(.jpg)**:
+    - Input images with bounding boxes and only for segment mode the overlayer.
+
+    ]]></help>
     <expand macro="citations" />
 </tool>