66 lines
2.7 KiB
Python
66 lines
2.7 KiB
Python
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() |