Spaces:

namchain
/

STAR

Sleeping

STAR / src /inference.py

Added Source Folder

af59080 verified over 1 year ago

4.52 kB

"""
Explainable AI (XAI) Inference for Nude Multi-Label Classification
==================================================================
This script performs inference using a trained Swin Transformer model for 
multi-label classification of nude images. It also integrates Class Activation 
Mapping (CAM) to provide visual explanations for the model's predictions.
Author: Ramaguru Radhakrishnan
Date: March 2025
"""

import torch
import torchvision.transforms as transforms
from PIL import Image
import json
from model import SwinTransformerMultiLabel
from torchcam.methods import SmoothGradCAMpp  # Explainability module
import matplotlib.pyplot as plt
import numpy as np


# Define the number of output classes (should match the trained model)
NUM_CLASSES = 18  

# Load the trained model with a correct classifier head
model = SwinTransformerMultiLabel(num_classes=NUM_CLASSES)

# Load model weights while ignoring mismatched layers
checkpoint_path = "../models/multi_nude_detector.pth"
checkpoint = torch.load(checkpoint_path, map_location="cpu")
model_dict = model.state_dict()

# Filter out layers that do not match
filtered_checkpoint = {
    k: v for k, v in checkpoint.items() if k in model_dict and v.shape == model_dict[k].shape
}
model_dict.update(filtered_checkpoint)
model.load_state_dict(model_dict, strict=False)

# Set the model to evaluation mode
model.eval()

# Define image preprocessing transformations
transform = transforms.Compose([
    transforms.Resize((224, 224)),  # Resize to model's input size
    transforms.ToTensor(),  # Convert to tensor
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),  # Normalize
])

# Load class labels from JSON file
with open("../data/labels.json", "r") as f:
    classes = sorted(set(tag for tags in json.load(f).values() for tag in tags))

# Validate that the number of classes matches
if len(classes) != NUM_CLASSES:
    raise ValueError(f"❌ Mismatch: Model expects {NUM_CLASSES} classes, but labels.json has {len(classes)} labels!")

# Load the test image
img_path = "C:\\Users\\RamaguruRadhakrishna\\Videos\\STAR-main\\STAR-main\\data\\images\\442_.jpeg"
image = Image.open(img_path).convert("RGB")  # Ensure RGB format
input_tensor = transform(image).unsqueeze(0)  # Add batch dimension

# Perform inference
with torch.no_grad():
    output = model(input_tensor)  # Forward pass through model
    print(f"🔹 Model Output Shape: {output.shape}")  # Debugging

    # Get predicted labels (threshold = 0.5)
    predicted_labels = [
        classes[i] for i in range(min(len(classes), output.shape[1])) if output[0][i] > 0.5
    ]
    predicted_indices = [i for i in range(output.shape[1]) if output[0][i] > 0.5]  # Store indices

# Display predicted labels
print("✅ Predicted Tags:", predicted_labels)

# ===============================
# Explainable AI: CAM Visualization
# ===============================

# Print model structure to find the correct target layer
print(model)

# Print model architecture to identify available layers
print("🔍 Model Architecture:\n")
for name, module in model.named_modules():
    print(name)  # Uncomment to see available layers

# Choose a valid convolutional layer from printed names
# Example: 'features.7.3' (Update this with an actual layer from print output)
valid_target_layer = "features.7.3"  # Modify based on your model structure

# Verify if the layer exists in the model
if valid_target_layer not in dict(model.named_modules()):
    raise ValueError(f"❌ Layer '{valid_target_layer}' not found in model. Choose from:\n{list(dict(model.named_modules()).keys())}")

# Initialize SmoothGradCAMpp with a valid layer
cam_extractor = SmoothGradCAMpp(model, target_layer=valid_target_layer)

print("✅ SmoothGradCAMpp initialized successfully!")

# Ensure model has processed the input before extracting CAM
output = model(input_tensor)  

