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()