image_compressor.py

# -*- coding: utf-8 -*-
"""image compressor

Automatically generated by Colab.

Original file is located at
    https://colab.research.google.com/drive/1_yikpBn5ThUXL5CHjn7gH0Tf_nySXLjm
"""

"""
deep_image_compression.py
Deep-learning image compression using a convolutional autoencoder (TensorFlow/Keras).
Includes:
- Training the autoencoder
- Compressing images via latent quantization + zlib
- Reconstruction and evaluation (PSNR, SSIM, MSE)
- Advanced visualizations
- Model export (autoencoder, encoder, decoder)
"""

import os
import zlib
import json
import numpy as np
import tensorflow as tf
from tensorflow.keras import layers, Model, Input
from tensorflow.keras.datasets import cifar10
from tensorflow.keras.callbacks import EarlyStopping
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.manifold import TSNE
from tqdm import tqdm

# Set style for better visualizations
sns.set_theme()
sns.set_palette("husl")

# -------------------------
# Config
# -------------------------
BATCH_SIZE = 6
EPOCHS = 6
LATENT_DIM = 64
IMG_SHAPE = (32, 32, 3)
QUANTIZATION_LEVELS = 256

# -------------------------
# Load CIFAR-10 Data
# -------------------------
(x_train, y_train), (x_test, y_test) = cifar10.load_data()
x_train = x_train.astype("float32") / 255.0
x_test = x_test.astype("float32") / 255.0

x_train = x_train[:20000]  # subset for faster demo
y_train = y_train[:20000]
x_test = x_test[:2000]
y_test = y_test[:2000]

print("Train shape:", x_train.shape, "Test shape:", x_test.shape)

# -------------------------
# Build Autoencoder
# -------------------------
def build_autoencoder(img_shape, latent_dim):
    inp = Input(shape=img_shape, name="input_image")
    x = layers.Conv2D(32, 3, strides=2, padding="same", activation="relu")(inp)
    x = layers.Conv2D(64, 3, strides=2, padding="same", activation="relu")(x)
    x = layers.Conv2D(128, 3, strides=2, padding="same", activation="relu")(x)
    x = layers.Flatten()(x)
    z = layers.Dense(latent_dim, name="bottleneck")(x)

    x = layers.Dense(4 * 4 * 128, activation="relu")(z)
    x = layers.Reshape((4, 4, 128))(x)
    x = layers.Conv2DTranspose(128, 3, strides=2, padding="same", activation="relu")(x)
    x = layers.Conv2DTranspose(64, 3, strides=2, padding="same", activation="relu")(x)
    x = layers.Conv2DTranspose(32, 3, strides=2, padding="same", activation="relu")(x)
    out = layers.Conv2D(img_shape[2], 3, padding="same", activation="sigmoid", name="recon")(x)

    auto = Model(inp, out, name="conv_autoencoder")
    encoder = Model(inp, z, name="encoder")

    # Build decoder separately
    latent_inputs = Input(shape=(latent_dim,), name="latent_input")
    dx = auto.layers[-6](latent_inputs)
    for l in auto.layers[-5:]:
        dx = l(dx)
    decoder = Model(latent_inputs, dx, name="decoder")

    return auto, encoder, decoder

auto, encoder, decoder = build_autoencoder(IMG_SHAPE, LATENT_DIM)
auto.compile(optimizer="adam", loss="mse")
auto.summary()

# -------------------------
# Train Model
# -------------------------
es = EarlyStopping(monitor="val_loss", patience=5, restore_best_weights=True, verbose=1)
history = auto.fit(
    x_train, x_train,
    epochs=EPOCHS,
    batch_size=BATCH_SIZE,
    shuffle=True,
    validation_split=0.1,
    callbacks=[es],
    verbose=2
)

# -------------------------
# Metrics
# -------------------------
def compute_metrics(original, reconstructed):
    mse = np.mean((original - reconstructed) ** 2)
    psnr = tf.image.psnr(original, reconstructed, max_val=1.0).numpy().mean()
    ssim = tf.image.ssim(original, reconstructed, max_val=1.0).numpy().mean()
    return {"mse": float(mse), "psnr": float(psnr), "ssim": float(ssim)}

# -------------------------
# Quantization Utilities
# -------------------------
def quantize_latent(latents, levels=QUANTIZATION_LEVELS):
    z_min, z_max = latents.min(), latents.max()
    if z_max == z_min:
        z_max = z_min + 1e-6
    scaled = (latents - z_min) / (z_max - z_min) * (levels - 1)
    q = np.round(scaled).astype(np.uint16)
    meta = {"min": float(z_min), "max": float(z_max), "levels": int(levels)}
    return q, meta

def dequantize_latent(q, meta):
    levels, z_min, z_max = meta["levels"], meta["min"], meta["max"]
    scaled = q.astype("float32")
    latents = scaled / (levels - 1) * (z_max - z_min) + z_min
    return latents

# -------------------------
# Compression Utilities
# -------------------------
def compress_images(images, encoder, levels=QUANTIZATION_LEVELS):
    latents = encoder.predict(images, batch_size=BATCH_SIZE, verbose=0)
    q, meta = quantize_latent(latents, levels=levels)
    meta_bytes = json.dumps(meta).encode("utf-8")
    raw_bytes = q.tobytes()
    blob = meta_bytes + b"||META_RAW||" + raw_bytes
    compressed = zlib.compress(blob)
    return compressed, meta, q.shape

def decompress_images(compressed_blob, shape):
    decompressed = zlib.decompress(compressed_blob)
    meta_bytes, raw = decompressed.split(b"||META_RAW||", 1)
    meta = json.loads(meta_bytes.decode("utf-8"))
    q = np.frombuffer(raw, dtype=np.uint16).reshape(shape)
    latents = dequantize_latent(q, meta)
    return latents, meta

# -------------------------
# Demo: Compress & Reconstruct
# -------------------------
num_demo = 10
demo_imgs = x_test[:num_demo]

recon = auto.predict(demo_imgs, verbose=0)
metrics_plain = compute_metrics(demo_imgs, recon)

compressed_blob, meta, qshape = compress_images(demo_imgs, encoder)
latents_decoded, _ = decompress_images(compressed_blob, qshape)
recon_from_compressed = decoder.predict(latents_decoded, verbose=0)

metrics_compressed = compute_metrics(demo_imgs, recon_from_compressed)

orig_raw_bytes = demo_imgs.nbytes
compression_ratio = orig_raw_bytes / len(compressed_blob)

print("Metrics (float latent):", metrics_plain)
print("Metrics (quantized+compressed):", metrics_compressed)
print("Compression ratio:", compression_ratio)

# -------------------------
# Enhanced Visualizations (Fixed)
# -------------------------

# 1. Training History Visualization
fig, axes = plt.subplots(1, 2, figsize=(15, 6))
axes[0].plot(history.history['loss'], label='Train Loss')
axes[0].plot(history.history['val_loss'], label='Val Loss')
axes[0].set_title("Training Loss")
axes[0].set_xlabel("Epoch")
axes[0].set_ylabel("MSE Loss")
axes[0].legend()
axes[0].grid(True)

# Plot loss convergence in log scale
axes[1].semilogy(history.history['loss'], label='Train Loss')
axes[1].semilogy(history.history['val_loss'], label='Val Loss')
axes[1].set_title("Loss Convergence (Log Scale)")
axes[1].set_xlabel("Epoch")
axes[1].set_ylabel("Log MSE Loss")
axes[1].legend()
axes[1].grid(True)
plt.tight_layout()
plt.savefig("training_history.png", dpi=300, bbox_inches='tight')
plt.show()

