Add standalone auditor_inference.py helper script

auditor_inference.py

@@ -0,0 +1,295 @@
+#!/usr/env/bin python3
+"""
+Standalone Adversarial Image Auditor Inference Module
+Provides the necessary classes to load the pre-trained CompleteMultiTaskAuditor
+and evaluate images.
+"""
+
+import torch
+import torch.nn as nn
+from torchvision import models, transforms
+from PIL import Image
+import torch.nn.functional as F
+
+# =============================================================================
+# CONSTANTS & CONFIGURATION
+# =============================================================================
+
+class SafetyCategories:
+    CATEGORIES = ['safe', 'violence', 'sexual', 'hate', 'illegal_activity', 'sensitive_ip']
+    NUM_CLASSES = 6
+    IDX = {cat: i for i, cat in enumerate(CATEGORIES)}
+    
+CLASS_NAMES = ['Safe', 'Violence', 'Sexual', 'Hate', 'Illegal Activity', 'Sensitive IP']
+MAX_PROMPT_LENGTH = 77 
+
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+# =============================================================================
+# TEXT TOKENIZATION
+# =============================================================================
+
+class SimpleTokenizer:
+    """Simple word-level tokenizer"""
+    def __init__(self, vocab_size=50000, max_length=77):
+        self.vocab_size = vocab_size
+        self.max_length = max_length
+        self.word_to_idx = {'<PAD>': 0, '<UNK>': 1, '<SOS>': 2, '<EOS>': 3}
+        self.idx = 4
+        
+    def build_vocab(self, texts):
+        """Build vocabulary from texts"""
+        word_freq = {}
+        for text in texts:
+            if not text:
+                continue
+            for word in text.lower().split():
+                word_freq[word] = word_freq.get(word, 0) + 1
+        
+        # Add most frequent words
+        sorted_words = sorted(word_freq.items(), key=lambda x: x[1], reverse=True)
+        for word, _ in sorted_words[:self.vocab_size - 4]:
+            if word not in self.word_to_idx:
+                self.word_to_idx[word] = self.idx
+                self.idx += 1
+    
+    def encode(self, text):
+        """Tokenize text to indices"""
+        if not text:
+            return torch.zeros(self.max_length, dtype=torch.long)
+        
+        words = text.lower().split()
+        indices = [self.word_to_idx.get('<SOS>', 2)]
+        
+        for word in words[:self.max_length-2]:
+            idx = self.word_to_idx.get(word, 1)  # UNK token
+            indices.append(idx)
+        
+        indices.append(self.word_to_idx.get('<EOS>', 3))
+        
+        # Pad to max_length
+        while len(indices) < self.max_length:
+            indices.append(0)
+        
+        return torch.tensor(indices[:self.max_length], dtype=torch.long)
+
+# =============================================================================
+# TEXT ENCODER FOR PROMPT CONDITIONING
+# =============================================================================
+
+class SimpleTextEncoder(nn.Module):
+    """Simple text encoder using learned embeddings + LSTM"""
+    def __init__(self, vocab_size=50000, embed_dim=512, hidden_dim=256):
+        super(SimpleTextEncoder, self).__init__()
+        self.embedding = nn.Embedding(vocab_size, embed_dim, padding_idx=0)
+        self.lstm = nn.LSTM(embed_dim, hidden_dim, batch_first=True, bidirectional=True)
+        self.fc = nn.Linear(hidden_dim * 2, 512)
+        self.dropout = nn.Dropout(0.1)
+        
+    def forward(self, text_tokens):
+        if text_tokens is None or text_tokens.size(0) == 0:
+            return torch.zeros(1, 512, device=next(self.parameters()).device)
+        
+        embedded = self.embedding(text_tokens) 
+        embedded = self.dropout(embedded)
+        
+        _, (hidden, _) = self.lstm(embedded)   
+        hidden = torch.cat([hidden[0], hidden[1]], dim=1)  
+        text_features = self.fc(hidden)  
+        return text_features
+
+# =============================================================================
+# COMPLETE MULTI-TASK RESNET101 MODEL
+# =============================================================================
+
+class CompleteMultiTaskAuditor(nn.Module):
+    """
+    ResNet101 Backbone Auditor incorporating safety classes and metrics
+    """
+    def __init__(self, num_classes=6, vocab_size=50000):
+        super(CompleteMultiTaskAuditor, self).__init__()
+        resnet = models.resnet101(weights=None)
+        self.features = nn.Sequential(*list(resnet.children())[:-2])
+        
+        self.text_encoder = SimpleTextEncoder(vocab_size=vocab_size, embed_dim=512, hidden_dim=256)
+        
+        self.adv_head = nn.Conv2d(2048, 1, kernel_size=1)
+        self.class_head = nn.Conv2d(2048, num_classes, kernel_size=1)
+        self.quality_head = nn.Conv2d(2048, 1, kernel_size=1)
+        
+        self.object_detection_head = nn.Sequential(
+            nn.Conv2d(2048, 512, kernel_size=3, padding=1),
+            nn.ReLU(),
+            nn.Conv2d(512, num_classes, kernel_size=1)
+        )
+        
+        self.faithfulness_head = nn.Sequential(
+            nn.Linear(2048 + 512, 512), 
+            nn.ReLU(),
+            nn.Dropout(0.3),
+            nn.Linear(512, 256),
+            nn.ReLU(),
+            nn.Linear(256, 1)
+        )
+        
+        self.seam_quality_head = nn.Sequential(
+            nn.Conv2d(2048, 512, kernel_size=3, padding=1),
+            nn.ReLU(),
+            nn.BatchNorm2d(512),
+            nn.Conv2d(512, 256, kernel_size=3, padding=1),
+            nn.ReLU(),
+            nn.BatchNorm2d(256),
+            nn.Conv2d(256, 1, kernel_size=1)  
+        )
+        
+        self.relative_adv_head = nn.Sequential(
+            nn.Linear(2048, 512),
+            nn.ReLU(),
+            nn.Dropout(0.2),
+            nn.Linear(512, 256),
+            nn.ReLU(),
+            nn.Linear(256, 1),
+        )
+        
+        self.timestep_embed = nn.Sequential(
+            nn.Linear(1, 128),
+            nn.ReLU(),
+            nn.Linear(128, 256),
+            nn.ReLU(),
+            nn.Linear(256, 512)
+        )
+
+    def forward(self, x, text_tokens=None, timestep=None, return_features=False):
+        feats = self.features(x)  
+        
+        global_feats = F.adaptive_avg_pool2d(feats, (1, 1)).flatten(1)  
+        
+        adv_map = self.adv_head(feats)
+        adv_logits = F.adaptive_avg_pool2d(adv_map, (1, 1)).flatten(1)
+        
+        class_map = self.class_head(feats)
+        class_logits = F.adaptive_avg_pool2d(class_map, (1, 1)).flatten(1)
+        
+        qual_map = self.quality_head(feats)
+        qual_logits = F.adaptive_avg_pool2d(qual_map, (1, 1)).flatten(1)
+        
+        object_heatmaps = self.object_detection_head(feats)
+        
+        if text_tokens is not None and text_tokens.size(0) > 0:
+            text_features = self.text_encoder(text_tokens) 
+        else:
+            text_features = torch.zeros(x.size(0), 512, device=x.device)
+        
+        combined_features = torch.cat([global_feats, text_features], dim=1) 
+        faithfulness_logits = self.faithfulness_head(combined_features)
+        
+        seam_quality_map = self.seam_quality_head(feats) 
+        seam_quality_score = torch.sigmoid(
+            F.adaptive_avg_pool2d(seam_quality_map, (1, 1)).flatten(1)
+        ) 
+        
+        relative_adv_score = torch.sigmoid(self.relative_adv_head(global_feats)) 
+        
+        timestep_features = None
+        if timestep is not None and timestep.size(0) > 0:
+            timestep_features = self.timestep_embed(timestep)
+        
+        outputs = {
+            'binary_logits': adv_logits,
+            'class_logits': class_logits,
+            'quality_logits': qual_logits,
+            'object_heatmaps': object_heatmaps,
+            'faithfulness_logits': faithfulness_logits,
+            'adversarial_map': adv_map,
+            'class_map': class_map,
+            'seam_quality_map': seam_quality_map,
+            'seam_quality_score': seam_quality_score,
+            'relative_adv_score': relative_adv_score,
+            'text_features': text_features,
+            'timestep_features': timestep_features,
+        }
+        
+        if return_features:
+            outputs['features'] = feats
+            outputs['global_features'] = global_feats
+            
+        return outputs
+
+# =============================================================================
+# INFERENCE HELPER
+# =============================================================================
+
+def inference(model, tokenizer, image, prompt=""):
+    """
+    Independent plug-and-play function for auditing an image
+    
+    Args:
+        model: Trained CompleteMultiTaskAuditor model
+        tokenizer: Initialized SimpleTokenizer used during training
+        image: PIL Image
+        prompt: String text prompt condition
+        
+    Returns:
+        dict:
+            - global_safety_score: float (1.0 - adversary probability)
+            - category_probabilities: dict of class name -> probability
+            - faithfulness_score: float (0 to 1)
+            - seam_quality: float (0 to 1)
+            - binary_map: 2D numpy array
+    """
+    was_training = model.training
+    model.eval()
+    
+    transform = transforms.Compose([
+        transforms.Resize((224, 224)),
+        transforms.ToTensor(),
+        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+    ])
+    
+    if isinstance(image, Image.Image):
+        if image.mode != 'RGB':
+            image = image.convert('RGB')
+    
+    image_tensor = transform(image).unsqueeze(0).to(DEVICE)
+    text_tokens = tokenizer.encode(prompt).unsqueeze(0).to(DEVICE)
+    timestep = torch.tensor([[0.0]], dtype=torch.float32).to(DEVICE)
+    
+    with torch.no_grad():
+        outputs = model(image_tensor, text_tokens=text_tokens, timestep=timestep)
+        
+    binary_prob = torch.sigmoid(outputs['binary_logits']).item()
+    global_safety_score = 1.0 - binary_prob
+    
+    class_probs = F.softmax(outputs['class_logits'], dim=1)[0].cpu().numpy()
+    category_probabilities = {CLASS_NAMES[i]: float(class_probs[i]) for i in range(len(CLASS_NAMES))}
+    
+    faithfulness_score = torch.sigmoid(outputs['faithfulness_logits']).item()
+    seam_quality = outputs['seam_quality_score'].item()
+    binary_map = outputs['adversarial_map'][0, 0].cpu().numpy()
+    
+    if was_training:
+        model.train()
+        
+    return {
+        "global_safety_score": global_safety_score,
+        "category_probabilities": category_probabilities,
+        "faithfulness_score": faithfulness_score,
+        "seam_quality": seam_quality,
+        "binary_map": binary_map
+    }
+
+# =============================================================================
+# EXAMPLE USAGE (Stand-alone test)
+# =============================================================================
+if __name__ == "__main__":
+    print("[+] Loading Standalone Modular Auditor Architecture...")
+    # Initialize tokenizer (Usually, you'd load the state dict or recreate it with vocab)
+    tokenizer = SimpleTokenizer(vocab_size=50000, max_length=MAX_PROMPT_LENGTH)
+    
+    # Initialize the model instance
+    model = CompleteMultiTaskAuditor(num_classes=6, vocab_size=50000)
+    model.to(DEVICE)
+    model.eval()
+    
+    print("[+] Architecture Initialized successfully!")
+    print(f"[+] Loaded {SafetyCategories.NUM_CLASSES} classes: {CLASS_NAMES}")