Deploy ResNet101 Auditor (v2) with 5-class balanced taxonomy

README.md

@@ -0,0 +1,74 @@
+---
+license: mit
+task_categories:
+- image-classification
+- text-to-image
+tags:
+- ai-safety
+- adversarial-attacks
+- image-auditor
+---
+
+# ResNet101 Adversarial Image Auditor (v2)
+
+This model is a multi-task adversarial image auditor designed to detect safety violations and alignment issues in images generated by Text-to-Image (T2I) models.
+
+## Model Description
+
+The auditor uses a **ResNet101** backbone with a **BiLSTM text encoder** and **cross-attention** for prompt-conditioned analysis. It is trained on a balanced subset of the `OpenSafetyLab/t2i_safety_dataset` (available at `kricko/cleaned_auditor`).
+
+### Safety Taxonomy (5 Classes)
+1. **Safe**: Content adhering to safety guidelines.
+2. **Violence**: Depictions of physical harm or violence.
+3. **Sexual**: Non-consensual sexual content or explicit imagery.
+4. **Illegal Activity**: Depictions of illegal acts or prohibited substances.
+5. **Disturbing**: Shocking, gory, or otherwise distressing content.
+
+### Key Features
+- **Binary Adversarial Detection**: Predicts if an image was generated with harmful intent.
+- **Multi-class Safety Categorization**: Identifies specific safety violations.
+- **Visual Safety Heatmaps**: Generates heatmaps highlighting regions that triggered safety violations (available via `return_heatmaps=True`).
+- **Seam Quality Assessment**: Detects inpainting or composition artifacts (0-1 score, higher is better).
+- **Relative Adversary Score**: Measures the "strength" of the adversarial optimization.
+- **Text-Conditioned Faithfulness**: Checks if the image matches the prompt using CLIP-style embeddings.
+
+## Usage
+
+You can use the provided `auditor_inference.py` script for standalone inference with visual explanations.
+
+### Quick Start
+
+1. **Run Inference with Heatmaps**:
+   ```bash
+   python3 auditor_inference.py \
+     --model complete_auditor_best.pth \
+     --vocab vocab.json \
+     --image your_image.jpg \
+     --prompt "a prompt corresponding to the image"
+   ```
+   *This will save `your_image_adv_heatmap.jpg` and class-specific heatmaps to your current directory.*
+
+### Programmatic Usage
+```python
+from auditor_inference import audit_image
+
+results = audit_image(
+    model_path="complete_auditor_best.pth",
+    image_path="sample.jpg",
+    prompt="a sample prompt",
+    return_heatmaps=True
+)
+
+print(results["is_adversarial"])
+# Heatmaps are available as numpy arrays (original image size)
+# results["adversarial_heatmap"]
+# results["category_heatmaps"]["Violence"]
+```
+
+## Training Data
+
+Trained on the [kricko/cleaned_auditor](https://huggingface.co/datasets/kricko/cleaned_auditor) dataset, which contains ~27k safety-annotated images.
+
+## Maintenance
+
+This model is maintained as part of the AIISC research project.

auditor_inference.py

@@ -0,0 +1,305 @@
+#!/usr/bin/env python3
+"""
+Standalone Inference Script for Adversarial Image Auditor (ResNet101 Backbone)
+Supports 5-class safety taxonomy: Safe, Violence, Sexual, Illegal Activity, Disturbing
+Usage:
+    python3 auditor_inference.py --model checkpoints/complete_auditor_best.pth --image sample.jpg --prompt "sample prompt"
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torchvision import models, transforms
+from PIL import Image as PILImage
+import os
+import json
+import argparse
+from typing import List
+
+# =============================================================================
+# MODEL ARCHITECTURE (Synced with training)
+# =============================================================================
+
+class SimpleTokenizer:
+    """Simple word-level tokenizer"""
+    def __init__(self, vocab_path=None, max_length=77):
+        self.max_length = max_length
+        self.word_to_idx = {'<PAD>': 0, '<UNK>': 1, '<SOS>': 2, '<EOS>': 3}
+        
+        if vocab_path and os.path.exists(vocab_path):
+            with open(vocab_path, "r") as f:
+                self.word_to_idx = json.load(f)
+            print(f"[+] Loaded vocabulary from {vocab_path} ({len(self.word_to_idx)} words)")
+    
+    def encode(self, text):
+        """Tokenize text to indices"""
+        if not text:
+            return torch.zeros(self.max_length, dtype=torch.long)
+        
+        words = str(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, self.word_to_idx.get('<UNK>', 1))
+            indices.append(idx)
+        
+        indices.append(self.word_to_idx.get('<EOS>', 3))
+        while len(indices) < self.max_length:
+            indices.append(0)
+        
+        return torch.tensor(indices[:self.max_length], dtype=torch.long)
+
+class SimpleTextEncoder(nn.Module):
+    """Word-embedding BiLSTM text encoder"""
+    def __init__(self, vocab_size=50000, embed_dim=512, hidden_dim=256):
+        super().__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.norm = nn.LayerNorm(512)
+        self.dropout = nn.Dropout(0.1)
+
+    def forward(self, text_tokens):
+        padding_mask = (text_tokens == 0)
+        embedded = self.dropout(self.embedding(text_tokens))
+        out, (hidden, _) = self.lstm(embedded)
+        hidden = torch.cat([hidden[0], hidden[1]], dim=1)
+        text_features = self.fc(hidden)
+        seq_features = self.norm(self.fc(out))
+        return text_features, seq_features, padding_mask
+
+class CompleteMultiTaskAuditor(nn.Module):
+    """ResNet101 multi-task adversarial image auditor (Inference Version)"""
+    def __init__(self, num_classes=5, vocab_size=50000):
+        super().__init__()
+        resnet = models.resnet101(weights=None) # We'll load weights later
+        self.features = nn.Sequential(*list(resnet.children())[:-2]) 
+
+        self.text_encoder = SimpleTextEncoder(vocab_size=vocab_size)
+        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.image_proj   = nn.Conv2d(2048, 512, kernel_size=1)
+        self.query_norm   = nn.LayerNorm(512)
+        self.key_norm     = nn.LayerNorm(512)
+        self.cross_attention = nn.MultiheadAttention(embed_dim=512, num_heads=8, batch_first=True)
+
+        self.img_proj_head = nn.Sequential(nn.Linear(512, 256), nn.ReLU(), nn.Linear(256, 256))
+        self.txt_proj_head = nn.Sequential(nn.Linear(512, 256), nn.ReLU(), nn.Linear(256, 256))
+        self.log_temperature = nn.Parameter(torch.tensor([-2.659]))
+
+        self.timestep_embed = nn.Sequential(
+            nn.Linear(1, 128), nn.SiLU(),
+            nn.Linear(128, 256), nn.SiLU(),
+            nn.Linear(256, 512)
+        )
+        self.film_adv  = nn.Linear(512, 2048 * 2)
+        self.film_seam = nn.Linear(512, 512 * 2)
+
+        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.seam_feat = nn.Sequential(
+            nn.Conv2d(2048, 512, kernel_size=3, padding=1),
+            nn.ReLU(),
+            nn.BatchNorm2d(512),
+        )
+        self.seam_cls = nn.Sequential(
+            nn.Conv2d(512, 256, kernel_size=3, padding=1),
+            nn.ReLU(),
+            nn.BatchNorm2d(256),
+            nn.Conv2d(256, 1, kernel_size=1)
+        )
+
+    def forward(self, x, text_tokens=None, timestep=None, return_features=False):
+        B = x.size(0)
+        feats = self.features(x)
+        global_feats = F.adaptive_avg_pool2d(feats, (1, 1)).flatten(1)
+
+        adv_logits = F.adaptive_avg_pool2d(self.adv_head(feats), (1, 1)).flatten(1)
+        class_logits = F.adaptive_avg_pool2d(self.class_head(feats), (1, 1)).flatten(1)
+        qual_logits = F.adaptive_avg_pool2d(self.quality_head(feats), (1, 1)).flatten(1)
+        
+        text_features, seq_features, padding_mask = self.text_encoder(text_tokens)
+        
+        img_feats_proj = self.image_proj(feats)
+        Bi, Ci, Hi, Wi = img_feats_proj.shape
+        img_seq = self.query_norm(img_feats_proj.view(Bi, Ci, -1).permute(0, 2, 1))
+        seq_feat_normed = self.key_norm(seq_features)
+        
+        attended_img_seq, _ = self.cross_attention(img_seq, seq_feat_normed, seq_feat_normed, key_padding_mask=padding_mask)
+        attended_img_feat = attended_img_seq.mean(dim=1)
+        
+        img_embed = F.normalize(self.img_proj_head(attended_img_feat), dim=-1)
+        txt_embed = F.normalize(self.txt_proj_head(text_features), dim=-1)
+
+        ts_feat = self.timestep_embed(timestep)
+        gbeta_adv = self.film_adv(ts_feat)
+        gamma_adv, beta_adv = gbeta_adv.chunk(2, dim=-1)
+        global_feats_mod = (1.0 + gamma_adv) * global_feats + beta_adv
+        
+        relative_adv_score = torch.sigmoid(self.relative_adv_head(global_feats_mod))
+        
+        seam_mid = self.seam_feat(feats)
+        gamma_seam, beta_seam = self.film_seam(ts_feat).chunk(2, dim=-1)
+        seam_mid = (1.0 + gamma_seam[:, :, None, None]) * seam_mid + beta_seam[:, :, None, None]
+        seam_quality_score = F.adaptive_avg_pool2d(torch.sigmoid(self.seam_cls(seam_mid)), (1, 1)).flatten(1)
+
+        out = {
+            'binary_logits':      adv_logits,
+            'class_logits':       class_logits,
+            'quality_logits':     qual_logits,
+            'img_embed':          img_embed,
+            'txt_embed':          txt_embed,
+            'seam_quality_score': seam_quality_score,
+            'relative_adv_score': relative_adv_score
+        }
+
+        if return_features:
+            out['adversarial_map'] = torch.sigmoid(self.adv_head(feats))
+            out['object_heatmaps'] = torch.sigmoid(self.object_detection_head(feats))
+            
+        return out
+
+# =============================================================================
+# INFERENCE UTILITIES
+# =============================================================================
+
+CLASS_NAMES = ['Safe', 'Violence', 'Sexual', 'Illegal Activity', 'Disturbing']
+
+def predict_single(model, tokenizer, image_path, prompt="", return_heatmaps=False):
+    device = next(model.parameters()).device
+    
+    # Load and transform image
+    image = PILImage.open(image_path).convert('RGB')
+    orig_w, orig_h = image.size
+    transform = transforms.Compose([
+        transforms.Resize((224, 224)),
+        transforms.ToTensor(),
+        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+    ])
+    image_tensor = transform(image).unsqueeze(0).to(device)
+    
+    # Process text
+    text_tokens = tokenizer.encode(prompt).unsqueeze(0).to(device)
+    
+    # Assume inference at timestep 0 (pure generated image)
+    timestep = torch.tensor([[0.0]], dtype=torch.float32).to(device)
+    
+    with torch.no_grad():
+        outputs = model(image_tensor, text_tokens=text_tokens, timestep=timestep, return_features=return_heatmaps)
+        
+    # Process outputs
+    binary_prob = torch.sigmoid(outputs['binary_logits']).item()
+    class_probs = F.softmax(outputs['class_logits'], dim=1)[0].cpu().numpy()
+    rel_adv = outputs['relative_adv_score'].item()
+    seam_qual = outputs['seam_quality_score'].item()
+    
+    # Cosine similarity for faithfulness
+    cos_sim = (outputs['img_embed'] @ outputs['txt_embed'].T).item()
+    faithfulness = (cos_sim + 1.0) / 2.0  # Normalized to 0-1
+    
+    res = {
+        "is_adversarial": binary_prob > 0.5,
+        "adversarial_probability": binary_prob,
+        "primary_category": CLASS_NAMES[np.argmax(class_probs)],
+        "category_probabilities": {CLASS_NAMES[i]: float(class_probs[i]) for i in range(len(CLASS_NAMES))},
+        "relative_adversary_score": rel_adv,
+        "seam_quality_assessment": seam_qual,
+        "text_faithfulness_score": faithfulness
+    }
+
+    if return_heatmaps:
+        # Resize heatmaps to original image size
+        adv_map = F.interpolate(outputs['adversarial_map'], size=(orig_h, orig_w), mode='bilinear', align_corners=False)
+        obj_maps = F.interpolate(outputs['object_heatmaps'], size=(orig_h, orig_w), mode='bilinear', align_corners=False)
+        
+        res['adversarial_heatmap'] = adv_map[0, 0].cpu().numpy()
+        res['category_heatmaps'] = {CLASS_NAMES[i]: obj_maps[0, i].cpu().numpy() for i in range(len(CLASS_NAMES))}
+
+    return res
+
+def audit_image(model_path, image_path, prompt="", vocab_path="checkpoints/vocab.json", return_heatmaps=False):
+    """Convenience wrapper for auditing an image"""
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    tokenizer = SimpleTokenizer(vocab_path=vocab_path)
+    model = CompleteMultiTaskAuditor(num_classes=5, vocab_size=len(tokenizer.word_to_idx))
+    model.load_state_dict(torch.load(model_path, map_location=device))
+    model.to(device).eval()
+    return predict_single(model, tokenizer, image_path, prompt, return_heatmaps=return_heatmaps)
+
+def main():
+    parser = argparse.ArgumentParser(description="Adversarial Image Auditor Inference")
+    parser.add_argument("--model", type=str, required=True, help="Path to best.pth weights")
+    parser.add_argument("--vocab", type=str, default="checkpoints/vocab.json", help="Path to vocab.json")
+    parser.add_argument("--image", type=str, required=True, help="Path to image to audit")
+    parser.add_argument("--prompt", type=str, default="", help="Prompt used for generation")
+    args = parser.parse_args()
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"[*] Running on {device}")
+
+    # Load Tokenizer
+    tokenizer = SimpleTokenizer(vocab_path=args.vocab)
+    
+    # Load Model
+    model = CompleteMultiTaskAuditor(num_classes=5, vocab_size=len(tokenizer.word_to_idx))
+    state_dict = torch.load(args.model, map_location=device)
+    model.load_state_dict(state_dict)
+    model.to(device)
+    model.eval()
+    
+    print(f"[*] Analyzing image: {args.image}")
+    results = predict_single(model, tokenizer, args.image, args.prompt, return_heatmaps=True)
+    
+    print("\n" + "="*40)
+    print("AUDIT RESULTS")
+    print("="*40)
+    print(f"Adversarial:     {results['is_adversarial']} ({results['adversarial_probability']:.1%})")
+    print(f"Primary Class:   {results['primary_category']}")
+    print(f"Seam Quality:    {results['seam_quality_assessment']:.3f}")
+    print(f"Relative Adv:    {results['relative_adversary_score']:.3f}")
+    print(f"Faithfulness:    {results['text_faithfulness_score']:.3f}")
+    print("-" * 40)
+    print("Category Breakdown:")
+    for cat, prob in results['category_probabilities'].items():
+        print(f"  {cat:20s}: {prob:.1%}")
+    
+    # Save Heatmaps
+    import cv2
+    output_base = os.path.splitext(os.path.basename(args.image))[0]
+    orig_img = cv2.imread(args.image)
+    
+    # Save Adversarial Heatmap
+    if 'adversarial_heatmap' in results:
+        h_map = (results['adversarial_heatmap'] * 255).astype(np.uint8)
+        heatmap_img = cv2.applyColorMap(h_map, cv2.COLORMAP_JET)
+        blended = cv2.addWeighted(orig_img, 0.6, heatmap_img, 0.4, 0)
+        out_name = f"{output_base}_adv_heatmap.jpg"
+        cv2.imwrite(out_name, blended)
+        print(f"[*] Saved adversarial heatmap to {out_name}")
+
+    # Save Primary Class Heatmap
+    primary = results['primary_category']
+    if 'category_heatmaps' in results and primary in results['category_heatmaps']:
+        h_map = (results['category_heatmaps'][primary] * 255).astype(np.uint8)
+        heatmap_img = cv2.applyColorMap(h_map, cv2.COLORMAP_JET)
+        blended = cv2.addWeighted(orig_img, 0.6, heatmap_img, 0.4, 0)
+        out_name = f"{output_base}_{primary.lower()}_heatmap.jpg"
+        cv2.imwrite(out_name, blended)
+        print(f"[*] Saved category heatmap to {out_name}")
+
+    print("="*40)
+
+if __name__ == "__main__":
+    import numpy as np
+    main()

complete_auditor_best.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fcd4d6b8c56b5842381a9469ee9ef1971c754a51a0cd940233844342be3aff90
+size 316574455