# Generate CAM heatmaps for each predicted label
for class_idx in predicted_indices:
    cam = cam_extractor(class_idx, output)
    cam = cam.squeeze().cpu().numpy()
    cam = (cam - np.min(cam)) / (np.max(cam) - np.min(cam))  # Normalize

    # Resize CAM to match input image dimensions
    cam_resized = np.array(Image.fromarray(cam * 255).resize(image.size))

    # Overlay CAM on the original image
    plt.figure(figsize=(6, 6))
    plt.imshow(image)
    plt.imshow(cam_resized, cmap='jet', alpha=0.5)  # Heatmap overlay
    plt.axis("off")
    plt.title(f"Explainability Heatmap for '{classes[class_idx]}'")
    plt.show()

	"""
	Explainable AI (XAI) Inference for Nude Multi-Label Classification
	==================================================================

	This script performs inference using a trained Swin Transformer model for
	multi-label classification of nude images. It also integrates Class Activation
	Mapping (CAM) to provide visual explanations for the model's predictions.

	Author: Ramaguru Radhakrishnan
	Date: March 2025
	"""

	import torch
	import torchvision.transforms as transforms
	from PIL import Image
	import json
	from model import SwinTransformerMultiLabel
	from torchcam.methods import SmoothGradCAMpp # Explainability module
	import matplotlib.pyplot as plt
	import numpy as np


	# Define the number of output classes (should match the trained model)
	NUM_CLASSES = 18

	# Load the trained model with a correct classifier head
	model = SwinTransformerMultiLabel(num_classes=NUM_CLASSES)

	# Load model weights while ignoring mismatched layers
	checkpoint_path = "../models/multi_nude_detector.pth"
	checkpoint = torch.load(checkpoint_path, map_location="cpu")
	model_dict = model.state_dict()

	# Filter out layers that do not match
	filtered_checkpoint = {
	k: v for k, v in checkpoint.items() if k in model_dict and v.shape == model_dict[k].shape
	}
	model_dict.update(filtered_checkpoint)
	model.load_state_dict(model_dict, strict=False)

	# Set the model to evaluation mode
	model.eval()

	# Define image preprocessing transformations
	transform = transforms.Compose([
	transforms.Resize((224, 224)), # Resize to model's input size
	transforms.ToTensor(), # Convert to tensor
	transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), # Normalize
	])

	# Load class labels from JSON file
	with open("../data/labels.json", "r") as f:
	classes = sorted(set(tag for tags in json.load(f).values() for tag in tags))

	# Validate that the number of classes matches
	if len(classes) != NUM_CLASSES:
	raise ValueError(f"❌ Mismatch: Model expects {NUM_CLASSES} classes, but labels.json has {len(classes)} labels!")

	# Load the test image
	img_path = "C:\\Users\\RamaguruRadhakrishna\\Videos\\STAR-main\\STAR-main\\data\\images\\442_.jpeg"
	image = Image.open(img_path).convert("RGB") # Ensure RGB format
	input_tensor = transform(image).unsqueeze(0) # Add batch dimension

	# Perform inference
	with torch.no_grad():
	output = model(input_tensor) # Forward pass through model
	print(f"🔹 Model Output Shape: {output.shape}") # Debugging

	# Get predicted labels (threshold = 0.5)
	predicted_labels = [
	classes[i] for i in range(min(len(classes), output.shape[1])) if output[0][i] > 0.5
	]
	predicted_indices = [i for i in range(output.shape[1]) if output[0][i] > 0.5] # Store indices

	# Display predicted labels
	print("✅ Predicted Tags:", predicted_labels)

	# ===============================
	# Explainable AI: CAM Visualization
	# ===============================

	# Print model structure to find the correct target layer
	print(model)

	# Print model architecture to identify available layers
	print("🔍 Model Architecture:\n")
	for name, module in model.named_modules():
	print(name) # Uncomment to see available layers

	# Choose a valid convolutional layer from printed names
	# Example: 'features.7.3' (Update this with an actual layer from print output)
	valid_target_layer = "features.7.3" # Modify based on your model structure

	# Verify if the layer exists in the model
	if valid_target_layer not in dict(model.named_modules()):
	raise ValueError(f"❌ Layer '{valid_target_layer}' not found in model. Choose from:\n{list(dict(model.named_modules()).keys())}")

	# Initialize SmoothGradCAMpp with a valid layer
	cam_extractor = SmoothGradCAMpp(model, target_layer=valid_target_layer)

	print("✅ SmoothGradCAMpp initialized successfully!")

	# Ensure model has processed the input before extracting CAM
	output = model(input_tensor)

	# Generate CAM heatmaps for each predicted label
	for class_idx in predicted_indices:
	cam = cam_extractor(class_idx, output)
	cam = cam.squeeze().cpu().numpy()
	cam = (cam - np.min(cam)) / (np.max(cam) - np.min(cam)) # Normalize

	# Resize CAM to match input image dimensions
	cam_resized = np.array(Image.fromarray(cam * 255).resize(image.size))

	# Overlay CAM on the original image
	plt.figure(figsize=(6, 6))
	plt.imshow(image)
	plt.imshow(cam_resized, cmap='jet', alpha=0.5) # Heatmap overlay
	plt.axis("off")
	plt.title(f"Explainability Heatmap for '{classes[class_idx]}'")
	plt.show()