import os
from torchvision import transforms
from PIL import Image
import cv2
import torch
import numpy as np

# Modify the plotting section to:
import matplotlib
matplotlib.use('Agg')  # Set non-interactive backend
import matplotlib.pyplot as plt

# First define or import RetinaNet class
# RetinaNet implementation:
class RetinaNet(torch.nn.Module):
    def __init__(self):
        super(RetinaNet, self).__init__()
        # Add your model architecture here

    def forward(self, x):
        # Add forward pass
        return x
# Add RetinaNet to safe globals before loading
torch.serialization.add_safe_globals([RetinaNet])
PATH = './20_feb_best_model_deployment.pth'
# Load model with proper safety measures
try:
    model = torch.load(PATH, 
                      map_location=torch.device("cpu"), 
                      weights_only=False)
except Exception as e:
    print(f"Error loading model: {e}")
    exit()

# Preprocessing function
def preprocess_frame(frame):
    rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
    resized_frame = cv2.resize(rgb_frame, (224, 224))
    normalized_frame = resized_frame / 255.0
    return np.expand_dims(normalized_frame, axis=0)

# Loop through images in the folder
for img_name in os.listdir('./images'):
    print(f"Processing image: {img_name}")

    # Load and preprocess the image
    # Define device
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    
    # Load image from correct path
    img_path = os.path.join('./images', img_name)
    img = Image.open(img_path)
    
    # Define transformation: resize and convert to tensor
    transform = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor()
    ])

    # Apply transformation to loaded image
    new_image = transform(img)
    # Add batch dimension
    new_image_batch = new_image.unsqueeze(0)

    print(new_image_batch.shape)  # Should output: torch.Size([1, 3, 224, 224])

 

    # Predict using the best model
    model.eval()
    with torch.no_grad():
        # Convert tensor to correct format for model
        new_image_batch = new_image_batch.permute(0, 3, 1, 2)
        prediction = model(new_image_batch.to(device).float())

        # print("Prediction:", abs(prediction))

    # Interpret the prediction
    # if prediction[0][0] > 0.5:
    # sum_value=torch.sum(abs(prediction[0]))
    # # print("Sum Value:", sum_value)
    # p_true=abs(prediction[0][0])
    # print("p_true:",p_true)
    # p_false=abs(prediction[0][1])
    # print("p_false:",p_false)
    # Interpret the prediction
    if prediction.shape[1] == 2:  # If binary classification output
        # Convert to probabilities using softmax
        probabilities = torch.nn.functional.softmax(prediction, dim=1)
        p_accept = probabilities[0][0].item() * 100  # Convert to percentage
        p_reject = probabilities[0][1].item() * 100
    else:  # If using regression/other output
        # Take mean of the output channels
        p_accept = abs(prediction[0][0]).mean().item() * 100
        p_reject = abs(prediction[0][1]).mean().item() * 100

    print(f"Accept probability: {p_accept:.2f}%")
    print(f"Reject probability: {p_reject:.2f}%")

    # Threshold for classification (now using percentage values)
    if p_accept > 35:  # 40% threshold
        result = "Acceptable"
    else:
        result = "Rejectable"



    # Threshold for classification (may need adjustment based on model training)
    if p_accept > 0.42:  # Consider making this a configurable parameter
        result = "Acceptable"
    else:
        result = "Rejectable"
    print(f"Predicted: {result}")

# Then in your loop:
fig = plt.figure()
plt.imshow(new_image.permute(1, 2, 0))
plt.title(f"Prediction: {result} ({p_accept:.1f}% vs {p_reject:.1f}%)")
plt.axis('off')
plt.savefig(f'./output_images/output_{img_name}')  # Save instead of show
plt.close(fig)  # Important for memory management