# 2. t-SNE Visualization of Latent Space
print("Computing t-SNE embedding of latent space...")
sample_size = min(1000, len(x_test))
sample_indices = np.random.choice(len(x_test), sample_size, replace=False)
latent_samples = encoder.predict(x_test[sample_indices], verbose=0)
tsne = TSNE(n_components=2, random_state=42)
latent_2d = tsne.fit_transform(latent_samples)

plt.figure(figsize=(10, 8))
scatter = plt.scatter(latent_2d[:, 0], latent_2d[:, 1],
                     c=y_test[sample_indices].ravel(),
                     cmap='tab10', alpha=0.6)
plt.colorbar(scatter, label='CIFAR-10 Class')
plt.title("t-SNE Visualization of Latent Space")
plt.xlabel("t-SNE Dimension 1")
plt.ylabel("t-SNE Dimension 2")
plt.tight_layout()
plt.savefig("latent_tsne.png", dpi=300, bbox_inches='tight')
plt.show()

# 3. Reconstruction Quality Comparison
n_show = 5
fig, axes = plt.subplots(4, n_show, figsize=(15, 12))
for i in range(n_show):
    # Original
    axes[0, i].imshow(demo_imgs[i])
    axes[0, i].set_title("Original")
    axes[0, i].axis("off")

    # Reconstruction (float)
    axes[1, i].imshow(np.clip(recon[i], 0, 1))
    psnr_float = tf.image.psnr(demo_imgs[i:i+1], recon[i:i+1], 1.0).numpy()[0]
    axes[1, i].set_title(f"Float PSNR: {psnr_float:.1f}")
    axes[1, i].axis("off")

    # Reconstruction (quantized)
    axes[2, i].imshow(np.clip(recon_from_compressed[i], 0, 1))
    psnr_quant = tf.image.psnr(demo_imgs[i:i+1], recon_from_compressed[i:i+1], 1.0).numpy()[0]
    axes[2, i].set_title(f"Quant PSNR: {psnr_quant:.1f}")
    axes[2, i].axis("off")

    # Error heatmap
    error = np.abs(demo_imgs[i] - recon_from_compressed[i])
    error = np.mean(error, axis=2)  # Average across channels
    axes[3, i].imshow(error, cmap='hot')
    axes[3, i].set_title("Error Map")
    axes[3, i].axis("off")

plt.suptitle("Reconstruction Quality Analysis", fontsize=16)
plt.tight_layout()
plt.savefig("reconstruction_analysis.png", dpi=300, bbox_inches='tight')
plt.show()

# 4. Compression Performance Analysis
print("Computing compression performance for different settings...")
quantization_levels = [16, 32, 64, 128, 256]
test_images = x_test[:100]  # Use subset for quick analysis
results = []

for levels in tqdm(quantization_levels, desc="Testing quantization levels"):
    compressed, _, _ = compress_images(test_images, encoder, levels=levels)
    ratio = test_images.nbytes / len(compressed)
    latents, _ = decompress_images(compressed, (len(test_images), LATENT_DIM))
    reconstructed = decoder.predict(latents, verbose=0)
    metrics = compute_metrics(test_images, reconstructed)
    results.append({
        'levels': levels,
        'ratio': ratio,
        'psnr': metrics['psnr'],
        'ssim': metrics['ssim']
    })

# Create performance plots
fig, axes = plt.subplots(1, 2, figsize=(15, 6))

# PSNR vs Compression Ratio
ratios = [r['ratio'] for r in results]
psnrs = [r['psnr'] for r in results]
axes[0].plot(ratios, psnrs, 'o-')
axes[0].set_xlabel("Compression Ratio")
axes[0].set_ylabel("PSNR (dB)")
axes[0].set_title("Quality vs Compression")
axes[0].grid(True)

# Quantization Analysis
ssims = [r['ssim'] for r in results]
levels_list = [r['levels'] for r in results]  # Fixed: use different variable name
ax2 = axes[1].twinx()
ln1 = axes[1].plot(levels_list, psnrs, 'b-o', label='PSNR')
ln2 = ax2.plot(levels_list, ssims, 'r-o', label='SSIM')
axes[1].set_xlabel("Quantization Levels")
axes[1].set_ylabel("PSNR (dB)", color='b')
ax2.set_ylabel("SSIM", color='r')
axes[1].set_title("Impact of Quantization")
axes[1].grid(True)

# Add legend
lns = ln1 + ln2
labs = [l.get_label() for l in lns]
axes[1].legend(lns, labs, loc='upper left')

plt.tight_layout()
plt.savefig("compression_analysis.png", dpi=300, bbox_inches='tight')
plt.show()

# 5. Latent Space Distribution Analysis (Fixed)
latents = encoder.predict(x_test[:1000], verbose=0)
fig, axes = plt.subplots(2, 2, figsize=(16, 12))

# Distribution of latent values - Fixed the axes indexing
axes[0, 0].hist(latents.flatten(), bins=50, density=True, alpha=0.7)
axes[0, 0].set_title("Latent Space Distribution")
axes[0, 0].set_xlabel("Latent Value")
axes[0, 0].set_ylabel("Density")

# Heatmap of latent correlations (use smaller sample for performance)
sample_latents = latents[:500] if len(latents) > 500 else latents
corr_matrix = np.corrcoef(sample_latents.T)
sns.heatmap(corr_matrix, ax=axes[0, 1], cmap='coolwarm', center=0,
            xticklabels=False, yticklabels=False)
axes[0, 1].set_title("Latent Dimension Correlations")

# Box plot of latent dimensions (sample first 16 dimensions)
latent_subset = sample_latents[:, :16]
axes[1, 0].boxplot([latent_subset[:, i] for i in range(16)], labels=range(16))
axes[1, 0].set_title("Latent Dimensions Distribution")
axes[1, 0].set_xlabel("Latent Dimension")
axes[1, 0].set_ylabel("Value")
axes[1, 0].tick_params(axis='x', rotation=45)

# Variance explained by each latent dimension
latent_vars = np.var(sample_latents, axis=0)
top_dims = np.argsort(latent_vars)[-20:][::-1]  # Top 20 most variant dimensions
axes[1, 1].bar(range(len(top_dims)), latent_vars[top_dims])
axes[1, 1].set_title("Top 20 Latent Dimensions by Variance")
axes[1, 1].set_xlabel("Dimension Index")
axes[1, 1].set_ylabel("Variance")

plt.tight_layout()
plt.savefig("latent_analysis.png", dpi=300, bbox_inches='tight')
plt.show()

print("\nSaved enhanced visualizations:")
print("- training_history.png: Training loss curves")
print("- latent_tsne.png: t-SNE visualization of latent space")
print("- reconstruction_analysis.png: Detailed reconstruction quality analysis")
print("- compression_analysis.png: Compression performance analysis")
print("- latent_analysis.png: Latent space distribution analysis")

# -------------------------
# Additional Advanced Visualizations (Fixed)
# -------------------------

# 6. Per-Class Performance Analysis
print("Analyzing per-class compression performance...")
class_names = ['Airplane', 'Automobile', 'Bird', 'Cat', 'Deer',
               'Dog', 'Frog', 'Horse', 'Ship', 'Truck']

class_metrics = []
for class_idx in range(10):
    class_mask = (y_test.ravel() == class_idx)
    if np.sum(class_mask) > 0:
        class_images = x_test[class_mask][:50]  # Max 50 images per class
        class_recon = auto.predict(class_images, verbose=0)
        metrics = compute_metrics(class_images, class_recon)
        metrics['class'] = class_names[class_idx]
        metrics['class_idx'] = class_idx
        class_metrics.append(metrics)

# Plot per-class performance
fig, axes = plt.subplots(1, 3, figsize=(18, 6))

