Predator911/MyData
0
1import os2import numpy as np3import torch4import torch.nn as nn5import torch.nn.functional as F6from torchvision import transforms7from PIL import Image8import re9from diffusers import StableDiffusionPipeline10import warnings11 12# Import improved utils13from stego_utils import (14 AESCipher,15 OBJECT_TABLE,16 COLOR_TABLE,17 STYLE_TABLE,18 MOOD_TABLE,19 INV_OBJECT_TABLE,20 INV_COLOR_TABLE,21 INV_STYLE_TABLE,22 INV_MOOD_TABLE,23 encode_binary_to_prompt,24 decode_prompt_to_binary,25 compress_message,26 validate_prompt_capacity,27 estimate_data_capacity28)29 30# Set device31device = torch.device("cuda" if torch.cuda.is_available() else "cpu")32 33class RGBStegoEncoder(nn.Module):34 """Improved GAN encoder with better capacity and stability."""35 def __init__(self, message_length=512):36 super(RGBStegoEncoder, self).__init__()37 self.message_length = message_length38 39 # Image processing branch40 self.img_conv1 = nn.Conv2d(3, 64, kernel_size=3, padding=1)41 self.img_bn1 = nn.BatchNorm2d(64)42 self.img_conv2 = nn.Conv2d(64, 128, kernel_size=3, padding=1)43 self.img_bn2 = nn.BatchNorm2d(128)44 self.img_conv3 = nn.Conv2d(128, 128, kernel_size=3, padding=1)45 self.img_bn3 = nn.BatchNorm2d(128)46 47 # Message processing branch48 self.msg_fc1 = nn.Linear(message_length, 1024)49 self.msg_fc2 = nn.Linear(1024, 2048)50 self.msg_fc3 = nn.Linear(2048, 4096)51 52 # Fusion layer53 self.fusion_conv = nn.Conv2d(128 + 16, 64, kernel_size=3, padding=1)54 self.fusion_bn = nn.BatchNorm2d(64)55 56 # Output layers with residual connection57 self.out_conv1 = nn.Conv2d(64, 32, kernel_size=3, padding=1)58 self.out_bn1 = nn.BatchNorm2d(32)59 self.out_conv2 = nn.Conv2d(32, 3, kernel_size=3, padding=1)60 61 # Adaptive scaling62 self.scale_factor = nn.Parameter(torch.tensor(0.1))63 64 self.dropout = nn.Dropout(0.1)65 66 def forward(self, image, message):67 batch_size, _, h, w = image.shape68 69 # Process image70 x = F.leaky_relu(self.img_bn1(self.img_conv1(image)), 0.2)71 x = F.leaky_relu(self.img_bn2(self.img_conv2(x)), 0.2)72 img_features = F.leaky_relu(self.img_bn3(self.img_conv3(x)), 0.2)73 74 # Process message75 msg = F.leaky_relu(self.msg_fc1(message))76 msg = self.dropout(msg)77 msg = F.leaky_relu(self.msg_fc2(msg))78 msg = self.dropout(msg)79 msg_features = F.leaky_relu(self.msg_fc3(msg))80 81 # Reshape message features to spatial format82 msg_features = msg_features.view(batch_size, 16, 16, 16)83 msg_features = F.interpolate(msg_features, size=(h, w), mode='bilinear', align_corners=False)84 85 # Fuse image and message features86 combined = torch.cat([img_features, msg_features], dim=1)87 fused = F.leaky_relu(self.fusion_bn(self.fusion_conv(combined)), 0.2)88 89 # Generate perturbation90 perturbation = F.leaky_relu(self.out_bn1(self.out_conv1(fused)), 0.2)91 perturbation = torch.tanh(self.out_conv2(perturbation))92 93 # Apply adaptive scaling and residual connection94 stego_image = image + self.scale_factor * perturbation95 stego_image = torch.clamp(stego_image, -1, 1)96 97 return stego_image98 99class RGBStegoDecoder(nn.Module):100 """Improved GAN decoder with attention mechanism."""101 def __init__(self, message_length=512):102 super(RGBStegoDecoder, self).__init__()103 self.message_length = message_length104 105 # Feature extraction106 self.conv1 = nn.Conv2d(3, 64, kernel_size=3, padding=1)107 self.bn1 = nn.BatchNorm2d(64)108 self.conv2 = nn.Conv2d(64, 128, kernel_size=3, padding=1)109 self.bn2 = nn.BatchNorm2d(128)110 self.conv3 = nn.Conv2d(128, 256, kernel_size=3, padding=1)111 self.bn3 = nn.BatchNorm2d(256)112 113 # Channel attention114 self.channel_att = nn.Sequential(115 nn.AdaptiveAvgPool2d(1),116 nn.Conv2d(256, 64, kernel_size=1),117 nn.ReLU(inplace=True),118 nn.Conv2d(64, 256, kernel_size=1),119 nn.Sigmoid()120 )121 122 # Spatial attention123 self.spatial_att = nn.Sequential(124 nn.Conv2d(256, 64, kernel_size=1),125 nn.ReLU(inplace=True),126 nn.Conv2d(64, 1, kernel_size=1),127 nn.Sigmoid()128 )129 130 # Message reconstruction131 self.global_pool = nn.AdaptiveAvgPool2d(8)132 self.fc1 = nn.Linear(256 * 8 * 8, 2048)133 self.dropout1 = nn.Dropout(0.3)134 self.fc2 = nn.Linear(2048, 1024)135 self.dropout2 = nn.Dropout(0.3)136 self.fc3 = nn.Linear(1024, message_length)137 138 def forward(self, stego_image):139 batch_size = stego_image.size(0)140 141 # Feature extraction142 x = F.leaky_relu(self.bn1(self.conv1(stego_image)), 0.2)143 x = F.max_pool2d(x, 2)144 x = F.leaky_relu(self.bn2(self.conv2(x)), 0.2)145 x = F.max_pool2d(x, 2)146 x = F.leaky_relu(self.bn3(self.conv3(x)), 0.2)147 148 # Apply attention mechanisms149 ch_att = self.channel_att(x)150 x = x * ch_att151 152 sp_att = self.spatial_att(x)153 x = x * sp_att154 155 # Global pooling and message reconstruction156 x = self.global_pool(x)157 x = x.view(batch_size, -1)158 159 x = F.relu(self.fc1(x))160 x = self.dropout1(x)161 x = F.relu(self.fc2(x))162 x = self.dropout2(x)163 message = torch.sigmoid(self.fc3(x))164 165 return message166 167class HybridStegoSystem:168 """Improved hybrid system with better capacity management."""169 def __init__(self, sd_model_id="stabilityai/stable-diffusion-2-1-base", gan_message_length=512):170 self.message_length = gan_message_length171 self.sd_model_id = sd_model_id172 self.sd_pipeline = None173 self.max_prompt_tokens = 60 # Conservative limit174 self.max_prompt_capacity = estimate_data_capacity(self.max_prompt_tokens)175 176 # Initialize GAN models177 self.encoder = RGBStegoEncoder(message_length=gan_message_length).to(device)178 self.decoder = RGBStegoDecoder(message_length=gan_message_length).to(device)179 self.encoder.eval()180 self.decoder.eval()181 182 print(f"Hybrid steganography system initialized on {device}")183 print(f"Prompt capacity: ~{self.max_prompt_capacity} bits")184 print(f"GAN capacity: {gan_message_length} bits")185 186 def _load_sd_pipeline(self):187 """Load Stable Diffusion pipeline with error handling."""188 if self.sd_pipeline is None:189 try:190 print("Loading Stable Diffusion model...")191 self.sd_pipeline = StableDiffusionPipeline.from_pretrained(192 self.sd_model_id,193 torch_dtype=torch.float16 if device.type == "cuda" else torch.float32,194 safety_checker=None,195 requires_safety_checker=False196 ).to(device)197 198 # Enable memory efficient attention if available199 if hasattr(self.sd_pipeline, 'enable_attention_slicing'):200 self.sd_pipeline.enable_attention_slicing()201 if hasattr(self.sd_pipeline, 'enable_xformers_memory_efficient_attention'):202 try:203 self.sd_pipeline.enable_xformers_memory_efficient_attention()204 except:205 pass206 207 print("Stable Diffusion model loaded successfully")208 except Exception as e:209 print(f"Error loading Stable Diffusion: {e}")210 raise211 212 def load_gan_models(self, encoder_path, decoder_path):213 """Load pre-trained GAN models."""214 try:215 self.encoder.load_state_dict(torch.load(encoder_path, map_location=device))216 self.decoder.load_state_dict(torch.load(decoder_path, map_location=device))217 self.encoder.eval()218 self.decoder.eval()219 print("GAN models loaded successfully")220 except Exception as e:221 print(f"Error loading GAN models: {e}")222 raise223 224 def _prepare_message_for_method(self, text, method):225 """Prepare message based on steganography method."""226 if method == "prompt":227 # Check if message fits in prompt capacity228 estimated_bits = len(text.encode('utf-8')) * 8229 if estimated_bits > self.max_prompt_capacity:230 compressed_text = compress_message(text, self.max_prompt_capacity)231 print(f"Message compressed from {len(text)} to {len(compressed_text)} characters")232 return compressed_text233 return text234 elif method == "gan":235 # Check if message fits in GAN capacity236 estimated_bits = len(text.encode('utf-8')) * 8237 if estimated_bits > self.message_length:238 max_chars = self.message_length // 8239 compressed_text = compress_message(text, self.message_length)240 print(f"Message truncated to fit GAN capacity ({max_chars} characters)")241 return compressed_text242 return text243 return text244 245 def encrypt_and_hide(self, text, key, cover_image=None, method="gan", seed=None):246 """Main encryption and hiding function with improved error handling."""247 try:248 # Prepare message249 text = self._prepare_message_for_method(text, method)250 251 # Encrypt252 cipher = AESCipher(key)253 encrypted_data = cipher.encrypt(text)254 binary_data = ''.join(format(byte, '08b') for byte in encrypted_data)255 256 if method == "prompt":257 return self._hide_with_prompt(binary_data, seed)258 elif method == "gan":259 if cover_image is None:260 raise ValueError("Cover image is required for GAN-based steganography")261 return self._hide_with_gan(cover_image, binary_data)262 else:263 raise ValueError("Method must be 'prompt' or 'gan'")264 265 except Exception as e:266 print(f"Error in encrypt_and_hide: {e}")267 raise268 269 def _hide_with_prompt(self, binary_data, seed=None):270 """Hide data using prompt-based method."""271 self._load_sd_pipeline()272 273 # Generate prompt with token limit consideration274 prompt = encode_binary_to_prompt(binary_data, max_tokens=self.max_prompt_tokens)275 276 # Validate prompt277 is_valid, token_count = validate_prompt_capacity(prompt, max_tokens=77)278 if not is_valid:279 print(f"Warning: Generated prompt has {token_count} tokens (>77). Image generation may fail.")280 281 # Generate image282 generator = torch.Generator(device=device).manual_seed(seed) if seed is not None else None283 284 with torch.no_grad():285 try:286 result = self.sd_pipeline(287 prompt,288 num_inference_steps=20, # Reduced for speed289 guidance_scale=7.5,290 generator=generator,291 height=512,292 width=512293 )294 image = result.images[0]295 296 return image, prompt297 298 except Exception as e:299 print(f"Image generation failed: {e}")300 # Fallback to simpler prompt301 simple_prompt = "a red cat, digital art"302 result = self.sd_pipeline(303 simple_prompt,304 num_inference_steps=20,305 guidance_scale=7.5,306 generator=generator307 )308 return result.images[0], prompt309 310 def _hide_with_gan(self, cover_image, binary_data, image_size=(256, 256)):311 """Hide data using GAN-based method."""312 # Prepare cover image313 if isinstance(cover_image, Image.Image):314 transform = transforms.Compose([315 transforms.Resize(image_size),316 transforms.ToTensor(),317 transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])318 ])319 cover_tensor = transform(cover_image).unsqueeze(0).to(device)320 else:321 raise ValueError("Cover image must be PIL Image")322 323 # Prepare message324 if isinstance(binary_data, str):325 message = torch.tensor([float(bit) for bit in binary_data], dtype=torch.float32)326 else:327 message = binary_data.float()328 329 # Pad or truncate message330 if len(message) > self.message_length:331 message = message[:self.message_length]332 elif len(message) < self.message_length:333 padding = torch.zeros(self.message_length - len(message))334 message = torch.cat([message, padding])335 336 message = message.unsqueeze(0).to(device)337 338 # Generate stego image339 with torch.no_grad():340 stego_tensor = self.encoder(cover_tensor, message)341 342 # Convert back to PIL Image343 stego_image = (stego_tensor.squeeze().permute(1, 2, 0).cpu().numpy() + 1.0) * 127.5344 stego_image = np.clip(stego_image, 0, 255).astype(np.uint8)345 346 return Image.fromarray(stego_image)347 348 def encrypt_and_hide_hybrid(self, text, key, cover_image=None, seed=None):349 """Hybrid method with intelligent data splitting."""350 try:351 cipher = AESCipher(key)352 encrypted_data = cipher.encrypt(text)353 binary_data = ''.join(format(byte, '08b') for byte in encrypted_data)354 355 # Calculate optimal split based on capacities356 prompt_capacity = self.max_prompt_capacity357 gan_capacity = self.message_length358 total_capacity = prompt_capacity + gan_capacity359 360 if len(binary_data) > total_capacity:361 print(f"Warning: Data exceeds hybrid capacity. Truncating from {len(binary_data)} to {total_capacity} bits")362 binary_data = binary_data[:total_capacity]363 364 # Smart splitting: more critical data in prompt (more reliable)365 prompt_data_len = min(len(binary_data) // 2, prompt_capacity)366 prompt_binary = binary_data[:prompt_data_len]367 gan_binary = binary_data[prompt_data_len:]368 369 # Hide first part in prompt370 stego_image, prompt = self._hide_with_prompt(prompt_binary, seed)371 372 # Use generated image as cover if none provided373 final_cover = cover_image if cover_image else stego_image374 375 # Hide second part in image376 if gan_binary:377 final_stego_image = self._hide_with_gan(final_cover, gan_binary)378 else:379 final_stego_image = stego_image380 381 return final_stego_image, prompt382 383 except Exception as e:384 print(f"Error in hybrid encryption: {e}")385 raise386 387 def extract_and_decrypt(self, stego_data, key, method="gan", prompt=None):388 """Extract and decrypt hidden message."""389 try:390 if method == "prompt":391 if prompt is None:392 raise ValueError("Prompt is required for prompt-based extraction")393 return self._extract_from_prompt(prompt, key)394 elif method == "gan":395 if not isinstance(stego_data, Image.Image):396 raise ValueError("Stego data must be PIL Image for GAN extraction")397 return self._extract_from_gan(stego_data, key)398 else:399 raise ValueError("Method must be 'prompt' or 'gan'")400 401 except Exception as e:402 print(f"Error in extract_and_decrypt: {e}")403 raise404 405 def _extract_from_prompt(self, prompt, key):406 """Extract data from prompt-based steganography."""407 try:408 # Decode binary data from prompt409 binary_data = decode_prompt_to_binary(prompt)410 411 if not binary_data:412 raise ValueError("No data found in prompt")413 414 # Convert binary to bytes415 byte_data = []416 for i in range(0, len(binary_data), 8):417 if i + 8 <= len(binary_data):418 byte_value = int(binary_data[i:i+8], 2)419 byte_data.append(byte_value)420 421 encrypted_data = bytes(byte_data)422 423 # Decrypt424 cipher = AESCipher(key)425 decrypted_text = cipher.decrypt(encrypted_data)426 427 return decrypted_text428 429 except Exception as e:430 print(f"Error extracting from prompt: {e}")431 raise432 433 def _extract_from_gan(self, stego_image, key, image_size=(256, 256)):434 """Extract data from GAN-based steganography."""435 try:436 # Prepare stego image437 transform = transforms.Compose([438 transforms.Resize(image_size),439 transforms.ToTensor(),440 transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])441 ])442 stego_tensor = transform(stego_image).unsqueeze(0).to(device)443 444 # Extract message445 with torch.no_grad():446 message_tensor = self.decoder(stego_tensor)447 448 # Convert to binary449 message_bits = (message_tensor.squeeze().cpu().numpy() > 0.5).astype(int)450 binary_data = ''.join(str(bit) for bit in message_bits)451 452 # Remove padding (find last meaningful bit)453 last_byte_idx = len(binary_data) - (len(binary_data) % 8)454 binary_data = binary_data[:last_byte_idx]455 456 # Convert to bytes457 byte_data = []458 for i in range(0, len(binary_data), 8):459 if i + 8 <= len(binary_data):460 byte_value = int(binary_data[i:i+8], 2)461 byte_data.append(byte_value)462 463 encrypted_data = bytes(byte_data)464 465 # Decrypt466 cipher = AESCipher(key)467 decrypted_text = cipher.decrypt(encrypted_data)468 469 return decrypted_text470 471 except Exception as e:472 print(f"Error extracting from GAN: {e}")473 raise474 475 def extract_and_decrypt_hybrid(self, stego_image, key, prompt):476 """Extract and decrypt from hybrid steganography."""477 try:478 # Extract from prompt (first part)479 prompt_text = self._extract_from_prompt(prompt, key + "_prompt")480 481 # Extract from GAN (second part) 482 gan_text = self._extract_from_gan(stego_image, key + "_gan")483 484 # Combine parts (implementation depends on how data was split)485 combined_text = prompt_text + gan_text486 487 return combined_text488 489 except Exception as e:490 print(f"Error in hybrid extraction: {e}")491 # Try individual methods as fallback492 try:493 return self._extract_from_prompt(prompt, key)494 except:495 try:496 return self._extract_from_gan(stego_image, key)497 except:498 raise e499 500 def calculate_capacity(self):501 """Calculate total system capacity."""502 return {503 'prompt_capacity_bits': self.max_prompt_capacity,504 'gan_capacity_bits': self.message_length,505 'hybrid_capacity_bits': self.max_prompt_capacity + self.message_length,506 'prompt_capacity_chars': self.max_prompt_capacity // 8,507 'gan_capacity_chars': self.message_length // 8,508 'hybrid_capacity_chars': (self.max_prompt_capacity + self.message_length) // 8509 }510 511 def benchmark_methods(self, test_text, key, cover_image=None, seed=42):512 """Benchmark different steganography methods."""513 results = {}514 515 try:516 # Test prompt method517 print("Testing prompt-based method...")518 start_time = torch.cuda.Event(enable_timing=True) if device.type == "cuda" else None519 end_time = torch.cuda.Event(enable_timing=True) if device.type == "cuda" else None520 521 if start_time:522 start_time.record()523 524 stego_img, prompt = self.encrypt_and_hide(test_text, key, method="prompt", seed=seed)525 extracted_text = self.extract_and_decrypt(None, key, method="prompt", prompt=prompt)526 527 if end_time:528 end_time.record()529 torch.cuda.synchronize()530 prompt_time = start_time.elapsed_time(end_time) / 1000.0531 else:532 prompt_time = 0533 534 results['prompt'] = {535 'success': extracted_text == test_text,536 'time_seconds': prompt_time,537 'capacity_used': len(test_text.encode('utf-8')) * 8538 }539 540 except Exception as e:541 results['prompt'] = {'success': False, 'error': str(e)}542 543 try:544 # Test GAN method545 if cover_image:546 print("Testing GAN-based method...")547 if start_time:548 start_time.record()549 550 stego_img = self.encrypt_and_hide(test_text, key, cover_image, method="gan")551 extracted_text = self.extract_and_decrypt(stego_img, key, method="gan")552 553 if end_time:554 end_time.record()555 torch.cuda.synchronize()556 gan_time = start_time.elapsed_time(end_time) / 1000.0557 else:558 gan_time = 0559 560 results['gan'] = {561 'success': extracted_text == test_text,562 'time_seconds': gan_time,563 'capacity_used': len(test_text.encode('utf-8')) * 8564 }565 else:566 results['gan'] = {'success': False, 'error': 'No cover image provided'}567 568 except Exception as e:569 results['gan'] = {'success': False, 'error': str(e)}570 571 return results572 573 def save_models(self, encoder_path, decoder_path):574 """Save trained GAN models."""575 try:576 torch.save(self.encoder.state_dict(), encoder_path)577 torch.save(self.decoder.state_dict(), decoder_path)578 print(f"Models saved to {encoder_path} and {decoder_path}")579 except Exception as e:580 print(f"Error saving models: {e}")581 raise