generate_remaining_scripts.py

"""
E-Raksha Training Script Generator

Automated generator for specialist model training scripts.
Creates comprehensive training pipelines for RR and TM models with
proper architecture definitions and training configurations.

Author: E-Raksha Team
"""

import os
from pathlib import Path

# Resolution/Re-recording specialist module architecture
RR_MODULE = '''class ResolutionModule(nn.Module):
    """
    Resolution artifact detection module for identifying upscaling and re-recording patterns.
    
    Detects:
    - Upscaling artifacts
    - Re-recording patterns  
    - Resolution inconsistencies
    - Interpolation artifacts
    """
    
    def __init__(self, in_channels=3):
        super().__init__()
        
        # Multi-scale resolution analyzer (10 channels)
        self.resolution_analyzer = nn.Sequential(
            nn.Conv2d(in_channels, 10, kernel_size=3, padding=1),
            nn.BatchNorm2d(10),
            nn.ReLU(),
            nn.Conv2d(10, 20, kernel_size=3, padding=1),
            nn.BatchNorm2d(20),
            nn.ReLU(),
            nn.Conv2d(20, 10, kernel_size=1)
        )
        
        # Upscaling artifact detector
        self.upscaling_detector = nn.Sequential(
            nn.Conv2d(in_channels, 8, kernel_size=5, padding=2),
            nn.BatchNorm2d(8),
            nn.ReLU(),
            nn.Conv2d(8, 16, kernel_size=5, padding=2),
            nn.BatchNorm2d(16),
            nn.ReLU(),
            nn.Conv2d(16, 8, kernel_size=1)
        )
        
        # Edge sharpness checker
        self.edge_checker = nn.Sequential(
            nn.Conv2d(in_channels, 8, kernel_size=3, padding=1),
            nn.BatchNorm2d(8),
            nn.ReLU(),
            nn.Conv2d(8, 16, kernel_size=3, padding=1),
            nn.BatchNorm2d(16),
            nn.ReLU(),
            nn.Conv2d(16, 8, kernel_size=1)
        )
        
        # Pixel interpolation detector
        self.interpolation_detector = nn.Sequential(
            nn.Conv2d(in_channels, 10, kernel_size=7, padding=3),
            nn.BatchNorm2d(10),
            nn.ReLU(),
            nn.Conv2d(10, 20, kernel_size=7, padding=3),
            nn.BatchNorm2d(20),
            nn.ReLU(),
            nn.Conv2d(20, 10, kernel_size=1)
        )
        
        # Attention fusion
        self.attention = nn.Sequential(
            nn.Conv2d(36, 32, kernel_size=1),
            nn.ReLU(),
            nn.Conv2d(32, 36, kernel_size=1),
            nn.Sigmoid()
        )
    
    def forward(self, x):
        # Extract features
        res_feat = self.resolution_analyzer(x)
        res_feat = F.adaptive_avg_pool2d(res_feat, (7, 7))
        
        upscale_feat = self.upscaling_detector(x)
        upscale_feat = F.adaptive_avg_pool2d(upscale_feat, (7, 7))
        
        edge_feat = self.edge_checker(x)
        edge_feat = F.adaptive_avg_pool2d(edge_feat, (7, 7))
        
        interp_feat = self.interpolation_detector(x)
        interp_feat = F.adaptive_avg_pool2d(interp_feat, (7, 7))
        
        # Combine and apply attention
        combined = torch.cat([res_feat, upscale_feat, edge_feat, interp_feat], dim=1)
        attention_weights = self.attention(combined)
        attended_features = combined * attention_weights
        
        return attended_features'''