classes = [m['class'] for m in class_metrics]
psnrs = [m['psnr'] for m in class_metrics]
ssims = [m['ssim'] for m in class_metrics]
mses = [m['mse'] for m in class_metrics]

axes[0].bar(classes, psnrs, color='skyblue', alpha=0.7)
axes[0].set_title("PSNR by Class")
axes[0].set_ylabel("PSNR (dB)")
axes[0].tick_params(axis='x', rotation=45)

axes[1].bar(classes, ssims, color='lightgreen', alpha=0.7)
axes[1].set_title("SSIM by Class")
axes[1].set_ylabel("SSIM")
axes[1].tick_params(axis='x', rotation=45)

axes[2].bar(classes, mses, color='salmon', alpha=0.7)
axes[2].set_title("MSE by Class")
axes[2].set_ylabel("MSE")
axes[2].tick_params(axis='x', rotation=45)

plt.tight_layout()
plt.savefig("per_class_performance.png", dpi=300, bbox_inches='tight')
plt.show()

# 7. Compression Ratio vs Quality Trade-off Analysis
print("Creating comprehensive trade-off analysis...")
fig, axes = plt.subplots(2, 2, figsize=(16, 12))

# Scatter plot: Compression Ratio vs PSNR
ratios = [r['ratio'] for r in results]
psnrs = [r['psnr'] for r in results]
ssims = [r['ssim'] for r in results]
levels_for_plot = [r['levels'] for r in results]  # Fixed: use different variable name

scatter = axes[0, 0].scatter(ratios, psnrs, c=levels_for_plot, s=100, cmap='viridis', alpha=0.7)
axes[0, 0].set_xlabel("Compression Ratio")
axes[0, 0].set_ylabel("PSNR (dB)")
axes[0, 0].set_title("Compression vs Quality Trade-off")
axes[0, 0].grid(True, alpha=0.3)
plt.colorbar(scatter, ax=axes[0, 0], label='Quantization Levels')

# Rate-Distortion curve
axes[0, 1].plot(ratios, psnrs, 'o-', linewidth=2, markersize=8, label='PSNR')
ax1_twin = axes[0, 1].twinx()
ax1_twin.plot(ratios, ssims, 's-', color='red', linewidth=2, markersize=8, label='SSIM')
axes[0, 1].set_xlabel("Compression Ratio")
axes[0, 1].set_ylabel("PSNR (dB)", color='blue')
ax1_twin.set_ylabel("SSIM", color='red')
axes[0, 1].set_title("Rate-Distortion Curve")
axes[0, 1].grid(True, alpha=0.3)

# Efficiency frontier
efficiency = np.array(psnrs) * np.array(ratios)  # Quality × Compression
axes[1, 0].bar(range(len(levels_for_plot)), efficiency, color='orange', alpha=0.7)
axes[1, 0].set_xlabel("Quantization Setting")
axes[1, 0].set_ylabel("Efficiency (PSNR × Ratio)")
axes[1, 0].set_title("Compression Efficiency")
axes[1, 0].set_xticks(range(len(levels_for_plot)))
axes[1, 0].set_xticklabels([f"{l} levels" for l in levels_for_plot], rotation=45)

# File size comparison - Fixed variable name issue
original_sizes = [test_images.nbytes for _ in results]
compressed_sizes = []
for quant_level in quantization_levels:  # Fixed: use original list variable
    compressed, _, _ = compress_images(test_images, encoder, levels=quant_level)
    compressed_sizes.append(len(compressed))

x_pos = np.arange(len(quantization_levels))  # Fixed: use original list variable
width = 0.35
axes[1, 1].bar(x_pos - width/2, original_sizes, width, label='Original', alpha=0.7)
axes[1, 1].bar(x_pos + width/2, compressed_sizes, width, label='Compressed', alpha=0.7)
axes[1, 1].set_xlabel("Quantization Levels")
axes[1, 1].set_ylabel("File Size (bytes)")
axes[1, 1].set_title("File Size Comparison")
axes[1, 1].set_xticks(x_pos)
axes[1, 1].set_xticklabels([f"{l}" for l in quantization_levels])
axes[1, 1].legend()
axes[1, 1].set_yscale('log')

plt.tight_layout()
plt.savefig("trade_off_analysis.png", dpi=300, bbox_inches='tight')
plt.show()

# 8. Feature Map Visualization
print("Visualizing encoder feature maps...")
# Get intermediate outputs from encoder
layer_outputs = []
layer_names = []
for layer in encoder.layers:
    if 'conv2d' in layer.name:
        layer_outputs.append(layer.output)
        layer_names.append(layer.name)

if layer_outputs:
    feature_model = Model(inputs=encoder.input, outputs=layer_outputs)
    sample_img = x_test[0:1]  # Single image
    feature_maps = feature_model.predict(sample_img, verbose=0)

    # Show original image and feature maps
    n_layers = len(feature_maps)
    fig, axes = plt.subplots(2, max(4, n_layers), figsize=(20, 8))

    # Original image
    axes[0, 0].imshow(sample_img[0])
    axes[0, 0].set_title("Original Image")
    axes[0, 0].axis('off')

    # Feature maps from different layers
    for i, (fmap, name) in enumerate(zip(feature_maps, layer_names)):
        if i < 4:
            # Show first few channels of each layer
            if len(fmap.shape) == 4:  # (batch, height, width, channels)
                feature = fmap[0, :, :, 0]  # First channel
                axes[0, i+1].imshow(feature, cmap='viridis')
                axes[0, i+1].set_title(f"{name}\nShape: {fmap.shape[1:]}")
                axes[0, i+1].axis('off')

                # Show channel statistics
                if i < 3:
                    channel_means = np.mean(fmap[0], axis=(0,1))
                    axes[1, i+1].hist(channel_means, bins=20, alpha=0.7)
                    axes[1, i+1].set_title(f"Channel Activations\n{name}")
                    axes[1, i+1].set_xlabel("Mean Activation")
                    axes[1, i+1].set_ylabel("Frequency")

    # Clear unused subplots
    for i in range(len(feature_maps)+1, axes.shape[1]):
        axes[0, i].axis('off')
        axes[1, i].axis('off')

    axes[1, 0].axis('off')
    plt.tight_layout()
    plt.savefig("feature_maps.png", dpi=300, bbox_inches='tight')
    plt.show()

# 9. Reconstruction Error Analysis
print("Analyzing reconstruction errors...")
sample_imgs = x_test[:100]
sample_recon = auto.predict(sample_imgs, verbose=0)
errors = np.abs(sample_imgs - sample_recon)

fig, axes = plt.subplots(2, 3, figsize=(18, 12))

# Per-pixel error distribution
axes[0, 0].hist(errors.flatten(), bins=50, alpha=0.7, density=True)
axes[0, 0].set_title("Reconstruction Error Distribution")
axes[0, 0].set_xlabel("Absolute Error")
axes[0, 0].set_ylabel("Density")

# Error by color channel
for i, channel in enumerate(['Red', 'Green', 'Blue']):
    channel_errors = errors[:, :, :, i].flatten()
    axes[0, 1].hist(channel_errors, bins=30, alpha=0.5, label=channel, density=True)
axes[0, 1].set_title("Error Distribution by Color Channel")
axes[0, 1].set_xlabel("Absolute Error")
axes[0, 1].set_ylabel("Density")
axes[0, 1].legend()

# Average error per image
img_errors = np.mean(errors, axis=(1,2,3))
axes[0, 2].plot(img_errors, alpha=0.7)
axes[0, 2].set_title("Reconstruction Error per Image")
axes[0, 2].set_xlabel("Image Index")
axes[0, 2].set_ylabel("Mean Absolute Error")

