Predator911/MyData
0
1import streamlit as st2import torch3import torch.nn as nn4import torch.nn.functional as F5import numpy as np6from PIL import Image7import torchvision.transforms as transforms8import io9import base6410import json11import hashlib12import matplotlib.pyplot as plt13import seaborn as sns14from datetime import datetime15import warnings16warnings.filterwarnings('ignore')17 18# Set page config19st.set_page_config(20 page_title="Prompt + GAN Steganography PoC",21 page_icon="🎭",22 layout="wide",23 initial_sidebar_state="expanded"24)25 26# Custom CSS for better styling27st.markdown("""28<style>29 .main-header {30 background: linear-gradient(90deg, #667eea 0%, #764ba2 100%);31 color: white;32 padding: 1rem;33 border-radius: 10px;34 text-align: center;35 margin-bottom: 2rem;36 }37 .metric-card {38 background: #f8f9fa;39 padding: 1rem;40 border-radius: 10px;41 border-left: 4px solid #667eea;42 margin: 0.5rem 0;43 }44 .success-box {45 background: #d4edda;46 color: #155724;47 padding: 1rem;48 border-radius: 10px;49 border-left: 4px solid #28a745;50 }51 .warning-box {52 background: #fff3cd;53 color: #856404;54 padding: 1rem;55 border-radius: 10px;56 border-left: 4px solid #ffc107;57 }58 .error-box {59 background: #f8d7da;60 color: #721c24;61 padding: 1rem;62 border-radius: 10px;63 border-left: 4px solid #dc3545;64 }65 .info-box {66 background: #d1ecf1;67 color: #0c5460;68 padding: 1rem;69 border-radius: 10px;70 border-left: 4px solid #17a2b8;71 }72 .stTabs [data-baseweb="tab-list"] {73 gap: 2px;74 }75 .stTabs [data-baseweb="tab"] {76 height: 50px;77 white-space: pre-wrap;78 background-color: #f0f2f6;79 border-radius: 4px 4px 0 0;80 color: #262730;81 font-size: 16px;82 font-weight: 500;83 }84 .stTabs [aria-selected="true"] {85 background-color: #667eea;86 color: white;87 }88</style>89""", unsafe_allow_html=True)90 91class StegoGANEncoder(nn.Module):92 """93 Fixed Encoder based on actual model analysis:94 - 6 Conv layers with batch normalization95 - Message processing layers (512->1024->2048->4096)96 - RGB channel processors (1->1 conv)97 - Spatial transformer (11 input channels->64->64->3)98 """99 def __init__(self):100 super().__init__()101 102 # Main convolutional layers (matches actual model)103 self.conv1 = nn.Conv2d(3, 64, 3, padding=1)104 self.bn1 = nn.BatchNorm2d(64)105 self.conv2 = nn.Conv2d(64, 64, 3, padding=1)106 self.bn2 = nn.BatchNorm2d(64)107 self.conv3 = nn.Conv2d(64, 128, 3, padding=1)108 self.bn3 = nn.BatchNorm2d(128)109 self.conv4 = nn.Conv2d(128, 128, 3, padding=1)110 self.bn4 = nn.BatchNorm2d(128)111 self.conv5 = nn.Conv2d(128, 64, 3, padding=1)112 self.bn5 = nn.BatchNorm2d(64)113 self.conv6 = nn.Conv2d(64, 3, 3, padding=1)114 115 # Message processing layers (matches actual dimensions)116 self.message_processor = nn.Sequential(117 nn.Linear(512, 1024),118 nn.ReLU(),119 nn.Linear(1024, 2048),120 nn.ReLU(),121 nn.Linear(2048, 4096),122 nn.ReLU()123 )124 125 # RGB channel processors (simplified to match actual model)126 self.r_channel_processor = nn.Conv2d(1, 1, 3, padding=1)127 self.g_channel_processor = nn.Conv2d(1, 1, 3, padding=1)128 self.b_channel_processor = nn.Conv2d(1, 1, 3, padding=1)129 130 # Spatial transformer (matches actual 11 input channels)131 self.spatial_transformer = nn.Sequential(132 nn.Conv2d(11, 64, 3, padding=1), # 3 (original) + 3 (RGB) + 5 (features)133 nn.ReLU(),134 nn.Conv2d(64, 64, 3, padding=1),135 nn.ReLU(),136 nn.Conv2d(64, 3, 3, padding=1)137 )138 139 def forward(self, x, message_features=None, prompt_features=None):140 batch_size, _, height, width = x.shape141 142 # Main convolution path143 x1 = F.relu(self.bn1(self.conv1(x)))144 x2 = F.relu(self.bn2(self.conv2(x1)))145 x3 = F.relu(self.bn3(self.conv3(x2)))146 x4 = F.relu(self.bn4(self.conv4(x3)))147 x5 = F.relu(self.bn5(self.conv5(x4)))148 x6 = self.conv6(x5)149 150 # If features are provided, process them151 if message_features is not None:152 # Process message features153 msg_processed = self.message_processor(message_features)154 155 # Process RGB channels separately156 r_processed = self.r_channel_processor(x[:, 0:1, :, :])157 g_processed = self.g_channel_processor(x[:, 1:2, :, :])158 b_processed = self.b_channel_processor(x[:, 2:3, :, :])159 160 # Create feature maps from processed message161 # Reshape to spatial dimensions162 feature_size = min(5, msg_processed.shape[1] // (height * width))163 if feature_size > 0:164 spatial_features = msg_processed[:, :feature_size * height * width].view(165 batch_size, feature_size, height, width166 )167 else:168 spatial_features = torch.zeros(batch_size, 5, height, width).to(x.device)169 170 # Combine all features for spatial transformer (total 11 channels)171 combined = torch.cat([172 x, # Original image (3 channels)173 r_processed, # R channel (1 channel)174 g_processed, # G channel (1 channel)175 b_processed, # B channel (1 channel)176 spatial_features # Message features (5 channels)177 ], dim=1)178 179 # Apply spatial transformer180 output = self.spatial_transformer(combined)181 return torch.tanh(output)182 183 return torch.tanh(x6)184 185class StegoGANDecoder(nn.Module):186 """187 Fixed Decoder based on actual model analysis:188 - 6 Conv layers with batch normalization189 - 3 Linear layers (16384->2048->1024->512)190 - Matches actual parameter counts191 """192 def __init__(self):193 super().__init__()194 195 # Convolutional layers (matches actual model)196 self.conv1 = nn.Conv2d(3, 64, 3, padding=1)197 self.bn1 = nn.BatchNorm2d(64)198 self.conv2 = nn.Conv2d(64, 64, 3, padding=1)199 self.bn2 = nn.BatchNorm2d(64)200 self.conv3 = nn.Conv2d(64, 128, 3, padding=1)201 self.bn3 = nn.BatchNorm2d(128)202 self.conv4 = nn.Conv2d(128, 128, 3, padding=1)203 self.bn4 = nn.BatchNorm2d(128)204 self.conv5 = nn.Conv2d(128, 64, 3, padding=1)205 self.bn5 = nn.BatchNorm2d(64)206 self.conv6 = nn.Conv2d(64, 64, 3, padding=1)207 self.bn6 = nn.BatchNorm2d(64)208 209 # Linear layers (matches actual dimensions: 16384->2048->1024->512)210 self.fc1 = nn.Linear(16384, 2048) # 64 * 16 * 16 = 16384211 self.fc2 = nn.Linear(2048, 1024)212 self.fc3 = nn.Linear(1024, 512)213 214 # Global average pooling to match expected input size215 self.adaptive_pool = nn.AdaptiveAvgPool2d((16, 16))216 217 def forward(self, x):218 # Convolutional feature extraction219 x1 = F.relu(self.bn1(self.conv1(x)))220 x2 = F.relu(self.bn2(self.conv2(x1)))221 x3 = F.relu(self.bn3(self.conv3(x2)))222 x4 = F.relu(self.bn4(self.conv4(x3)))223 x5 = F.relu(self.bn5(self.conv5(x4)))224 x6 = F.relu(self.bn6(self.conv6(x5)))225 226 # Pool and flatten to match linear layer input227 pooled = self.adaptive_pool(x6) # 64 x 16 x 16228 flattened = pooled.view(pooled.size(0), -1) # 16384229 230 # Linear layers231 features = F.relu(self.fc1(flattened))232 features = F.relu(self.fc2(features))233 features = self.fc3(features)234 235 # Split features for message and prompt (simplified)236 message_features = features # All 512 features for message237 prompt_features = features[:, :256] # First 256 for prompt238 239 return message_features, prompt_features240 241@st.cache_resource242def load_models():243 """Load the trained models with proper error handling"""244 try:245 encoder = StegoGANEncoder()246 decoder = StegoGANDecoder()247 248 model_loaded = False249 model_info = {}250 251 # Try to load encoder weights252 try:253 encoder_checkpoint = torch.load('stego_gan_encoder.pth', map_location='cpu')254 encoder.load_state_dict(encoder_checkpoint, strict=False)255 model_info['encoder'] = {256 'loaded': True,257 'size': '43.5 MB',258 'params': '55 groups, 16 layers'259 }260 st.success("✅ Encoder loaded successfully!")261 except FileNotFoundError:262 model_info['encoder'] = {'loaded': False, 'error': 'File not found'}263 st.warning("⚠️ Encoder file not found - using random weights")264 except Exception as e:265 model_info['encoder'] = {'loaded': False, 'error': str(e)}266 st.error(f"❌ Error loading encoder: {e}")267 268 # Try to load decoder weights269 try:270 decoder_checkpoint = torch.load('stego_gan_decoder.pth', map_location='cpu')271 decoder.load_state_dict(decoder_checkpoint, strict=False)272 model_info['decoder'] = {273 'loaded': True,274 'size': '139.6 MB',275 'params': '60 groups, 20 layers'276 }277 st.success("✅ Decoder loaded successfully!")278 except FileNotFoundError:279 model_info['decoder'] = {'loaded': False, 'error': 'File not found'}280 st.warning("⚠️ Decoder file not found - using random weights")281 except Exception as e:282 model_info['decoder'] = {'loaded': False, 'error': str(e)}283 st.error(f"❌ Error loading decoder: {e}")284 285 model_loaded = model_info.get('encoder', {}).get('loaded', False) and \286 model_info.get('decoder', {}).get('loaded', False)287 288 encoder.eval()289 decoder.eval()290 291 return encoder, decoder, model_loaded, model_info292 293 except Exception as e:294 st.error(f"Critical error in model loading: {e}")295 return None, None, False, {}296 297def text_to_features(text, feature_dim=512):298 """Convert text to feature vector with improved encoding"""299 if not text:300 return torch.zeros(1, feature_dim)301 302 # Create a more robust feature representation303 text_hash = hashlib.sha256(text.encode('utf-8')).hexdigest()304 text_bytes = text.encode('utf-8')305 306 features = np.zeros(feature_dim)307 308 # Use hash for initial seeding (first 32 positions)309 hash_bytes = bytes.fromhex(text_hash)310 for i, byte in enumerate(hash_bytes[:min(32, feature_dim)]):311 features[i] = (byte - 128) / 128.0 # Normalize to [-1, 1]312 313 # Add actual text content (next positions)314 for i, byte in enumerate(text_bytes[:min(feature_dim - 50, len(text_bytes))]):315 if i + 32 < feature_dim - 20:316 features[i + 32] = (byte - 128) / 128.0317 318 # Add text statistics (last 20 positions)319 if feature_dim >= 20:320 stats_start = feature_dim - 20321 features[stats_start:stats_start + 5] = [322 min(len(text) / 1000.0, 1.0), # Text length (normalized)323 text.count(' ') / max(len(text), 1), # Space ratio324 text.count('\n') / max(len(text), 1), # Newline ratio325 len(set(text)) / max(len(text), 1), # Unique char ratio326 sum(c.isupper() for c in text) / max(len(text), 1) # Uppercase ratio327 ]328 329 return torch.tensor(features, dtype=torch.float32).unsqueeze(0)330 331def prompt_to_features(prompt, feature_dim=256):332 """Convert prompt to feature vector with keyword detection"""333 if not prompt:334 return torch.zeros(1, feature_dim)335 336 prompt_lower = prompt.lower()337 features = np.zeros(feature_dim)338 339 # Enhanced keywords for prompt classification340 keywords = [341 'hide', 'secret', 'stealth', 'invisible', 'embed', 'conceal',342 'subtle', 'obvious', 'strong', 'weak', 'noise', 'clear',343 'robust', 'fragile', 'secure', 'visible', 'artifact', 'quality',344 'minimal', 'maximum', 'compression', 'resistant'345 ]346 347 # Set keyword presence flags348 for i, keyword in enumerate(keywords[:min(len(keywords), feature_dim // 2)]):349 if keyword in prompt_lower:350 features[i] = 1.0351 352 # Add prompt hash for uniqueness353 prompt_hash = hashlib.md5(prompt.encode('utf-8')).hexdigest()354 hash_bytes = bytes.fromhex(prompt_hash)355 hash_start = len(keywords)356 for i, byte in enumerate(hash_bytes[:min(16, feature_dim - hash_start)]):357 if hash_start + i < feature_dim:358 features[hash_start + i] = (byte - 128) / 128.0359 360 # Add prompt characteristics361 char_start = hash_start + 16362 if char_start + 10 < feature_dim:363 features[char_start:char_start + 5] = [364 min(len(prompt) / 100.0, 1.0), # Length365 prompt.count(' ') / max(len(prompt), 1), # Word density366 sum(c.isupper() for c in prompt) / max(len(prompt), 1), # Uppercase367 sum(c.isdigit() for c in prompt) / max(len(prompt), 1), # Digits368 len(set(prompt)) / max(len(prompt), 1) # Uniqueness369 ]370 371 return torch.tensor(features, dtype=torch.float32).unsqueeze(0)372 373def features_to_text(features, max_length=1000):374 """Convert feature vector back to text with improved decoding"""375 if features is None:376 return "No message found"377 378 features = features.squeeze().cpu().numpy()379 380 # Try to decode from the content section381 text_bytes = []382 start_idx = 32383 end_idx = min(len(features) - 20, max_length + start_idx)384 385 for i in range(start_idx, end_idx):386 if i < len(features):387 # Convert back from normalized range388 byte_val = int((features[i] * 128) + 128)389 if 0 <= byte_val <= 255 and byte_val != 0:390 text_bytes.append(byte_val)391 else:392 break393 394 if not text_bytes:395 return "Unable to decode message - no valid content found"396 397 try:398 decoded = bytes(text_bytes).decode('utf-8', errors='ignore')399 # Clean up the decoded text400 decoded = decoded.strip('\x00\x01\x02\x03\x04\x05\x06\x07\x08')401 return decoded if decoded.strip() else "Message decoded but appears empty"402 except Exception as e:403 return f"Decoding error: {str(e)}"404 405def features_to_prompt(features):406 """Convert feature vector back to prompt with keyword detection"""407 if features is None:408 return "No prompt found"409 410 features = features.squeeze().cpu().numpy()411 412 # Check for keyword indicators413 keywords = [414 'hide', 'secret', 'stealth', 'invisible', 'embed', 'conceal',415 'subtle', 'obvious', 'strong', 'weak', 'noise', 'clear',416 'robust', 'fragile', 'secure', 'visible', 'artifact', 'quality',417 'minimal', 'maximum', 'compression', 'resistant'418 ]419 420 detected_keywords = []421 for i, keyword in enumerate(keywords[:min(len(keywords), len(features))]):422 if features[i] > 0.5:423 detected_keywords.append(keyword)424 425 if detected_keywords:426 return f"Style keywords: {', '.join(detected_keywords)}"427 else:428 # Try to extract some characteristics429 if len(features) > 50:430 length_indicator = features[32] if len(features) > 32 else 0431 if length_indicator > 0.5:432 return "Long prompt detected"433 elif length_indicator > 0.2:434 return "Medium prompt detected"435 else:436 return "Short prompt detected"437 return "Prompt style unclear"438 439def preprocess_image(image, target_size=(256, 256)):440 """Preprocess image for model input with proper resizing"""441 # Convert to RGB if needed442 if image.mode != 'RGB':443 image = image.convert('RGB')444 445 # FIXED: Use resize instead of thumbnail to ensure exact dimensions446 image = image.resize(target_size, Image.Resampling.LANCZOS)447 448 # Convert to tensor with normalization449 transform = transforms.Compose([450 transforms.ToTensor(),451 transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])452 ])453 return transform(image).unsqueeze(0)454 455def postprocess_image(tensor):456 """Convert tensor back to PIL Image with proper denormalization"""457 tensor = tensor.squeeze(0).detach()458 459 # Denormalize from [-1, 1] to [0, 1]460 tensor = tensor * 0.5 + 0.5461 tensor = torch.clamp(tensor, 0, 1)462 463 # Convert to PIL Image464 transform = transforms.ToPILImage()465 return transform(tensor)466 467def calculate_metrics(img1, img2):468 """Calculate comprehensive image quality metrics with proper shape handling"""469 # Ensure both images are the same size470 if img1.size != img2.size:471 # Resize the second image to match the first472 img2 = img2.resize(img1.size, Image.Resampling.LANCZOS)473 474 img1_np = np.array(img1).astype(np.float64)475 img2_np = np.array(img2).astype(np.float64)476 477 # Verify shapes match478 if img1_np.shape != img2_np.shape:479 raise ValueError(f"Shape mismatch: {img1_np.shape} vs {img2_np.shape}")480 481 # Mean Squared Error482 mse = np.mean((img1_np - img2_np) ** 2)483 484 # Peak Signal-to-Noise Ratio485 if mse == 0:486 psnr = float('inf')487 else:488 max_pixel = 255.0489 psnr = 20 * np.log10(max_pixel / np.sqrt(mse))490 491 # Simplified SSIM calculation492 def calculate_ssim(img1, img2):493 mu1 = np.mean(img1)494 mu2 = np.mean(img2)495 sigma1 = np.var(img1)496 sigma2 = np.var(img2)497 sigma12 = np.mean((img1 - mu1) * (img2 - mu2))498 499 c1 = (0.01 * 255) ** 2500 c2 = (0.03 * 255) ** 2501 502 ssim_val = ((2 * mu1 * mu2 + c1) * (2 * sigma12 + c2)) / \503 ((mu1 ** 2 + mu2 ** 2 + c1) * (sigma1 + sigma2 + c2))504 return max(0, min(1, ssim_val))505 506 ssim_val = calculate_ssim(img1_np, img2_np)507 508 # Histogram difference509 hist1 = np.histogram(img1_np, bins=256, range=(0, 255))[0]510 hist2 = np.histogram(img2_np, bins=256, range=(0, 255))[0]511 hist_diff = np.sum(np.abs(hist1 - hist2)) / np.sum(hist1 + hist2)512 513 # Additional metrics514 mae = np.mean(np.abs(img1_np - img2_np)) # Mean Absolute Error515 516 return {517 'psnr': psnr,518 'ssim': ssim_val,519 'mse': mse,520 'mae': mae,521 'hist_diff': hist_diff522 }523 524def create_analysis_plot(cover_img, stego_img):525 """Create comprehensive analysis plots with proper size handling"""526 # Ensure both images are the same size527 if cover_img.size != stego_img.size:528 stego_img = stego_img.resize(cover_img.size, Image.Resampling.LANCZOS)529 530 fig, axes = plt.subplots(2, 3, figsize=(15, 10))531 532 # Convert to numpy arrays533 cover_np = np.array(cover_img)534 stego_np = np.array(stego_img)535 536 # Original images537 axes[0, 0].imshow(cover_img)538 axes[0, 0].set_title('Cover Image', fontsize=12, fontweight='bold')539 axes[0, 0].axis('off')540 541 axes[0, 1].imshow(stego_img)542 axes[0, 1].set_title('Stego Image', fontsize=12, fontweight='bold')543 axes[0, 1].axis('off')544 545 # Difference image (enhanced)546 diff = np.abs(cover_np.astype(np.float32) - stego_np.astype(np.float32))547 diff_enhanced = np.clip(diff * 5, 0, 255).astype(np.uint8) # Enhance differences548 axes[0, 2].imshow(diff_enhanced)549 axes[0, 2].set_title('Difference (5x Enhanced)', fontsize=12, fontweight='bold')550 axes[0, 2].axis('off')551 552 # Histograms553 colors = ['red', 'blue']554 labels = ['Cover', 'Stego']555 556 for i, (img, color, label) in enumerate(zip([cover_np, stego_np], colors, labels)):557 axes[1, i].hist(img.flatten(), bins=50, alpha=0.7, color=color, label=label)558 axes[1, i].set_title(f'{label} Histogram', fontsize=12, fontweight='bold')559 axes[1, i].set_xlabel('Pixel Value')560 axes[1, i].set_ylabel('Frequency')561 axes[1, i].grid(True, alpha=0.3)562 563 # Difference histogram564 axes[1, 2].hist(diff.flatten(), bins=50, alpha=0.7, color='green')565 axes[1, 2].set_title('Difference Distribution', fontsize=12, fontweight='bold')566 axes[1, 2].set_xlabel('Difference Value')567 axes[1, 2].set_ylabel('Frequency')568 axes[1, 2].grid(True, alpha=0.3)569 570 plt.tight_layout()571 return fig572 573def main():574 # Header575 st.markdown("""576 <div class="main-header">577 <h1>🎭 Prompt + GAN Steganography</h1>578 <p>Next-Generation Image Steganography with Intelligent Prompt Integration</p>579 <p><small>✨ Fixed Shape Handling & Architecture Matching</small></p>580 </div>581 """, unsafe_allow_html=True)582 583 # Load models with detailed information584 encoder, decoder, model_loaded, model_info = load_models()585 586 if encoder is None or decoder is None:587 st.markdown("""588 <div class="error-box">589 <h3>⚠️ Critical Error</h3>590 <p>Unable to initialize models. Please check the error messages above.</p>591 </div>592 """, unsafe_allow_html=True)593 return594 595 # Display model status596 if model_loaded:597 st.markdown(f"""598 <div class="success-box">599 <h3>✅ Models Loaded Successfully</h3>600 <p><strong>Encoder:</strong> {model_info['encoder']['size']}, {model_info['encoder']['params']}</p>601 <p><strong>Decoder:</strong> {model_info['decoder']['size']}, {model_info['decoder']['params']}</p>602 <p>All models loaded with pre-trained weights and ready for operation!</p>603 </div>604 """, unsafe_allow_html=True)605 else:606 st.markdown("""607 <div class="warning-box">608 <h3>⚠️ Demo Mode</h3>609 <p>Running with randomly initialized weights. For best results, ensure model files are present:</p>610 <ul>611 <li><code>stego_gan_encoder.pth</code> (Expected: 43.5 MB)</li>612 <li><code>stego_gan_decoder.pth</code> (Expected: 139.6 MB)</li>613 </ul>614 </div>615 """, unsafe_allow_html=True)616 617 # Sidebar configuration618 st.sidebar.title("🎛️ Configuration")619 620 # Model information621 with st.sidebar.expander("📊 Architecture Details"):622 st.markdown(f"""623 **Encoder Architecture:**624 - 6 Convolutional layers with BatchNorm625 - Message processor: 512→1024→2048→4096626 - RGB channel processors: 1→1 conv627 - Spatial transformer: 11→64→64→3628 - Total params: ~{model_info.get('encoder', {}).get('params', 'Unknown')}629 630 **Decoder Architecture:**631 - 6 Convolutional layers with BatchNorm 632 - Linear layers: 16384→2048→1024→512633 - Adaptive pooling: Any size → 16×16634 - Total params: ~{model_info.get('decoder', {}).get('params', 'Unknown')}635 636 **Key Features:**637 - ✅ Architecture matches your trained models638 - ✅ Proper parameter dimensions639 - ✅ Compatible layer structures640 - ✅ Fixed shape handling issues641 """)642 643 # Operation modes644 tab1, tab2, tab3, tab4 = st.tabs([645 "🔐 Encode Message", 646 "🔍 Decode Message", 647 "📊 Analysis", 648 "🧪 Experiments"649 ])650 651 with tab1:652 st.header("🔐 Encode Secret Message with Prompt")653 654 col1, col2 = st.columns([1, 1])655 656 with col1:657 st.subheader("📥 Input")658 659 # Upload cover image660 # Upload cover image661 cover_file = st.file_uploader(662 "Choose a cover image...",663 type=['png', 'jpg', 'jpeg'],664 key="cover_upload"665 )666 667 if cover_file is not None:668 cover_img = Image.open(cover_file)669 st.image(cover_img, caption="Cover Image", use_column_width=True)670 671 # Display image info672 st.info(f"📏 Size: {cover_img.size[0]}×{cover_img.size[1]} | Mode: {cover_img.mode}")673 674 # Secret message input675 secret_message = st.text_area(676 "🔤 Secret Message",677 placeholder="Enter your secret message here...",678 height=100,679 max_chars=500680 )681 682 # Steganography prompt683 stego_prompt = st.text_area(684 "🎯 Steganography Prompt",685 placeholder="hide message subtly with minimal artifacts...",686 height=80,687 max_chars=200,688 help="Describe how the message should be hidden (e.g., 'hide strongly', 'minimal changes', 'robust to compression')"689 )690 691 # Encoding parameters692 st.subheader("⚙️ Parameters")693 strength = st.slider("Encoding Strength", 0.1, 2.0, 1.0, 0.1)694 use_prompt = st.checkbox("Use Prompt Guidance", value=True)695 696 with col2:697 st.subheader("📤 Output")698 699 if st.button("🚀 Encode Message", type="primary", use_container_width=True):700 if cover_file is not None and secret_message:701 with st.spinner("Encoding message..."):702 try:703 # Preprocess image704 cover_tensor = preprocess_image(cover_img)705 706 # Generate features707 message_features = text_to_features(secret_message)708 prompt_features = prompt_to_features(stego_prompt if use_prompt else "")709 710 # Encode711 with torch.no_grad():712 if use_prompt:713 stego_tensor = encoder(cover_tensor, message_features, prompt_features)714 else:715 stego_tensor = encoder(cover_tensor, message_features)716 717 # Convert back to image718 stego_img = postprocess_image(stego_tensor)719 720 # Display result721 st.image(stego_img, caption="Steganographic Image", use_column_width=True)722 723 # Calculate metrics724 metrics = calculate_metrics(cover_img, stego_img)725 726 # Display metrics727 col_m1, col_m2 = st.columns(2)728 with col_m1:729 st.metric("PSNR", f"{metrics['psnr']:.2f} dB")730 st.metric("SSIM", f"{metrics['ssim']:.4f}")731 with col_m2:732 st.metric("MSE", f"{metrics['mse']:.2f}")733 st.metric("MAE", f"{metrics['mae']:.2f}")734 735 # Download button736 buf = io.BytesIO()737 stego_img.save(buf, format='PNG')738 st.download_button(739 label="💾 Download Stego Image",740 data=buf.getvalue(),741 file_name=f"stego_{datetime.now().strftime('%Y%m%d_%H%M%S')}.png",742 mime="image/png"743 )744 745 # Store in session state for analysis746 st.session_state.cover_img = cover_img747 st.session_state.stego_img = stego_img748 st.session_state.secret_message = secret_message749 st.session_state.stego_prompt = stego_prompt750 751 except Exception as e:752 st.error(f"❌ Encoding failed: {str(e)}")753 st.error("Please check your inputs and try again.")754 else:755 st.warning("⚠️ Please upload a cover image and enter a secret message.")756 757 with tab2:758 st.header("🔍 Decode Hidden Message")759 760 col1, col2 = st.columns([1, 1])761 762 with col1:763 st.subheader("📥 Input")764 765 # Upload stego image766 stego_file = st.file_uploader(767 "Choose a steganographic image...",768 type=['png', 'jpg', 'jpeg'],769 key="stego_upload"770 )771 772 if stego_file is not None:773 stego_img = Image.open(stego_file)774 st.image(stego_img, caption="Steganographic Image", use_column_width=True)775 st.info(f"📏 Size: {stego_img.size[0]}×{stego_img.size[1]} | Mode: {stego_img.mode}")776 777 with col2:778 st.subheader("📤 Output")779 780 if st.button("🔓 Decode Message", type="primary", use_container_width=True):781 if stego_file is not None:782 with st.spinner("Decoding message..."):783 try:784 # Preprocess image785 stego_tensor = preprocess_image(stego_img)786 787 # Decode788 with torch.no_grad():789 message_features, prompt_features = decoder(stego_tensor)790 791 # Convert features back to text792 decoded_message = features_to_text(message_features)793 decoded_prompt = features_to_prompt(prompt_features)794 795 # Display results796 st.success("✅ Decoding completed!")797 798 st.subheader("🔤 Decoded Message:")799 st.text_area("", decoded_message, height=100, key="decoded_msg")800 801 st.subheader("🎯 Detected Prompt Style:")802 st.text_area("", decoded_prompt, height=60, key="decoded_prompt")803 804 # Confidence indicators805 st.subheader("📊 Decoding Confidence")806 807 # Simple confidence estimation based on feature variance808 msg_confidence = min(1.0, np.var(message_features.numpy()) * 10)809 prompt_confidence = min(1.0, np.var(prompt_features.numpy()) * 10)810 811 col_c1, col_c2 = st.columns(2)812 with col_c1:813 st.metric("Message Confidence", f"{msg_confidence:.2%}")814 with col_c2:815 st.metric("Prompt Confidence", f"{prompt_confidence:.2%}")816 817 except Exception as e:818 st.error(f"❌ Decoding failed: {str(e)}")819 st.error("The image may not contain a hidden message or may be corrupted.")820 else:821 st.warning("⚠️ Please upload a steganographic image.")822 823 with tab3:824 st.header("📊 Steganographic Analysis")825 826 if hasattr(st.session_state, 'cover_img') and hasattr(st.session_state, 'stego_img'):827 cover_img = st.session_state.cover_img828 stego_img = st.session_state.stego_img829 830 # Analysis options831 analysis_type = st.selectbox(832 "Choose Analysis Type:",833 ["Visual Comparison", "Statistical Analysis", "Histogram Analysis", "Difference Analysis"]834 )835 836 if analysis_type == "Visual Comparison":837 col1, col2 = st.columns(2)838 with col1:839 st.subheader("Cover Image")840 st.image(cover_img, use_column_width=True)841 with col2:842 st.subheader("Stego Image")843 st.image(stego_img, use_column_width=True)844 845 # Side-by-side comparison846 st.subheader("Side-by-Side Comparison")847 comparison_img = Image.new('RGB', (cover_img.width * 2, cover_img.height))848 comparison_img.paste(cover_img, (0, 0))849 comparison_img.paste(stego_img, (cover_img.width, 0))850 st.image(comparison_img, caption="Cover (Left) vs Stego (Right)", use_column_width=True)851 852 elif analysis_type == "Statistical Analysis":853 metrics = calculate_metrics(cover_img, stego_img)854 855 st.subheader("📈 Quality Metrics")856 857 col1, col2, col3 = st.columns(3)858 with col1:859 st.metric("PSNR", f"{metrics['psnr']:.2f} dB", 860 help="Peak Signal-to-Noise Ratio (higher is better)")861 st.metric("SSIM", f"{metrics['ssim']:.4f}", 862 help="Structural Similarity Index (1.0 = identical)")863 with col2:864 st.metric("MSE", f"{metrics['mse']:.2f}", 865 help="Mean Squared Error (lower is better)")866 st.metric("MAE", f"{metrics['mae']:.2f}", 867 help="Mean Absolute Error (lower is better)")868 with col3:869 st.metric("Histogram Diff", f"{metrics['hist_diff']:.4f}", 870 help="Histogram difference (lower is better)")871 872 # Quality assessment873 st.subheader("🏆 Quality Assessment")874 if metrics['psnr'] > 40:875 quality = "Excellent"876 color = "success"877 elif metrics['psnr'] > 30:878 quality = "Good"879 color = "info"880 elif metrics['psnr'] > 20:881 quality = "Fair"882 color = "warning"883 else:884 quality = "Poor"885 color = "error"886 887 st.markdown(f"""888 <div class="{color}-box">889 <h4>Overall Quality: {quality}</h4>890 <p>PSNR: {metrics['psnr']:.2f} dB | SSIM: {metrics['ssim']:.4f}</p>891 </div>892 """, unsafe_allow_html=True)893 894 elif analysis_type == "Histogram Analysis":895 fig = create_analysis_plot(cover_img, stego_img)896 st.pyplot(fig)897 898 elif analysis_type == "Difference Analysis":899 # Calculate and display difference900 cover_np = np.array(cover_img)901 stego_np = np.array(stego_img.resize(cover_img.size))902 903 diff = np.abs(cover_np.astype(np.float32) - stego_np.astype(np.float32))904 diff_enhanced = np.clip(diff * 5, 0, 255).astype(np.uint8)905 906 col1, col2 = st.columns(2)907 with col1:908 st.subheader("Raw Difference")909 st.image(Image.fromarray(diff.astype(np.uint8)), use_column_width=True)910 with col2:911 st.subheader("Enhanced Difference (5x)")912 st.image(Image.fromarray(diff_enhanced), use_column_width=True)913 914 # Difference statistics915 st.subheader("📊 Difference Statistics")916 col1, col2, col3 = st.columns(3)917 with col1:918 st.metric("Max Difference", f"{np.max(diff):.2f}")919 with col2:920 st.metric("Mean Difference", f"{np.mean(diff):.2f}")921 with col3:922 st.metric("Std Difference", f"{np.std(diff):.2f}")923 924 else:925 st.info("🔍 Encode a message first to see analysis results.")926 927 with tab4:928 st.header("🧪 Experimental Features")929 930 # Batch processing931 st.subheader("📦 Batch Processing")932 933 batch_files = st.file_uploader(934 "Upload multiple images for batch processing:",935 type=['png', 'jpg', 'jpeg'],936 accept_multiple_files=True,937 key="batch_upload"938 )939 940 if batch_files:941 batch_message = st.text_input("Message for all images:", "Secret batch message")942 batch_prompt = st.text_input("Prompt for all images:", "hide subtly")943 944 if st.button("🚀 Process Batch"):945 progress_bar = st.progress(0)946 results = []947 948 for i, file in enumerate(batch_files):949 with st.spinner(f"Processing {file.name}..."):950 try:951 img = Image.open(file)952 img_tensor = preprocess_image(img)953 954 msg_features = text_to_features(batch_message)955 prompt_features = prompt_to_features(batch_prompt)956 957 with torch.no_grad():958 stego_tensor = encoder(img_tensor, msg_features, prompt_features)959 960 stego_img = postprocess_image(stego_tensor)961 metrics = calculate_metrics(img, stego_img)962 963 results.append({964 'filename': file.name,965 'psnr': metrics['psnr'],966 'ssim': metrics['ssim'],967 'mse': metrics['mse']968 })969 970 except Exception as e:971 st.error(f"Error processing {file.name}: {e}")972 973 progress_bar.progress((i + 1) / len(batch_files))974 975 # Display results976 if results:977 st.subheader("📊 Batch Results")978 import pandas as pd979 df = pd.DataFrame(results)980 st.dataframe(df)981 982 # Summary statistics983 st.subheader("📈 Summary Statistics")984 col1, col2, col3 = st.columns(3)985 with col1:986 st.metric("Avg PSNR", f"{df['psnr'].mean():.2f} dB")987 with col2:988 st.metric("Avg SSIM", f"{df['ssim'].mean():.4f}")989 with col3:990 st.metric("Avg MSE", f"{df['mse'].mean():.2f}")991 992 # Model debugging993 st.subheader("🔧 Model Debugging")994 995 if st.button("🔍 Analyze Model Architecture"):996 st.subheader("Encoder Architecture")997 st.text(str(encoder))998 999 st.subheader("Decoder Architecture")1000 st.text(str(decoder))1001 1002 # Parameter count1003 encoder_params = sum(p.numel() for p in encoder.parameters())1004 decoder_params = sum(p.numel() for p in decoder.parameters())1005 1006 st.metric("Encoder Parameters", f"{encoder_params:,}")1007 st.metric("Decoder Parameters", f"{decoder_params:,}")1008 st.metric("Total Parameters", f"{encoder_params + decoder_params:,}")1009 1010 # Feature visualization1011 st.subheader("🎨 Feature Visualization")1012 1013 test_message = st.text_input("Test message for feature analysis:", "Hello World!")1014 test_prompt = st.text_input("Test prompt for feature analysis:", "hide strongly")1015 1016 if st.button("📊 Analyze Features"):1017 if test_message:1018 msg_features = text_to_features(test_message)1019 prompt_features = prompt_to_features(test_prompt)1020 1021 # Plot features1022 fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 4))1023 1024 ax1.plot(msg_features.numpy().flatten())1025 ax1.set_title("Message Features")1026 ax1.set_xlabel("Feature Index")1027 ax1.set_ylabel("Value")1028 ax1.grid(True)1029 1030 ax2.plot(prompt_features.numpy().flatten())1031 ax2.set_title("Prompt Features")1032 ax2.set_xlabel("Feature Index")1033 ax2.set_ylabel("Value")1034 ax2.grid(True)1035 1036 plt.tight_layout()1037 st.pyplot(fig)1038 1039 # Feature statistics1040 col1, col2 = st.columns(2)1041 with col1:1042 st.metric("Message Feature Range", 1043 f"{msg_features.min():.3f} to {msg_features.max():.3f}")1044 st.metric("Message Feature Mean", f"{msg_features.mean():.3f}")1045 with col2:1046 st.metric("Prompt Feature Range", 1047 f"{prompt_features.min():.3f} to {prompt_features.max():.3f}")1048 st.metric("Prompt Feature Mean", f"{prompt_features.mean():.3f}")1049 1050 # Footer1051 st.markdown("---")1052 st.markdown("""1053 <div style="text-align: center; color: #666; margin-top: 2rem;">1054 <p>🎭 Prompt + GAN Steganography PoC | Built with Streamlit & PyTorch</p>1055 <p><small>⚠️ This is a proof-of-concept implementation for research purposes</small></p>1056 </div>1057 """, unsafe_allow_html=True)1058 1059if __name__ == "__main__":1060 main()