Clear repo for new ResNet101 auditor deployment

README.md

@@ -1,74 +0,0 @@
----
-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

@@ -1,305 +0,0 @@
-#!/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

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