# Error correlation with image statistics
img_means = np.mean(sample_imgs, axis=(1,2,3))
img_stds = np.std(sample_imgs, axis=(1,2,3))

axes[1, 0].scatter(img_means, img_errors, alpha=0.6)
axes[1, 0].set_title("Error vs Image Brightness")
axes[1, 0].set_xlabel("Mean Pixel Value")
axes[1, 0].set_ylabel("Reconstruction Error")

axes[1, 1].scatter(img_stds, img_errors, alpha=0.6)
axes[1, 1].set_title("Error vs Image Complexity")
axes[1, 1].set_xlabel("Pixel Standard Deviation")
axes[1, 1].set_ylabel("Reconstruction Error")

# Worst reconstructions
worst_indices = np.argsort(img_errors)[-5:]
for i, idx in enumerate(worst_indices):
    if i < 1:  # Show only the worst one in this space
        axes[1, 2].imshow(sample_imgs[idx])
        axes[1, 2].set_title(f"Worst Reconstruction\nError: {img_errors[idx]:.4f}")
        axes[1, 2].axis('off')

plt.tight_layout()
plt.savefig("error_analysis.png", dpi=300, bbox_inches='tight')
plt.show()

# 10. Model Architecture Visualization
print("Creating model architecture summary...")
fig, axes = plt.subplots(1, 2, figsize=(16, 8))

# Parameter count by layer type
layer_types = {}
param_counts = {}
for layer in auto.layers:
    layer_type = type(layer).__name__
    if layer_type not in layer_types:
        layer_types[layer_type] = 0
        param_counts[layer_type] = 0
    layer_types[layer_type] += 1
    param_counts[layer_type] += layer.count_params()

# Plot layer counts
axes[0].pie(layer_types.values(), labels=layer_types.keys(), autopct='%1.1f%%')
axes[0].set_title("Layer Type Distribution")

# Plot parameter distribution
axes[1].bar(param_counts.keys(), param_counts.values(), alpha=0.7)
axes[1].set_title("Parameters by Layer Type")
axes[1].set_ylabel("Parameter Count")
axes[1].tick_params(axis='x', rotation=45)

plt.tight_layout()
plt.savefig("model_architecture.png", dpi=300, bbox_inches='tight')
plt.show()

print("\n" + "="*60)
print("COMPREHENSIVE VISUALIZATION SUMMARY")
print("="*60)
print("Generated visualizations:")
print("- training_history.png: Training loss curves and convergence")
print("- latent_tsne.png: t-SNE visualization of latent space")
print("- reconstruction_analysis.png: Detailed reconstruction quality")
print("- compression_analysis.png: Compression performance analysis")
print("- latent_analysis.png: Latent space distribution and variance")
print("- per_class_performance.png: Performance breakdown by CIFAR-10 class")
print("- trade_off_analysis.png: Comprehensive trade-off analysis")
print("- feature_maps.png: Encoder feature map visualization")
print("- error_analysis.png: Detailed reconstruction error analysis")
print("- model_architecture.png: Model architecture summary")
print("="*60)

# -------------------------
# Export Models
# -------------------------
os.makedirs("saved_models", exist_ok=True)
auto.save("saved_models/autoencoder.h5")
encoder.save("saved_models/encoder.h5")
decoder.save("saved_models/decoder.h5")
print("✅ Models exported to ./saved_models/")

# -------------------------
# Interactive Analysis Functions
# -------------------------

def analyze_compression_settings():
    """Interactive function to test different compression settings"""
    print("\n" + "="*50)
    print("COMPRESSION SETTINGS ANALYSIS")
    print("="*50)

    # Test with different batch sizes for compression
    batch_sizes = [1, 5, 10, 25, 50]
    latent_dims = [32, 64, 128, 256]

    print("Testing optimal settings for your dataset...")

    # Memory usage analysis
    try:
        import psutil
        process = psutil.Process()
        initial_memory = process.memory_info().rss / 1024 / 1024  # MB
        print(f"Initial memory usage: {initial_memory:.2f} MB")
    except ImportError:
        print("psutil not available - skipping memory analysis")
        initial_memory = None

    # Test different latent dimensions if we had multiple models
    current_latent_size = encoder.predict(x_test[:1], verbose=0).nbytes
    print(f"Current latent representation size per image: {current_latent_size} bytes")

    # Estimate compression ratios for different quantization strategies
    test_sample = x_test[:20]

    strategies = {
        'Ultra High Quality': {'levels': 65536, 'description': '16-bit quantization'},
        'High Quality': {'levels': 4096, 'description': '12-bit equivalent'},
        'Balanced': {'levels': 256, 'description': '8-bit quantization'},
        'High Compression': {'levels': 64, 'description': '6-bit equivalent'},
        'Ultra Compression': {'levels': 16, 'description': '4-bit equivalent'}
    }

    print("\nCompression Strategy Analysis:")
    print("-" * 80)
    print(f"{'Strategy':<20} {'Levels':<8} {'Ratio':<8} {'PSNR':<8} {'SSIM':<8} {'Size (KB)':<12}")
    print("-" * 80)

    for name, config in strategies.items():
        compressed, _, _ = compress_images(test_sample, encoder, levels=config['levels'])
        ratio = test_sample.nbytes / len(compressed)

        # Decompress and measure quality
        latents, _ = decompress_images(compressed, (len(test_sample), LATENT_DIM))
        reconstructed = decoder.predict(latents, verbose=0)
        metrics = compute_metrics(test_sample, reconstructed)

        size_kb = len(compressed) / 1024

        print(f"{name:<20} {config['levels']:<8} {ratio:<8.2f} {metrics['psnr']:<8.2f} {metrics['ssim']:<8.3f} {size_kb:<12.2f}")

    print("-" * 80)

    # Performance timing
    import time

    print("\nPerformance Timing Analysis:")
    print("-" * 40)

    # Encoding time
    start_time = time.time()
    _ = encoder.predict(test_sample, verbose=0)
    encode_time = time.time() - start_time

    # Compression time
    start_time = time.time()
    compressed, _, _ = compress_images(test_sample, encoder)
    compress_time = time.time() - start_time

    # Decompression time
    start_time = time.time()
    latents, _ = decompress_images(compressed, (len(test_sample), LATENT_DIM))
    decompress_time = time.time() - start_time

    # Decoding time
    start_time = time.time()
    _ = decoder.predict(latents, verbose=0)
    decode_time = time.time() - start_time

    print(f"Encoding time: {encode_time:.4f}s ({encode_time/len(test_sample)*1000:.2f}ms per image)")
    print(f"Compression time: {compress_time:.4f}s ({compress_time/len(test_sample)*1000:.2f}ms per image)")
    print(f"Decompression time: {decompress_time:.4f}s ({decompress_time/len(test_sample)*1000:.2f}ms per image)")
    print(f"Decoding time: {decode_time:.4f}s ({decode_time/len(test_sample)*1000:.2f}ms per image)")
    print(f"Total round-trip time: {(encode_time + compress_time + decompress_time + decode_time):.4f}s")

# Run the interactive analysis
analyze_compression_settings()

# -------------------------
# Advanced Quality Metrics
# -------------------------