# TM Module code
TM_MODULE = '''class TemporalModule(nn.Module):
    """Temporal inconsistency detection module"""
    
    def __init__(self, in_channels=3):
        super().__init__()
        
        # Frame-to-frame consistency checker
        self.frame_consistency = nn.Sequential(
            nn.Conv2d(in_channels, 14, kernel_size=5, padding=2),
            nn.BatchNorm2d(14),
            nn.ReLU(),
            nn.Conv2d(14, 28, kernel_size=5, padding=2),
            nn.BatchNorm2d(28),
            nn.ReLU(),
            nn.Conv2d(28, 14, kernel_size=1)
        )
        
        # Motion flow analyzer
        self.motion_analyzer = nn.Sequential(
            nn.Conv2d(in_channels, 12, kernel_size=7, padding=3),
            nn.BatchNorm2d(12),
            nn.ReLU(),
            nn.Conv2d(12, 24, kernel_size=7, padding=3),
            nn.BatchNorm2d(24),
            nn.ReLU(),
            nn.Conv2d(24, 12, kernel_size=1)
        )
        
        # Temporal artifact detector
        self.temporal_detector = nn.Sequential(
            nn.Conv2d(in_channels, 12, kernel_size=3, padding=1),
            nn.BatchNorm2d(12),
            nn.ReLU(),
            nn.Conv2d(12, 24, kernel_size=3, padding=1),
            nn.BatchNorm2d(24),
            nn.ReLU(),
            nn.Conv2d(24, 12, kernel_size=1)
        )
        
        # Optical flow validator
        self.optical_flow = nn.Sequential(
            nn.Conv2d(in_channels, 14, kernel_size=9, padding=4),
            nn.BatchNorm2d(14),
            nn.ReLU(),
            nn.Conv2d(14, 28, kernel_size=9, padding=4),
            nn.BatchNorm2d(28),
            nn.ReLU(),
            nn.Conv2d(28, 14, kernel_size=1)
        )
        
        # Attention fusion
        self.attention = nn.Sequential(
            nn.Conv2d(52, 32, kernel_size=1),
            nn.ReLU(),
            nn.Conv2d(32, 52, kernel_size=1),
            nn.Sigmoid()
        )
    
    def forward(self, x):
        # Extract features
        frame_feat = self.frame_consistency(x)
        frame_feat = F.adaptive_avg_pool2d(frame_feat, (7, 7))
        
        motion_feat = self.motion_analyzer(x)
        motion_feat = F.adaptive_avg_pool2d(motion_feat, (7, 7))
        
        temporal_feat = self.temporal_detector(x)
        temporal_feat = F.adaptive_avg_pool2d(temporal_feat, (7, 7))
        
        optical_feat = self.optical_flow(x)
        optical_feat = F.adaptive_avg_pool2d(optical_feat, (7, 7))
        
        # Combine and apply attention
        combined = torch.cat([frame_feat, motion_feat, temporal_feat, optical_feat], dim=1)
        attention_weights = self.attention(combined)
        attended_features = combined * attention_weights
        
        return attended_features'''

def generate_script(template_file, output_file, replacements):
    """Generate a new script from template with replacements"""
    print(f"Generating {output_file}...")
    
    with open(template_file, 'r', encoding='utf-8') as f:
        content = f.read()
    
    # Apply replacements
    for old, new in replacements.items():
        content = content.replace(old, new)
    
    with open(output_file, 'w', encoding='utf-8') as f:
        f.write(content)
    
    print(f"✅ Created {output_file}")

