image_model.py

"""
image_model.py — A real CNN (Convolutional Neural Network) built from scratch in PyTorch.

Architecture:
  Input (3 x 64 x 64 image)
      → Conv2d(3, 32, 3)  + BatchNorm + ReLU + MaxPool    →  32 x 31 x 31
      → Conv2d(32, 64, 3) + BatchNorm + ReLU + MaxPool    →  64 x 14 x 14
      → Conv2d(64,128, 3) + BatchNorm + ReLU + MaxPool    → 128 x  6 x  6
      → Flatten → Linear(128*6*6, 512) → ReLU → Dropout
      → Linear(512, 128) → ReLU
      → Linear(128, 8 categories)

This CNN learns to LOOK at images the same way your text network learns to READ text.
"""

import torch
import torch.nn as nn
import torch.optim as optim
import threading
_CNN_LOCK = threading.Lock()
import numpy as np
import os
import json
from datetime import datetime, UTC
from pathlib import Path
from collections import defaultdict

try:
    from PIL import Image
    PIL_AVAILABLE = True
except ImportError:
    PIL_AVAILABLE = False

IMG_SIZE   = 64       # resize all images to 64x64 (small = faster on CPU)
CATEGORIES = [
    'nature', 'technology', 'science', 'people',
    'animals', 'food', 'sports', 'architecture'
]

CNN_CHECKPOINT  = 'cnn_checkpoint.pt'
CNN_STATS_FILE  = 'cnn_stats.json'


# ── IMAGE PREPROCESSING ───────────────────────────────────────────────────────
def load_image_tensor(path: str, size: int = IMG_SIZE):
    """Load image from disk → normalised float tensor (3, size, size)."""
    if not PIL_AVAILABLE:
        raise RuntimeError("Pillow not installed. Add 'Pillow' to requirements.txt")
    img = Image.open(path).convert('RGB')
    img = img.resize((size, size), Image.BILINEAR)
    arr = np.array(img, dtype=np.float32) / 255.0
    # Normalize with ImageNet mean/std (works well even for non-ImageNet data)
    mean = np.array([0.485, 0.456, 0.406])
    std  = np.array([0.229, 0.224, 0.225])
    arr  = (arr - mean) / std
    return torch.tensor(arr).permute(2, 0, 1)   # HWC → CHW