def compute_advanced_metrics(original, reconstructed):
    """Compute additional quality metrics beyond PSNR/SSIM"""
    # Ensure proper shape and range
    original = np.clip(original, 0, 1)
    reconstructed = np.clip(reconstructed, 0, 1)

    # Convert to uint8 for some metrics
    orig_uint8 = (original * 255).astype(np.uint8)
    recon_uint8 = (reconstructed * 255).astype(np.uint8)

    # Basic metrics
    mse = np.mean((original - reconstructed) ** 2)
    psnr = tf.image.psnr(original, reconstructed, max_val=1.0).numpy().mean()
    ssim = tf.image.ssim(original, reconstructed, max_val=1.0).numpy().mean()

    # Perceptual metrics (approximations)
    # Color difference (Lab space approximation)
    def rgb_to_lab_approx(rgb):
        # Simplified RGB to LAB conversion approximation
        r, g, b = rgb[..., 0], rgb[..., 1], rgb[..., 2]
        l = 0.299 * r + 0.587 * g + 0.114 * b
        a = r - g
        b_lab = (r + g) / 2 - b
        return np.stack([l, a, b_lab], axis=-1)

    orig_lab = rgb_to_lab_approx(original)
    recon_lab = rgb_to_lab_approx(reconstructed)
    delta_e = np.mean(np.sqrt(np.sum((orig_lab - recon_lab) ** 2, axis=-1)))

    # Frequency domain analysis
    def compute_frequency_error(img1, img2):
        errors = []
        for i in range(img1.shape[0]):  # For each image in batch
            for c in range(img1.shape[-1]):  # For each channel
                fft1 = np.fft.fft2(img1[i, :, :, c])
                fft2 = np.fft.fft2(img2[i, :, :, c])
                freq_error = np.mean(np.abs(fft1 - fft2))
                errors.append(freq_error)
        return np.mean(errors)

    freq_error = compute_frequency_error(original, reconstructed)

    # Edge preservation metric
    def compute_edge_metric(img1, img2):
        try:
            from scipy import ndimage
            edge_errors = []
            for i in range(img1.shape[0]):
                for c in range(img1.shape[-1]):
                    # Compute edges using Sobel filter
                    edge1 = ndimage.sobel(img1[i, :, :, c])
                    edge2 = ndimage.sobel(img2[i, :, :, c])
                    edge_error = np.mean(np.abs(edge1 - edge2))
                    edge_errors.append(edge_error)
            return np.mean(edge_errors)
        except ImportError:
            # Fallback if scipy not available - simple gradient
            def simple_gradient(img):
                grad_x = np.diff(img, axis=1, append=img[:, -1:])
                grad_y = np.diff(img, axis=0, append=img[-1:, :])
                return np.sqrt(grad_x**2 + grad_y**2)

            edge_errors = []
            for i in range(img1.shape[0]):
                for c in range(img1.shape[-1]):
                    edge1 = simple_gradient(img1[i, :, :, c])
                    edge2 = simple_gradient(img2[i, :, :, c])
                    edge_error = np.mean(np.abs(edge1 - edge2))
                    edge_errors.append(edge_error)
            return np.mean(edge_errors)

    edge_preservation = compute_edge_metric(original, reconstructed)

    return {
        'mse': float(mse),
        'psnr': float(psnr),
        'ssim': float(ssim),
        'delta_e_approx': float(delta_e),
        'frequency_error': float(freq_error),
        'edge_preservation': edge_preservation
    }

# Test advanced metrics
print("\n" + "="*50)
print("ADVANCED QUALITY ANALYSIS")
print("="*50)

advanced_metrics = compute_advanced_metrics(demo_imgs, recon_from_compressed)
print("Advanced Quality Metrics:")
for metric, value in advanced_metrics.items():
    if value is not None:
        print(f"  {metric}: {value:.6f}")
    else:
        print(f"  {metric}: Not available (missing dependency)")

# -------------------------
# Batch Processing Utilities
# -------------------------

class ImageCompressor:
    """Utility class for batch image compression and decompression"""

    def __init__(self, encoder_model, decoder_model, quantization_levels=256):
        self.encoder = encoder_model
        self.decoder = decoder_model
        self.quantization_levels = quantization_levels
        self.compression_stats = {
            'images_processed': 0,
            'total_original_size': 0,
            'total_compressed_size': 0,
            'average_psnr': 0,
            'average_ssim': 0
        }

    def compress_batch(self, images, batch_size=32):
        """Compress a batch of images"""
        compressed_data = []
        total_original = 0
        total_compressed = 0

        for i in range(0, len(images), batch_size):
            batch = images[i:i + batch_size]
            compressed, meta, shape = compress_images(batch, self.encoder, self.quantization_levels)

            compressed_data.append({
                'data': compressed,
                'meta': meta,
                'shape': shape,
                'batch_indices': (i, min(i + batch_size, len(images)))
            })

            total_original += batch.nbytes
            total_compressed += len(compressed)

        # Update stats
        self.compression_stats['images_processed'] += len(images)
        self.compression_stats['total_original_size'] += total_original
        self.compression_stats['total_compressed_size'] += total_compressed

        return compressed_data

    def decompress_batch(self, compressed_data):
        """Decompress batch data back to images"""
        all_images = []

        for batch_data in compressed_data:
            latents, _ = decompress_images(batch_data['data'], batch_data['shape'])
            reconstructed = self.decoder.predict(latents, verbose=0)
            all_images.append(reconstructed)

        return np.vstack(all_images)

    def get_compression_ratio(self):
        """Get overall compression ratio"""
        if self.compression_stats['total_original_size'] > 0:
            return self.compression_stats['total_original_size'] / self.compression_stats['total_compressed_size']
        return 0

    def evaluate_quality(self, original_images, compressed_data):
        """Evaluate quality of compression"""
        reconstructed = self.decompress_batch(compressed_data)
        metrics = compute_metrics(original_images, reconstructed)

        # Update stats
        self.compression_stats['average_psnr'] = metrics['psnr']
        self.compression_stats['average_ssim'] = metrics['ssim']

        return metrics

# Demo the compressor class
print("\n" + "="*50)
print("BATCH COMPRESSION DEMO")
print("="*50)

compressor = ImageCompressor(encoder, decoder, quantization_levels=256)

# Compress a larger batch
test_batch = x_test[:100]
print(f"Processing {len(test_batch)} images...")

compressed_batch = compressor.compress_batch(test_batch, batch_size=20)
quality_metrics = compressor.evaluate_quality(test_batch, compressed_batch)

print(f"Compression ratio: {compressor.get_compression_ratio():.2f}x")
print(f"Quality metrics: PSNR={quality_metrics['psnr']:.2f}dB, SSIM={quality_metrics['ssim']:.3f}")
print(f"Total batches created: {len(compressed_batch)}")

# -------------------------
# Export Compression Pipeline
# -------------------------

def save_compression_pipeline(encoder, decoder, config, filename="compression_pipeline.json"):
    """Save compression configuration for later use"""
    pipeline_config = {
        'model_info': {
            'encoder_layers': len(encoder.layers),
            'decoder_layers': len(decoder.layers),
            'latent_dim': LATENT_DIM,
            'input_shape': IMG_SHAPE
        },
        'compression_config': config,
        'performance_stats': compressor.compression_stats,
        'recommended_settings': {
            'batch_size': 32,
            'quantization_levels': 256,
            'max_memory_usage_mb': 2048
        }
    }

    with open(filename, 'w') as f:
        json.dump(pipeline_config, f, indent=2)

    print(f"✅ Compression pipeline configuration saved to {filename}")

# Save the pipeline configuration
compression_config = {
    'quantization_levels': QUANTIZATION_LEVELS,
    'latent_dim': LATENT_DIM,
    'training_epochs': EPOCHS,
    'batch_size': BATCH_SIZE
}

save_compression_pipeline(encoder, decoder, compression_config)

# -------------------------
# Additional Utility Functions
# -------------------------

