import os

import cv2

from ultralytics import YOLO

import numpy as np

import pandas as pd

from sklearn.metrics import classification_report, accuracy_score, f1_score

from tqdm import tqdm

# Paths to your images and labels folders

images_folder = r'D:\CSE\Semester - 5\Extra\IIT jammu\YOLO\Code\Data\valid\images'

labels_folder = r'D:\CSE\Semester - 5\Extra\IIT jammu\YOLO\Code\Data\valid\labels'

# Initialize lists to store results

true_labels = []

predicted_labels = []

# Initialize lists to store IoU scores

iou_scores = []

# Initialize the YOLOv8 model

model = YOLO("best_75.pt")

# Initialize lists to store detection results

detection_results = []

def calculate_iou(box1_tensor, box2_list, image_width, image_height):

    # Extract values from the tensor

    x1, y1, x2, y2, confidence, class_id = box1_tensor.tolist()

    # Convert values from the list to strings

    class_id2, x2_gt, y2_gt, w2_gt, h2_gt = map(str, box2_list)

    # Convert coordinates to float

    x2_gt, y2_gt, w2_gt, h2_gt = float(x2_gt), float(y2_gt), float(w2_gt), float(h2_gt)

    # Convert the model's bounding box coordinates to YOLO format

    x1, y1, x2, y2 = x1 / image_width, y1 / image_height, x2 / image_width, y2 / image_height

    # Calculate the intersection coordinates

    xA = max(x1, x2_gt)

    yA = max(y1, y2_gt)

    xB = min(x1 + x2, x2_gt + w2_gt)

    yB = min(y1 + y2, y2_gt + h2_gt)

    # Calculate the intersection area

    interArea = max(0, xB - xA) * max(0, yB - yA)

    # Calculate the areas of the boxes

    box1Area = x2 * y2

    box2Area = w2_gt * h2_gt

    # Calculate the IoU

    iou = interArea / float(box1Area + box2Area - interArea)

    return iou

# Iterate through image files and corresponding label files

for image_file in tqdm(os.listdir(images_folder)):

    if image_file.endswith('.png'):

        image_path = os.path.join(images_folder, image_file)

        # Determine class (bleeding or non-bleeding) from image file name

        is_bleeding = image_file.startswith('img-')

        # Load the image

        image = cv2.imread(image_path)

        # Perform inference

        results = model.predict(image)

        # Check if bleeding is detected

        detected_bleeding = results[0].boxes.shape[0] > 0

        # Append true and predicted labels

        true_labels.append(is_bleeding)

        predicted_labels.append(detected_bleeding)

        # Extract ground truth bounding boxes from label file

        label_file_name = os.path.splitext(image_file)[0] + '.txt'

        label_file_path = os.path.join(labels_folder, label_file_name)

        if os.path.exists(label_file_path):

            with open(label_file_path, 'r') as label_file:

                lines = label_file.readlines()

                ground_truth_boxes = [list(map(float, line.strip().split())) for line in lines]

            # Calculate IoU for each detected box

            if results[0].boxes.shape[0] > 0:

                for box in results[0].boxes.boxes:

                    iou_values = []

                    for gt_box in ground_truth_boxes:

                        iou = calculate_iou(box, gt_box, 224, 224)

                        iou_values.append(iou)

                    max_iou = max(iou_values) if iou_values else 0

                    iou_scores.append(max_iou)

            # Append detection results

            if results[0].boxes.shape[0] > 0:

                box = results[0].boxes[0]

                detection_results.append({

                    'image_path': image_path,

                    'is_bleeding': is_bleeding,

                    'detection_confidence': box.conf[0].item()

                })

# Calculate classification metrics

accuracy = accuracy_score(true_labels, predicted_labels)

recall = classification_report(true_labels, predicted_labels, target_names=['Non-Bleeding', 'Bleeding'], output_dict=True)

f1 = f1_score(true_labels, predicted_labels)

# Create a table of achieved evaluation metrics

classification_metrics = {

    'Accuracy': [accuracy],

    'Recall (Non-Bleeding)': [recall['Non-Bleeding']['recall']],

    'Recall (Bleeding)': [recall['Bleeding']['recall']],

    'F1-Score': [f1]

}

classification_df = pd.DataFrame(classification_metrics)

print("Classification Metrics:")

print(classification_df)

print()

print("Detection Metrics:")

print("IoU Metrics:")

# Calculate and print IoU statistics

iou_scores = np.array(iou_scores)

iou_mean = np.mean(iou_scores)

iou_median = np.median(iou_scores)

iou_75th_percentile = np.percentile(iou_scores, 75)

print("IoU Mean:", iou_mean)

print("IoU Median:", iou_median)

print("IoU 75th Percentile:", iou_75th_percentile)

# Calculate Detection Average Precision (AP) and Mean Average Precision (mAP)

detection_results_df = pd.DataFrame(detection_results)

detection_results_df['detection_confidence'] = detection_results_df['detection_confidence'].apply(lambda x: x.item() if isinstance(x, np.float32) else x)

detection_results_df['is_detected'] = detection_results_df['detection_confidence'] > 0.5  # You can adjust the confidence threshold as needed

detection_results_df['is_detected'] = detection_results_df['is_detected'].astype(int)

detection_ap = detection_results_df.groupby('is_bleeding')['is_detected'].apply(lambda x: np.mean(x))

mAP = detection_ap.mean()

print("Detection Average Precision (AP):")

print(detection_ap)

print("Mean Average Precision (mAP):", mAP)