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