Delete image_model.py

image_model.py

@@ -1,363 +0,0 @@
-"""
-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 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