def main():
    print("🚀 Generating remaining specialist model training scripts...")
    print("="*80)
    
    # RR Model Scripts
    print("\n📊 Creating RR (Resolution) Model Scripts...")
    
    # RR Stage 1
    generate_script(
        'train_cm_stage1_faceforensics.py',
        'train_rr_stage1_faceforensics.py',
        {
            'COMPRESSION (CM)': 'RESOLUTION (RR)',
            'Compression': 'Resolution',
            'compression': 'resolution',
            'CM MODEL': 'RR MODEL',
            'CM Specialist': 'RR Specialist',
            'cm_specialist': 'rr_specialist',
            'cm_stage': 'rr_stage',
            'CompressionModule': 'ResolutionModule',
            '40 * 7 * 7  # 1960': '36 * 7 * 7  # 1764',
            'class CompressionModule(nn.Module):\n    """Compression artifact detection module"""': RR_MODULE
        }
    )
    
    # RR Stage 2
    generate_script(
        'train_cm_stage2_celebdf.py',
        'train_rr_stage2_celebdf.py',
        {
            'COMPRESSION (CM)': 'RESOLUTION (RR)',
            'Compression': 'Resolution',
            'compression': 'resolution',
            'CM MODEL': 'RR MODEL',
            'CM Specialist': 'RR Specialist',
            'cm_specialist': 'rr_specialist',
            'cm_stage': 'rr_stage',
            'cm-stage': 'rr-stage',
            'CompressionModule': 'ResolutionModule',
            '40 * 7 * 7  # 1960': '36 * 7 * 7  # 1764',
            'class CompressionModule(nn.Module):\n    """Compression artifact detection module"""': RR_MODULE
        }
    )
    
    # RR Stage 4
    generate_script(
        'train_cm_stage4_dfdc.py',
        'train_rr_stage4_dfdc.py',
        {
            'COMPRESSION (CM)': 'RESOLUTION (RR)',
            'Compression': 'Resolution',
            'compression': 'resolution',
            'CM MODEL': 'RR MODEL',
            'CM Specialist': 'RR Specialist',
            'cm_specialist': 'rr_specialist',
            'cm_stage': 'rr_stage',
            'cm-stage': 'rr-stage',
            'CompressionModule': 'ResolutionModule',
            '40 * 7 * 7  # 1960': '36 * 7 * 7  # 1764',
            'class CompressionModule(nn.Module):\n    """Compression artifact detection module"""': RR_MODULE
        }
    )
    
    # TM Model Scripts
    print("\n📊 Creating TM (Temporal) Model Scripts...")
    
    # TM Stage 1
    generate_script(
        'train_cm_stage1_faceforensics.py',
        'train_tm_stage1_faceforensics.py',
        {
            'COMPRESSION (CM)': 'TEMPORAL (TM)',
            'Compression': 'Temporal',
            'compression': 'temporal',
            'CM MODEL': 'TM MODEL',
            'CM Specialist': 'TM Specialist',
            'cm_specialist': 'tm_specialist',
            'cm_stage': 'tm_stage',
            'CompressionModule': 'TemporalModule',
            '40 * 7 * 7  # 1960': '52 * 7 * 7  # 2548',
            'class CompressionModule(nn.Module):\n    """Compression artifact detection module"""': TM_MODULE
        }
    )
    
    # TM Stage 2
    generate_script(
        'train_cm_stage2_celebdf.py',
        'train_tm_stage2_celebdf.py',
        {
            'COMPRESSION (CM)': 'TEMPORAL (TM)',
            'Compression': 'Temporal',
            'compression': 'temporal',
            'CM MODEL': 'TM MODEL',
            'CM Specialist': 'TM Specialist',
            'cm_specialist': 'tm_specialist',
            'cm_stage': 'tm_stage',
            'cm-stage': 'tm-stage',
            'CompressionModule': 'TemporalModule',
            '40 * 7 * 7  # 1960': '52 * 7 * 7  # 2548',
            'class CompressionModule(nn.Module):\n    """Compression artifact detection module"""': TM_MODULE
        }
    )
    
    # TM Stage 4
    generate_script(
        'train_cm_stage4_dfdc.py',
        'train_tm_stage4_dfdc.py',
        {
            'COMPRESSION (CM)': 'TEMPORAL (TM)',
            'Compression': 'Temporal',
            'compression': 'temporal',
            'CM MODEL': 'TM MODEL',
            'CM Specialist': 'TM Specialist',
            'cm_specialist': 'tm_specialist',
            'cm_stage': 'tm_stage',
            'cm-stage': 'tm-stage',
            'CompressionModule': 'TemporalModule',
            '40 * 7 * 7  # 1960': '52 * 7 * 7  # 2548',
            'class CompressionModule(nn.Module):\n    """Compression artifact detection module"""': TM_MODULE
        }
    )
    
    print("\n🎉 ALL SCRIPTS GENERATED SUCCESSFULLY!")
    print("="*80)
    print("\n📊 Summary:")
    print("✅ RR Model: 3 scripts created")
    print("✅ TM Model: 3 scripts created")
    print("✅ Total: 6 new scripts")
    print("\n🚀 All 15 specialist model training scripts are now complete!")

if __name__ == "__main__":
    main()