# ── CNN MODEL ─────────────────────────────────────────────────────────────────
class ImageCNN(nn.Module):
    """
    A real Convolutional Neural Network.
    Learns to detect edges → shapes → textures → objects, layer by layer.
    Each Conv2d layer is looking for patterns the previous layer found.
    """

    def __init__(self, num_classes: int = 8):
        super().__init__()
        self.num_classes = num_classes

        # Convolutional feature extractor
        self.features = nn.Sequential(
            # Block 1 — learns basic edges and colours
            nn.Conv2d(3, 32, kernel_size=3, padding=1),
            nn.BatchNorm2d(32),
            nn.ReLU(),
            nn.Conv2d(32, 32, kernel_size=3, padding=1),
            nn.BatchNorm2d(32),
            nn.ReLU(),
            nn.MaxPool2d(2, 2),        # 64→32
            nn.Dropout2d(0.1),

            # Block 2 — learns corners, curves, textures
            nn.Conv2d(32, 64, kernel_size=3, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.Conv2d(64, 64, kernel_size=3, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.MaxPool2d(2, 2),        # 32→16
            nn.Dropout2d(0.15),

            # Block 3 — learns complex shapes and object parts
            nn.Conv2d(64, 128, kernel_size=3, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(),
            nn.Conv2d(128, 128, kernel_size=3, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(),
            nn.MaxPool2d(2, 2),        # 16→8
            nn.Dropout2d(0.2),
        )

        # Classifier head
        self.classifier = nn.Sequential(
            nn.Flatten(),
            nn.Linear(128 * 8 * 8, 512),
            nn.ReLU(),
            nn.Dropout(0.4),
            nn.Linear(512, 128),
            nn.ReLU(),
            nn.Linear(128, num_classes),
        )

        # Activation tracking for visualization
        self._activations = {}
        self._register_hooks()
        self._initialize_weights()

    def _initialize_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
            elif isinstance(m, nn.BatchNorm2d):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.Linear):
                nn.init.xavier_normal_(m.weight)
                nn.init.constant_(m.bias, 0)

    def _register_hooks(self):
        def make_hook(name):
            def hook(module, inp, out):
                if isinstance(out, torch.Tensor):
                    v = out.detach().float()
                    if v.dim() > 1:
                        v = v.mean(0)
                    if v.dim() > 1:
                        v = v.mean(-1).mean(-1)  # spatial mean for conv layers
                    self._activations[name] = v[:16].tolist()
            return hook
        for i, layer in enumerate(self.features):
            layer.register_forward_hook(make_hook(f'conv_{i}'))
        for i, layer in enumerate(self.classifier):
            layer.register_forward_hook(make_hook(f'fc_{i}'))

    def forward(self, x):
        x = self.features(x)
        return self.classifier(x)

    def get_activations(self):
        return dict(self._activations)

    def get_feature_maps(self, x):
        """Return intermediate feature maps for visualization."""
        maps = {}
        for i, layer in enumerate(self.features):
            x = layer(x)
            if isinstance(layer, nn.ReLU):
                maps[f'relu_{i}'] = x.detach()
        return maps


# ── LIVING IMAGE NETWORK ──────────────────────────────────────────────────────
class LivingImageNetwork:
    """
    Wraps the CNN with training loop, data loading, and stats.
    Trains on images downloaded by ImageFetcher.
    """

    def __init__(self):
        self.model      = ImageCNN(num_classes=len(CATEGORIES))
        self.optimizer  = optim.Adam(self.model.parameters(), lr=0.001, weight_decay=1e-4)
        self.scheduler  = optim.lr_scheduler.StepLR(self.optimizer, step_size=50, gamma=0.8)
        self.criterion  = nn.CrossEntropyLoss()

        self.epoch           = 0
        self.total_images    = 0
        self.loss_history    = []
        self.acc_history     = []
        self.category_counts = defaultdict(int)

        self.stats = {
            'epoch':        0,
            'loss':         '—',
            'accuracy':     '—',
            'total_images': 0,
            'lr':           0.001,
            'last_image':   '(none yet)',
            'status':       'idle',
        }

        self._load_checkpoint()

    # ── DATA LOADING ──────────────────────────────────────────────────────────
    def _load_batch(self, image_dir: Path, batch_size: int = 16):
        """
        Load a random batch of images from image_data/ folder.
        Returns (tensor_batch, label_batch) or None if not enough images.
        """
        if not PIL_AVAILABLE:
            return None

        all_paths = []
        for cat_idx, cat in enumerate(CATEGORIES):
            cat_dir = image_dir / cat
            if cat_dir.exists():
                for p in cat_dir.iterdir():
                    if p.suffix.lower() in ('.jpg', '.jpeg', '.png', '.webp'):
                        all_paths.append((str(p), cat_idx))

        if len(all_paths) < batch_size:
            return None

        import random
        batch_paths = random.sample(all_paths, batch_size)
        tensors, labels = [], []

        for path, label in batch_paths:
            try:
                t = load_image_tensor(path)
                tensors.append(t)
                labels.append(label)
                self.category_counts[CATEGORIES[label]] += 1
            except Exception:
                continue

        if not tensors:
            return None

        return torch.stack(tensors), torch.tensor(labels, dtype=torch.long)

    # ── TRAINING ──────────────────────────────────────────────────────────────
    def train_step(self, image_dir: Path, batch_size: int = 16):
        """One training step — load images, forward pass, backprop."""
        batch = self._load_batch(image_dir, batch_size)
        if batch is None:
            return None

        x, y = batch
        x = x.float()  # ensure float32
        with _CNN_LOCK:
            self.model.train()
            self.model.zero_grad(set_to_none=True)
            logits = self.model(x)
            loss   = self.criterion(logits, y)
            loss.backward()
            for p in self.model.parameters():
                if p.grad is not None:
                    p.grad.data.clamp_(-1.0, 1.0)
            self.optimizer.step()
            loss_val = loss.detach().item()
            acc      = (logits.detach().argmax(1) == y).float().mean().item()

        self.epoch       += 1
        self.total_images += len(x)
        self.scheduler.step()

        self.loss_history.append(round(loss_val, 5))
        self.acc_history.append(round(acc, 4))
        if len(self.loss_history) > 500:
            self.loss_history = self.loss_history[-500:]
            self.acc_history  = self.acc_history[-500:]

        last_path = batch[0]  # just the paths string
        self.stats.update({
            'epoch':        self.epoch,
            'loss':         round(loss_val, 4),
            'accuracy':     round(acc * 100, 1),
            'total_images': self.total_images,
            'lr':           round(self.optimizer.param_groups[0]['lr'], 7),
        })

        if self.epoch % 20 == 0:
            self._save_checkpoint()
        self._write_stats()
        return loss_val

    def train_n_steps(self, image_dir: Path, n: int = 20):
        losses = []
        for _ in range(n):
            l = self.train_step(image_dir)
            if l is not None:
                losses.append(l)
        return {
            'steps':    len(losses),
            'avg_loss': round(sum(losses)/len(losses), 5) if losses else None,
        }

    # ── INFERENCE ─────────────────────────────────────────────────────────────
    def predict_image(self, image_path: str) -> dict:
        """Predict category of a single image."""
        if not PIL_AVAILABLE:
            return {'error': 'Pillow not installed'}
        try:
            t = load_image_tensor(image_path).unsqueeze(0)
            self.model.eval()
            with torch.no_grad():
                logits = self.model(t)
                probs  = torch.softmax(logits, dim=1)[0].tolist()
                pred   = int(logits.argmax(1).item())
            return {
                'prediction': CATEGORIES[pred],
                'confidence': round(probs[pred] * 100, 1),
                'all_probs':  {c: round(p*100, 2) for c, p in zip(CATEGORIES, probs)},
            }
        except Exception as e:
            return {'error': str(e)}

    # ── VIZ STATE ─────────────────────────────────────────────────────────────
    def get_viz_state(self) -> dict:
        self.model.eval()
        with torch.no_grad():
            dummy = torch.zeros(1, 3, IMG_SIZE, IMG_SIZE)
            self.model(dummy)
        return {
            'layer_sizes':  [3, 32, 64, 128, 512, 128, len(CATEGORIES)],
            'activations':  self.model.get_activations(),
            'loss_history': self.loss_history[-100:],
            'acc_history':  self.acc_history[-100:],
            'stats':        self.stats,
            'type':         'cnn',
        }

    # ── PERSISTENCE ───────────────────────────────────────────────────────────
    def _save_checkpoint(self):
        try:
            torch.save({
                'model':           self.model.state_dict(),
                'optimizer':       self.optimizer.state_dict(),
                'epoch':           self.epoch,
                'total_images':    self.total_images,
                'loss_history':    self.loss_history,
                'acc_history':     self.acc_history,
                'category_counts': dict(self.category_counts),
            }, CNN_CHECKPOINT)
        except Exception:
            pass

    def _load_checkpoint(self):
        if not os.path.exists(CNN_CHECKPOINT):
            return
        try:
            ck = torch.load(CNN_CHECKPOINT, map_location='cpu')
            self.model.load_state_dict(ck['model'])
            self.optimizer.load_state_dict(ck['optimizer'])
            self.epoch           = ck.get('epoch', 0)
            self.total_images    = ck.get('total_images', 0)
            self.loss_history    = ck.get('loss_history', [])
            self.acc_history     = ck.get('acc_history', [])
            self.category_counts = defaultdict(int, ck.get('category_counts', {}))
            if self.epoch > 0:
                self.stats.update({
                    'epoch':        self.epoch,
                    'total_images': self.total_images,
                })
        except Exception:
            pass

    def _write_stats(self):
        try:
            with open(CNN_STATS_FILE, 'w') as f:
                json.dump(self.stats, f, indent=2)
        except Exception:
            pass