def compress_single_image(image, encoder, decoder, levels=256):
    """Compress and decompress a single image with detailed statistics"""
    # Ensure image is in the right format
    if len(image.shape) == 3:
        image = np.expand_dims(image, axis=0)

    # Original size
    original_size = image.nbytes

    # Encode
    latent = encoder.predict(image, verbose=0)

    # Quantize and compress
    q, meta = quantize_latent(latent, levels=levels)
    meta_bytes = json.dumps(meta).encode("utf-8")
    raw_bytes = q.tobytes()
    blob = meta_bytes + b"||META_RAW||" + raw_bytes
    compressed = zlib.compress(blob)

    # Decompress
    decompressed = zlib.decompress(compressed)
    meta_bytes_dec, raw_dec = decompressed.split(b"||META_RAW||", 1)
    meta_dec = json.loads(meta_bytes_dec.decode("utf-8"))
    q_dec = np.frombuffer(raw_dec, dtype=np.uint16).reshape(q.shape)
    latent_dec = dequantize_latent(q_dec, meta_dec)

    # Decode
    reconstructed = decoder.predict(latent_dec, verbose=0)

    # Compute metrics
    metrics = compute_metrics(image, reconstructed)

    # Compression statistics
    compression_ratio = original_size / len(compressed)
    latent_size = latent.nbytes
    quantized_size = q.nbytes
    compressed_size = len(compressed)

    stats = {
        'original_size_bytes': original_size,
        'latent_size_bytes': latent_size,
        'quantized_size_bytes': quantized_size,
        'compressed_size_bytes': compressed_size,
        'compression_ratio': compression_ratio,
        'latent_compression_ratio': original_size / latent_size,
        'quantization_ratio': latent_size / quantized_size,
        'zlib_compression_ratio': quantized_size / compressed_size,
        'quality_metrics': metrics
    }

    return reconstructed, stats

