""" 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