Computer vision YOLO11 model

Test_image.py

@@ -0,0 +1,66 @@
+import os
+import cv2
+import pandas as pd
+from ultralytics import YOLO
+
+# Define input and output directories
+input_dir = r'D:\AIM\pecan\test_image\input_image'
+output_dir = r'D:\AIM\pecan\test_image\output_image'
+
+# Create the output directory if it doesn't exist
+os.makedirs(output_dir, exist_ok=True)
+
+# Load a YOLO model (you can specify your model path here)
+model = YOLO(r"D:\AIM\pecan\runs\detect\train2\weights\best.pt")  # Load a custom-trained model
+
+# Initialize a threshold and a list for classifications
+threshold = 0.5
+all_classifications = []
+
+# Iterate over all images in the input directory
+for image_filename in os.listdir(input_dir):
+    if image_filename.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.tif', '.tiff')):  # Check for image files
+        image_path = os.path.join(input_dir, image_filename)
+        image = cv2.imread(image_path)
+        if image is None:
+            continue  # Skip files that cannot be read
+
+        # Perform inference
+        results = model(image)[0]
+        classifications = []
+
+        # Create a semi-transparent overlay for the bounding boxes
+        for result in results.boxes.data.tolist():
+            x1, y1, x2, y2, score, class_id = result
+
+            if score > threshold:
+                # Create a semi-transparent overlay for the bounding box area
+                overlay = image.copy()
+                cv2.rectangle(overlay, (int(x1), int(y1)), (int(x2), int(y2)), (0, 255, 0), -1)
+                alpha = 0.15  # Transparency factor (0 - transparent, 1 - opaque)
+                # Apply the overlay
+                image = cv2.addWeighted(overlay, alpha, image, 1 - alpha, 0)
+
+                # Draw bounding box outline
+                cv2.rectangle(image, (int(x1), int(y1)), (int(x2), int(y2)), (0, 255, 0), 2)
+                # Prepare label with class name and confidence score
+                label = f"{results.names[int(class_id)].upper()} {score:.2f}"
+                # Put text (class label with confidence score)
+                cv2.putText(image, label, (int(x1), int(y1 - 10)),
+                            cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2, cv2.LINE_AA)
+
+                # Add classification to the list
+                classifications.append(f"{results.names[int(class_id)]} ({score:.2f})")
+
+        # Save the output image
+        output_path = os.path.join(output_dir, image_filename)
+        cv2.imwrite(output_path, image)
+
+        # Save classifications for each image
+        all_classifications.append({"image": image_filename, "classes": classifications})
+
+# Create a DataFrame from the classifications
+df = pd.DataFrame(all_classifications)
+df.to_excel(os.path.join(output_dir, 'output_classifications.xlsx'), index=False)
+
+cv2.destroyAllWindows()