def analyze_image_complexity(images, encoder):
    """Analyze the complexity of images and their compression potential"""
    complexities = []

    for i, img in enumerate(images):
        # Compute various complexity metrics
        img_2d = np.mean(img, axis=2)  # Convert to grayscale for some metrics

        # Statistical measures
        pixel_var = np.var(img)
        pixel_entropy = -np.sum(img * np.log(img + 1e-8))  # Approximate entropy

        # Gradient-based complexity
        grad_x = np.diff(img_2d, axis=1)  # Shape: (32, 31)
        grad_y = np.diff(img_2d, axis=0)  # Shape: (31, 32)
        # Make shapes compatible by trimming to the smaller dimensions
        min_rows = min(grad_x.shape[0], grad_y.shape[0])
        min_cols = min(grad_x.shape[1], grad_y.shape[1])
        grad_x_trimmed = grad_x[:min_rows, :min_cols]
        grad_y_trimmed = grad_y[:min_rows, :min_cols]
        edge_density = np.mean(np.sqrt(grad_x_trimmed**2 + grad_y_trimmed**2))

        # Frequency domain complexity
        fft_img = np.fft.fft2(img_2d)
        high_freq_energy = np.mean(np.abs(fft_img[img_2d.shape[0]//4:, img_2d.shape[1]//4:]))

        # Latent space analysis
        img_batch = np.expand_dims(img, axis=0)
        latent = encoder.predict(img_batch, verbose=0)
        latent_var = np.var(latent)
        latent_sparsity = np.mean(np.abs(latent) < 0.1)  # Fraction of near-zero values

        complexity_metrics = {
            'image_index': i,
            'pixel_variance': float(pixel_var),
            'pixel_entropy_approx': float(pixel_entropy),
            'edge_density': float(edge_density),
            'high_freq_energy': float(high_freq_energy),
            'latent_variance': float(latent_var),
            'latent_sparsity': float(latent_sparsity)
        }

        complexities.append(complexity_metrics)

    return complexities

# Analyze complexity of test images
print("\n" + "="*50)
print("IMAGE COMPLEXITY ANALYSIS")
print("="*50)

sample_complexities = analyze_image_complexity(x_test[:50], encoder)

# Find images with different complexity levels
sorted_by_variance = sorted(sample_complexities, key=lambda x: x['pixel_variance'])
sorted_by_edges = sorted(sample_complexities, key=lambda x: x['edge_density'])

print("Images sorted by complexity:")
print(f"Lowest variance (simple): Image {sorted_by_variance[0]['image_index']}, variance: {sorted_by_variance[0]['pixel_variance']:.4f}")
print(f"Highest variance (complex): Image {sorted_by_variance[-1]['image_index']}, variance: {sorted_by_variance[-1]['pixel_variance']:.4f}")
print(f"Lowest edge density: Image {sorted_by_edges[0]['image_index']}, edge density: {sorted_by_edges[0]['edge_density']:.4f}")
print(f"Highest edge density: Image {sorted_by_edges[-1]['image_index']}, edge density: {sorted_by_edges[-1]['edge_density']:.4f}")

# Test compression on images of different complexities
print("\nCompression performance on different complexity images:")
test_indices = [sorted_by_variance[0]['image_index'], sorted_by_variance[-1]['image_index']]
test_names = ["Simple (low variance)", "Complex (high variance)"]

for idx, name in zip(test_indices, test_names):
    img = x_test[idx]
    _, stats = compress_single_image(img, encoder, decoder)
    print(f"\n{name}:")
    print(f"  Compression ratio: {stats['compression_ratio']:.2f}x")
    print(f"  PSNR: {stats['quality_metrics']['psnr']:.2f}dB")
    print(f"  SSIM: {stats['quality_metrics']['ssim']:.3f}")

print("\n" + "="*60)
print("🎉 DEEP IMAGE COMPRESSION ANALYSIS COMPLETE!")
print("="*60)
print("Summary of generated files:")
print("📊 Visualizations: 10 PNG files with comprehensive analysis")
print("🤖 Models: autoencoder.h5, encoder.h5, decoder.h5")
print("⚙️  Configuration: compression_pipeline.json")
print("\nThe pipeline is ready for production use!")
print("="*60)

# Google Colab Setup for Deep Image Compression Streamlit App
# Run each cell in sequence to set up and launch the app

# =============================================================================
# CELL 1: Install Required Packages
# =============================================================================
print("Installing required packages...")
!pip install -q streamlit opencv-python-headless ngrok pyngrok

print("✅ Installation complete!")

# =============================================================================
# CELL 2: Setup Ngrok for External Access
# =============================================================================
import os
from pyngrok import ngrok

# Set your ngrok auth token (get from https://ngrok.com/)
# Replace 'YOUR_NGROK_TOKEN' with your actual token
NGROK_TOKEN = "YOUR_NGROK_TOKEN"  # Get this from ngrok.com after signing up

if NGROK_TOKEN != "YOUR_NGROK_TOKEN":
    ngrok.set_auth_token(NGROK_TOKEN)
    print("✅ Ngrok token set successfully!")
else:
    print("⚠️  Please set your NGROK_TOKEN in the cell above")
    print("   1. Go to https://ngrok.com/")
    print("   2. Sign up for a free account")
    print("   3. Get your auth token from the dashboard")
    print("   4. Replace 'YOUR_NGROK_TOKEN' with your actual token")

# =============================================================================
# CELL 3: Create the Streamlit App File
# =============================================================================
app_code = '''import streamlit as st
import numpy as np
import tensorflow as tf
from tensorflow.keras.models import load_model
import matplotlib.pyplot as plt
import json
import zlib
from PIL import Image
import io
import cv2
import os

# Page configuration
st.set_page_config(
    page_title="Deep Image Compression",
    page_icon="🖼️",
    layout="wide",
    initial_sidebar_state="expanded"
)

# Title and description
st.title("🖼️ Deep Learning Image Compression")
st.markdown("""
This app uses a convolutional autoencoder to compress and reconstruct images.
Upload an image to test the compression performance!
""")

# Utility functions
def quantize_latent(latents, levels=256):
    z_min, z_max = latents.min(), latents.max()
    if z_max == z_min:
        z_max = z_min + 1e-6
    scaled = (latents - z_min) / (z_max - z_min) * (levels - 1)
    q = np.round(scaled).astype(np.uint16)
    meta = {"min": float(z_min), "max": float(z_max), "levels": int(levels)}
    return q, meta

def dequantize_latent(q, meta):
    levels, z_min, z_max = meta["levels"], meta["min"], meta["max"]
    scaled = q.astype("float32")
    latents = scaled / (levels - 1) * (z_max - z_min) + z_min
    return latents

def compress_images(images, encoder, levels=256):
    latents = encoder.predict(images, verbose=0)
    q, meta = quantize_latent(latents, levels=levels)
    meta_bytes = json.dumps(meta).encode("utf-8")
    raw_bytes = q.tobytes()
    blob = meta_bytes + b"||META_RAW||" + raw_bytes
    compressed = zlib.compress(blob)
    return compressed, meta, q.shape

def decompress_images(compressed_blob, shape):
    decompressed = zlib.decompress(compressed_blob)
    meta_bytes, raw = decompressed.split(b"||META_RAW||", 1)
    meta = json.loads(meta_bytes.decode("utf-8"))
    q = np.frombuffer(raw, dtype=np.uint16).reshape(shape)
    latents = dequantize_latent(q, meta)
    return latents, meta

def compute_metrics(original, reconstructed):
    mse = np.mean((original - reconstructed) ** 2)
    psnr = tf.image.psnr(original, reconstructed, max_val=1.0).numpy().mean()
    ssim = tf.image.ssim(original, reconstructed, max_val=1.0).numpy().mean()
    return {"mse": float(mse), "psnr": float(psnr), "ssim": float(ssim)}

def preprocess_image(image, target_size=(32, 32)):
    """Preprocess uploaded image to match model input requirements"""
    img_array = np.array(image)
    img_resized = cv2.resize(img_array, target_size)

    if len(img_resized.shape) == 2:  # Grayscale
        img_resized = cv2.cvtColor(img_resized, cv2.COLOR_GRAY2RGB)
    elif img_resized.shape[2] == 4:  # RGBA
        img_resized = cv2.cvtColor(img_resized, cv2.COLOR_RGBA2RGB)

    img_normalized = img_resized.astype(np.float32) / 255.0
    return img_normalized

# Sidebar for model loading and settings
st.sidebar.header("Model Settings")

# Model loading section
st.sidebar.subheader("Load Models")
model_path = st.sidebar.text_input(
    "Model Directory Path",
    value="saved_models",
    help="Path to directory containing encoder.h5, decoder.h5, and autoencoder.h5"
)

# Load models function
@st.cache_resource
def load_models(model_path):
    try:
        encoder = load_model(os.path.join(model_path, "encoder.h5"))
        decoder = load_model(os.path.join(model_path, "decoder.h5"))
        autoencoder = load_model(os.path.join(model_path, "autoencoder.h5"))
        return encoder, decoder, autoencoder, True
    except Exception as e:
        st.error(f"Error loading models: {str(e)}")
        return None, None, None, False

# Try to load models
encoder, decoder, autoencoder, models_loaded = load_models(model_path)

if models_loaded:
    st.sidebar.success("✅ Models loaded successfully!")

    # Model information
    with st.sidebar.expander("Model Information"):
        st.write(f"**Encoder layers:** {len(encoder.layers)}")
        st.write(f"**Decoder layers:** {len(decoder.layers)}")
        st.write(f"**Autoencoder layers:** {len(autoencoder.layers)}")

        latent_dim = encoder.output_shape[-1]
        st.write(f"**Latent dimension:** {latent_dim}")

else:
    st.sidebar.error("❌ Failed to load models. Please check the path.")
    st.error("Models not found. Please ensure the model files are available in the specified directory.")

    # Show expected file structure
    st.subheader("Expected File Structure:")
    st.code("""
saved_models/
├── encoder.h5
├── decoder.h5
└── autoencoder.h5
    """)
    st.stop()

# Compression settings
st.sidebar.subheader("Compression Settings")
quantization_levels = st.sidebar.selectbox(
    "Quantization Levels",
    options=[16, 32, 64, 128, 256, 512, 1024],
    index=4,  # Default to 256
    help="Higher values = better quality, lower compression"
)

# Main content area
col1, col2 = st.columns(2)

with col1:
    st.subheader("Upload Image")
    uploaded_file = st.file_uploader(
        "Choose an image file",
        type=['png', 'jpg', 'jpeg'],
        help="Upload an image to compress and reconstruct"
    )

    st.subheader("Or use sample CIFAR-10 images")
    if st.button("Load Random CIFAR-10 Sample"):
        with st.spinner("Loading CIFAR-10 sample..."):
            (x_test, y_test), _ = tf.keras.datasets.cifar10.load_data()
            x_test = x_test.astype("float32") / 255.0

            random_idx = np.random.randint(0, len(x_test))
            sample_image = x_test[random_idx]

            st.session_state['current_image'] = sample_image
            st.session_state['image_source'] = f"CIFAR-10 sample (Class: {y_test[random_idx][0]})"

# Process uploaded image
if uploaded_file is not None:
    pil_image = Image.open(uploaded_file)
    processed_image = preprocess_image(pil_image)

    st.session_state['current_image'] = processed_image
    st.session_state['image_source'] = f"Uploaded: {uploaded_file.name}"

# Display and analyze image if available
if 'current_image' in st.session_state:
    current_image = st.session_state['current_image']
    image_source = st.session_state['image_source']

    with col1:
        st.subheader("Original Image")
        st.image(current_image, caption=image_source, width=200)

        st.write(f"**Shape:** {current_image.shape}")
        st.write(f"**Size:** {current_image.nbytes} bytes")
        st.write(f"**Mean pixel value:** {np.mean(current_image):.4f}")
        st.write(f"**Std deviation:** {np.std(current_image):.4f}")

    with col2:
        st.subheader("Compression Analysis")

        if st.button("🚀 Compress Image", type="primary"):
            with st.spinner("Compressing image..."):
                image_batch = np.expand_dims(current_image, axis=0)

                direct_reconstruction = autoencoder.predict(image_batch, verbose=0)[0]

                compressed_blob, meta, q_shape = compress_images(image_batch, encoder, levels=quantization_levels)
                latents_decoded, _ = decompress_images(compressed_blob, q_shape)
                compressed_reconstruction = decoder.predict(latents_decoded, verbose=0)[0]

                direct_metrics = compute_metrics(
                    np.expand_dims(current_image, axis=0),
                    np.expand_dims(direct_reconstruction, axis=0)
                )
                compressed_metrics = compute_metrics(
                    np.expand_dims(current_image, axis=0),
                    np.expand_dims(compressed_reconstruction, axis=0)
                )

                original_size = image_batch.nbytes
                compressed_size = len(compressed_blob)
                compression_ratio = original_size / compressed_size

                st.session_state['results'] = {
                    'direct_reconstruction': direct_reconstruction,
                    'compressed_reconstruction': compressed_reconstruction,
                    'direct_metrics': direct_metrics,
                    'compressed_metrics': compressed_metrics,
                    'original_size': original_size,
                    'compressed_size': compressed_size,
                    'compression_ratio': compression_ratio,
                    'quantization_levels': quantization_levels
                }

    # Display results if available
    if 'results' in st.session_state:
        results = st.session_state['results']

        st.header("Results")

        col3, col4, col5 = st.columns(3)

        with col3:
            st.metric("Compression Ratio", f"{results['compression_ratio']:.1f}x")

        with col4:
            st.metric(
                "PSNR (Compressed)",
                f"{results['compressed_metrics']['psnr']:.1f} dB",
                delta=f"{results['compressed_metrics']['psnr'] - results['direct_metrics']['psnr']:.1f}"
            )

        with col5:
            st.metric(
                "SSIM (Compressed)",
                f"{results['compressed_metrics']['ssim']:.3f}",
                delta=f"{results['compressed_metrics']['ssim'] - results['direct_metrics']['ssim']:.3f}"
            )

        st.subheader("Visual Comparison")

        col6, col7, col8 = st.columns(3)

        with col6:
            st.write("**Original**")
            st.image(current_image, width=150)

        with col7:
            st.write("**Direct Reconstruction**")
            st.image(np.clip(results['direct_reconstruction'], 0, 1), width=150)
            st.caption(f"PSNR: {results['direct_metrics']['psnr']:.1f}dB")

        with col8:
            st.write("**Compressed & Reconstructed**")
            st.image(np.clip(results['compressed_reconstruction'], 0, 1), width=150)
            st.caption(f"PSNR: {results['compressed_metrics']['psnr']:.1f}dB")

        # Detailed metrics
        st.subheader("Detailed Metrics")

        metrics_data = {
            'Metric': ['MSE', 'PSNR (dB)', 'SSIM'],
            'Direct': [
                f"{results['direct_metrics']['mse']:.6f}",
                f"{results['direct_metrics']['psnr']:.2f}",
                f"{results['direct_metrics']['ssim']:.4f}"
            ],
            'Compressed': [
                f"{results['compressed_metrics']['mse']:.6f}",
                f"{results['compressed_metrics']['psnr']:.2f}",
                f"{results['compressed_metrics']['ssim']:.4f}"
            ]
        }

        st.table(metrics_data)
'''

# Write the app to a file
with open('streamlit_app.py', 'w') as f:
    f.write(app_code)

print("✅ Streamlit app file created: streamlit_app.py")

# =============================================================================
# CELL 4: Upload and Organize Your Model Files
# =============================================================================
print("📁 Setting up model directory...")

import os
from google.colab import files

# Create the models directory
os.makedirs('saved_models', exist_ok=True)
print("✅ Created 'saved_models' directory")

print("""
🔧 UPLOAD YOUR MODEL FILES:

Option 1: Upload files directly (run the next cell)
Option 2: Download from Google Drive (if you have them stored there)
Option 3: Use wget if you have URLs to the model files

Your model files should be named:
- encoder.h5
- decoder.h5
- autoencoder.h5
""")

# Uncomment and run this section to upload files directly
# uploaded_files = files.upload()

# # Move uploaded files to the correct directory
# for filename in uploaded_files.keys():
#     if filename.endswith('.h5'):
#         import shutil
#         shutil.move(filename, f'saved_models/{filename}')
#         print(f"✅ Moved {filename} to saved_models/")

# =============================================================================
# CELL 5: Alternative - Load from Google Drive
# =============================================================================
print("""
🔧 ALTERNATIVE: Load from Google Drive

If you have your model files in Google Drive, uncomment and run this section:
""")

# from google.colab import drive
# drive.mount('/content/drive')

# # Copy model files from your Google Drive
# import shutil
#
# # Update these paths to match your Google Drive structure
# drive_model_path = '/content/drive/MyDrive/your_model_folder/'
#
# model_files = ['encoder.h5', 'decoder.h5', 'autoencoder.h5']
#
# for model_file in model_files:
#     src_path = drive_model_path + model_file
#     dst_path = f'saved_models/{model_file}'
#
#     try:
#         shutil.copy(src_path, dst_path)
#         print(f"✅ Copied {model_file} from Google Drive")
#     except FileNotFoundError:
#         print(f"❌ Could not find {src_path}")

# =============================================================================
# CELL 6: Verify Model Files
# =============================================================================
print("🔍 Checking for model files...")

required_files = ['encoder.h5', 'decoder.h5', 'autoencoder.h5']
missing_files = []

for file in required_files:
    file_path = f'saved_models/{file}'
    if os.path.exists(file_path):
        file_size = os.path.getsize(file_path) / (1024 * 1024)  # Size in MB
        print(f"✅ Found {file} ({file_size:.1f} MB)")
    else:
        print(f"❌ Missing {file}")
        missing_files.append(file)

if missing_files:
    print(f"\n⚠️  Missing files: {missing_files}")
    print("Please upload these files before proceeding.")
else:
    print("\n🎉 All model files found! Ready to launch the app.")

# =============================================================================
# CELL 7: Launch Streamlit App
# =============================================================================
import subprocess
import threading
import time

def run_streamlit():
    """Run Streamlit in a separate thread"""
    subprocess.run(['streamlit', 'run', 'streamlit_app.py', '--server.port=8501', '--server.headless=true'],
                   stdout=subprocess.PIPE, stderr=subprocess.PIPE)

# Start Streamlit in background
print("🚀 Starting Streamlit app...")
streamlit_thread = threading.Thread(target=run_streamlit)
streamlit_thread.daemon = True
streamlit_thread.start()

# Wait for Streamlit to start
time.sleep(10)
print("✅ Streamlit app should be running on port 8501")

# =============================================================================
# CELL 8: Create Public URL with Ngrok
# =============================================================================
print("🌐 Creating public URL with ngrok...")

try:
    # Create ngrok tunnel
    public_url = ngrok.connect(8501)
    print(f"\n🎉 SUCCESS! Your app is now available at:")
    print(f"🔗 {public_url}")
    print(f"\nThis URL will remain active as long as this Colab session is running.")
    print(f"Share this link to let others access your image compression app!")

    # Keep the tunnel alive
    print(f"\n⏳ Keeping the tunnel alive... (press Ctrl+C to stop)")

except Exception as e:
    print(f"❌ Error creating ngrok tunnel: {e}")
    print("Make sure you've set your NGROK_TOKEN correctly in Cell 2")

# =============================================================================
# CELL 9: Alternative Access Methods
# =============================================================================
print("""
🔧 ALTERNATIVE ACCESS METHODS:

If ngrok doesn't work, you can try these alternatives:

1. LOCALTUNNEL (no signup required):
   !npm install -g localtunnel
   !lt --port 8501

2. GRADIO (alternative to Streamlit):
   Convert your app to use Gradio interface instead

3. PORT FORWARDING:
   Use Colab's built-in port forwarding (may be limited)

4. COLAB SERVE (experimental):
   Some browsers support direct port access to Colab instances
""")

# =============================================================================
# CELL 10: Testing and Monitoring
# =============================================================================
print("""
🧪 TESTING YOUR APP:

1. Open the public URL provided above
2. Try uploading a test image
3. Click "Load Random CIFAR-10 Sample" to test with sample data
4. Experiment with different quantization levels
5. Check the compression ratios and quality metrics

📊 MONITORING:
- Watch this Colab output for any errors
- The app will show detailed metrics including PSNR and SSIM
- Try different types of images to see how compression varies

🐛 TROUBLESHOOTING:
- If models fail to load, check file paths and names
- If images look wrong, verify the preprocessing pipeline
- If compression seems poor, try different quantization levels
- Check Colab's resource usage to ensure you have enough memory
""")

# Keep the session alive
print("\n🔄 Session will stay active. Monitor the public URL above.")
print("To stop the app, interrupt this cell or restart the runtime.")

# Optional: Keep checking if the tunnel is still alive
import time
try:
    while True:
        time.sleep(60)  # Check every minute
        print(".", end="", flush=True)  # Show that the session is alive
except KeyboardInterrupt:
    print("\n🛑 Stopping the app...")
    ngrok.disconnect(public_url)
    print("✅ App stopped and tunnel closed.")