Upload spa_ensemble.py

new_approach/spa_ensemble.py

@@ -0,0 +1,351 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+import cv2
+from PIL import Image
+import os
+from pathlib import Path
+from scipy import stats
+from scipy.fftpack import dct
+from sklearn.preprocessing import StandardScaler
+import torchvision.transforms as transforms
+import open_clip
+import joblib
+from huggingface_hub import hf_hub_download
+
+# --- CONFIGURATION ---
+CONFIDENCE_THRESHOLD = 0.99
+# The list directory remains in the root of the Space
+LIST_DIR = Path("list") 
+
+# ==============================================================================
+# 1. FEATURE EXTRACTOR
+# ==============================================================================
+class FeatureExtractor:
+    @staticmethod
+    def extract_color_features(img):
+        img_np = np.array(img); features = {}
+        for i, channel in enumerate(['R', 'G', 'B']):
+            ch = img_np[:, :, i].flatten()
+            if len(ch) > 0:
+                features.update({f'color_{channel}_mean': float(np.mean(ch)), f'color_{channel}_std': float(np.std(ch)), f'color_{channel}_skew': float(stats.skew(ch)), f'color_{channel}_min': float(np.min(ch)), f'color_{channel}_max': float(np.max(ch))})
+            else:
+                 features.update({f'color_{channel}_mean': 0.0, f'color_{channel}_std': 0.0, f'color_{channel}_skew': 0.0, f'color_{channel}_min': 0.0, f'color_{channel}_max': 0.0})
+            
+            # --- FIX: Removed Histogram extraction (9 features) to match the 40 features expected by your .pth files ---
+            # hist, _ = np.histogram(ch, bins=3, range=(0, 256)); hist = hist / (hist.sum() + 1e-8); 
+            # for j, v in enumerate(hist): features[f'color_{channel}_hist_bin{j}'] = float(v)
+            
+        try:
+            hsv = cv2.cvtColor(img_np, cv2.COLOR_RGB2HSV)
+            features.update({'color_hue_mean': float(np.mean(hsv[:, :, 0])), 'color_saturation_mean': float(np.mean(hsv[:, :, 1])), 'color_value_mean': float(np.mean(hsv[:, :, 2]))})
+        except: 
+            features.update({'color_hue_mean': 0.0, 'color_saturation_mean': 0.0, 'color_value_mean': 0.0})
+        return features
+
+    @staticmethod
+    def extract_texture_features(img):
+        img_np = np.array(img); gray = cv2.cvtColor(img_np, cv2.COLOR_RGB2GRAY); features = {}
+        # Optimization: Canny/Sobel can be slow on huge images.
+        # We assume image is resized in extract_all_features
+        edges = cv2.Canny(gray, 50, 150)
+        gx, gy = cv2.Sobel(gray, cv2.CV_64F, 1, 0, ksize=3), cv2.Sobel(gray, cv2.CV_64F, 0, 1, ksize=3)
+        features.update({
+            'texture_edge_density': float(np.sum(edges > 0) / edges.size) if edges.size > 0 else 0.0, 
+            'texture_gradient_mean': float(np.mean(np.sqrt(gx**2 + gy**2))), 
+            'texture_gradient_std': float(np.std(np.sqrt(gx**2 + gy**2))), 
+            'texture_laplacian_var': float(np.var(cv2.Laplacian(gray, cv2.CV_64F)))
+        })
+        return features
+
+    @staticmethod
+    def extract_shape_features(img):
+        w, h = img.size; features = {}; features.update({'shape_height': h, 'shape_width': w, 'shape_aspect_ratio': w / h if h > 0 else 0.0, 'shape_area': w * h}); return features
+
+    @staticmethod
+    def extract_brightness_features(img):
+        img_np = np.array(img); gray = cv2.cvtColor(img_np, cv2.COLOR_RGB2GRAY); features = {}; features.update({'brightness_mean': float(np.mean(gray)), 'brightness_std': float(np.std(gray))}); return features
+
+    @staticmethod
+    def extract_frequency_features(img):
+        img_np = np.array(img)
+        gray = cv2.cvtColor(img_np, cv2.COLOR_RGB2GRAY)
+        gray_small = cv2.resize(gray, (64, 64))
+        dct_coeffs = dct(dct(gray_small.T, norm='ortho').T, norm='ortho')
+        features = {}
+        # FIX: Loop must finish before returning!
+        for i, v in enumerate(dct_coeffs.flatten()[:10]): 
+            features[f'freq_dct_{i}'] = float(v)
+        return features
+
+    @staticmethod
+    def extract_statistical_features(img):
+        img_np = np.array(img); gray = cv2.cvtColor(img_np, cv2.COLOR_RGB2GRAY); hist, _ = np.histogram(gray.flatten(), bins=256, range=(0, 256)); hist = hist / (hist.sum() + 1e-8)
+        hist_nonzero = hist[hist > 0]; entropy = -np.sum(hist_nonzero * np.log2(hist_nonzero)) if hist_nonzero.size > 0 else 0.0; features = {}; features.update({'stat_entropy': entropy, 'stat_uniformity': float(np.sum(hist**2))}); return features
+
+    @staticmethod
+    def extract_all_features(img):
+        img = img.convert('RGB')
+        # OPTIMIZATION: Resize for Handcrafted Features to speed up Canny/Sobel
+        max_size = 1024
+        if max(img.size) > max_size:
+             img.thumbnail((max_size, max_size), Image.Resampling.LANCZOS)
+             
+        features = {}
+        features.update(FeatureExtractor.extract_color_features(img))
+        features.update(FeatureExtractor.extract_texture_features(img))
+        features.update(FeatureExtractor.extract_shape_features(img))
+        features.update(FeatureExtractor.extract_brightness_features(img))
+        features.update(FeatureExtractor.extract_frequency_features(img))
+        features.update(FeatureExtractor.extract_statistical_features(img))
+        return features
+
+# ==============================================================================
+# 2. MODEL ARCHITECTURE
+# ==============================================================================
+class BioCLIP2ZeroShot:
+    def __init__(self, device, class_to_idx, id_to_name):
+        self.device = device; self.num_classes = len(class_to_idx); self.idx_to_class = {v: k for k, v in class_to_idx.items()}; self.id_to_name = id_to_name
+        print("Loading BioCLIP-2 model...")
+        try:
+            self.model, _, self.preprocess = open_clip.create_model_and_transforms('hf-hub:imageomics/bioclip-2')
+            self.tokenizer = open_clip.get_tokenizer('hf-hub:imageomics/bioclip-2')
+        except:
+            print("Warning: BioCLIP-2 load failed, trying base BioCLIP...")
+            self.model, _, self.preprocess = open_clip.create_model_and_transforms('hf-hub:imageomics/bioclip')
+            self.tokenizer = open_clip.get_tokenizer('hf-hub:imageomics/bioclip')
+        self.model.to(self.device).eval()
+        self.text_features_prototypes = self._precompute_text_features()
+
+    def _precompute_text_features(self):
+        templates = [ "a photo of {}", "a herbarium specimen of {}", "a botanical photograph of {}", "{} plant species", "leaves and flowers of {}" ]
+        class_ids = [self.idx_to_class[i] for i in range(self.num_classes)]
+        class_names = [self.id_to_name.get(str(cid), str(cid)) for cid in class_ids]
+        all_emb = []; bs = 64
+        text_inputs = [t.format(name) for name in class_names for t in templates]
+        with torch.no_grad():
+            for i in range(0, len(text_inputs), bs):
+                tokens = self.tokenizer(text_inputs[i:i+bs]).to(self.device)
+                emb = self.model.encode_text(tokens)
+                all_emb.append(emb)
+        all_text_embs = torch.cat(all_emb, dim=0).cpu().numpy()
+        prototypes = np.zeros((self.num_classes, all_text_embs.shape[1]), dtype=np.float32)
+        for idx in range(self.num_classes):
+            start = idx * len(templates)
+            avg = np.mean(all_text_embs[start:start + len(templates)], axis=0)
+            norm = np.linalg.norm(avg)
+            prototypes[idx] = avg / norm if norm > 0 else avg
+        return torch.from_numpy(prototypes).to(self.device)
+
+    def predict_zero_shot_logits(self, img):
+        processed = self.preprocess(img).unsqueeze(0).to(self.device)
+        with torch.no_grad():
+            image_features = self.model.encode_image(processed)
+            image_features = image_features / image_features.norm(dim=-1, keepdim=True)
+        prototypes = self.text_features_prototypes
+        try: logit_scale = self.model.logit_scale.exp()
+        except: logit_scale = torch.tensor(100.0).to(self.device)
+        # --- FIX: Added .detach() before .numpy() ---
+        logits = (logit_scale * image_features @ prototypes.T).detach().cpu().numpy().squeeze()
+        return logits
+
+class EnsembleClassifier(nn.Module):
+    def __init__(self, num_handcrafted_features=40, dinov2_dim=1024, bioclip2_dim=100,
+                 num_classes=100, hidden_dim=512, dropout_rate=0.3, prototype_dim=768):
+        super().__init__()
+        self.dinov2_proj = nn.Sequential(nn.Linear(dinov2_dim, hidden_dim), nn.ReLU(), nn.Dropout(dropout_rate))
+        
+        # --- FIX: Removed 3rd layer to match training checkpoint (Size mismatch error) ---
+        self.handcraft_branch = nn.Sequential(
+            nn.Linear(num_handcrafted_features, 128), nn.BatchNorm1d(128), nn.ReLU(), nn.Dropout(dropout_rate),
+            nn.Linear(128, hidden_dim // 2), nn.BatchNorm1d(hidden_dim // 2), nn.ReLU(), nn.Dropout(dropout_rate)
+        )
+        
+        self.bioclip2_branch = nn.Sequential(
+            nn.Linear(bioclip2_dim, hidden_dim // 4), nn.BatchNorm1d(hidden_dim // 4), nn.ReLU(), nn.Dropout(dropout_rate * 0.5))
+        fusion_input_dim = hidden_dim + hidden_dim // 2 + hidden_dim // 4
+        self.fusion = nn.Sequential(
+            nn.Linear(fusion_input_dim, hidden_dim), nn.BatchNorm1d(hidden_dim), nn.ReLU(), nn.Dropout(dropout_rate))
+        self.classifier = nn.Linear(hidden_dim, num_classes)
+        self.prototype_proj = nn.Linear(hidden_dim, prototype_dim) 
+
+    def forward(self, handcrafted_features, dinov2_features, bioclip2_logits):
+        dinov2_out = self.dinov2_proj(dinov2_features)
+        handcraft_out = self.handcraft_branch(handcrafted_features)
+        bioclip2_out = self.bioclip2_branch(bioclip2_logits)
+        shared_features = self.fusion(torch.cat([dinov2_out, handcraft_out, bioclip2_out], dim=1))
+        class_output = self.classifier(shared_features)
+        projected_feature = self.prototype_proj(shared_features)
+        return class_output, projected_feature
+
+# ==============================================================================
+# 3. MANAGER CLASS & EXPORTED FUNCTION
+# ==============================================================================
+class ModelManager:
+    def __init__(self):
+        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        print(f"Initializing SPA Ensemble on {self.device}...")
+        
+        # --- CONFIG: YOUR MODEL REPO ID ---
+        # Using the correct repo ID provided
+        self.REPO_ID = "FrAnKu34t23/ensemble_models_plant" 
+        
+        self.class_to_idx, self.idx_to_class, self.id_to_name = self.load_class_info()
+        self.num_classes = len(self.class_to_idx)
+        print(f"SPA Ensemble: Loaded {self.num_classes} classes.")
+
+        # 1. Load DINOv2
+        print("SPA Ensemble: Loading DINOv2...")
+        self.dinov2 = torch.hub.load('facebookresearch/dinov2', 'dinov2_vitl14')
+        self.dinov2.to(self.device).eval()
+        self.dinov2_transform = transforms.Compose([
+            transforms.Resize(256), transforms.CenterCrop(224), 
+            transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+        ])
+
+        # 2. Load BioCLIP
+        self.bioclip = BioCLIP2ZeroShot(self.device, self.class_to_idx, self.id_to_name)
+        
+        # 3. Download & Load Scaler
+        print("SPA Ensemble: Downloading Scaler...")
+        try:
+            # Now fetching scaler.joblib from the Model Repo
+            scaler_path = hf_hub_download(repo_id=self.REPO_ID, filename="scaler.joblib")
+            self.scaler = joblib.load(scaler_path)
+            print("✓ Scaler downloaded and loaded.")
+        except Exception as e:
+            print(f"Warning: Could not download scaler from {self.REPO_ID}: {e}.")
+            print("Using dummy scaler (predictions may be inaccurate).")
+            self.scaler = StandardScaler()
+            # FIX: Fit on 40 zeros instead of 49 to match the feature reduction
+            self.scaler.fit(np.zeros((1, 40))) 
+
+        # 4. Download & Load Ensemble Models
+        self.models = []
+        hidden_dims = [384, 448, 512, 576, 640]
+        dropout_rates = [0.2, 0.25, 0.3, 0.35, 0.4]
+        
+        print(f"SPA Ensemble: Downloading Models from {self.REPO_ID}...")
+        for i in range(5):
+            filename = f"ensemble_model_{i}.pth"
+            try:
+                # Download
+                model_path = hf_hub_download(repo_id=self.REPO_ID, filename=filename)
+                
+                # Load
+                # FIX: Passed num_handcrafted_features=40 and prototype_dim=768 to match weights
+                model = EnsembleClassifier(
+                    num_handcrafted_features=40, dinov2_dim=1024, bioclip2_dim=self.num_classes,
+                    num_classes=self.num_classes, hidden_dim=hidden_dims[i], dropout_rate=dropout_rates[i],
+                    prototype_dim=768
+                )
+                state_dict = torch.load(model_path, map_location=self.device)
+                model.load_state_dict(state_dict)
+                model.to(self.device).eval()
+                self.models.append(model)
+                print(f"✓ Loaded {filename}")
+            except Exception as e:
+                print(f"Failed to load {filename}: {e}")
+
+    def load_class_info(self):
+        class_to_idx = {}
+        id_to_name = {}
+        
+        species_path = LIST_DIR / "species_list.txt"
+        train_path = LIST_DIR / "train.txt"
+        
+        classes_set = set()
+        
+        if train_path.exists():
+            with open(train_path, 'r') as f:
+                for line in f:
+                    parts = line.strip().split()
+                    if len(parts) >= 2: classes_set.add(parts[1])
+        elif species_path.exists():
+             with open(species_path, 'r') as f:
+                for line in f:
+                    parts = line.strip().split(";", 1)
+                    classes_set.add(parts[0].strip())
+        else:
+            classes_set = {str(i) for i in range(100)}
+
+        sorted_classes = sorted(list(classes_set))
+        class_to_idx = {cls: idx for idx, cls in enumerate(sorted_classes)}
+        idx_to_class = {idx: cls for cls, idx in class_to_idx.items()}
+        
+        if species_path.exists():
+            with open(species_path, 'r') as f:
+                for line in f:
+                    if ";" in line:
+                        parts = line.strip().split(";", 1)
+                        id_to_name[parts[0].strip()] = parts[1].strip()
+        return class_to_idx, idx_to_class, id_to_name
+
+    def predict(self, image):
+        if image is None: return {}
+        img_pil = image.convert("RGB")
+        
+        # 1. Handcrafted Features
+        hc_feats = FeatureExtractor.extract_all_features(img_pil)
+        hc_vector = np.array([hc_feats[k] for k in sorted(hc_feats.keys())]).reshape(1, -1)
+        hc_vector = self.scaler.transform(hc_vector) 
+        hc_tensor = torch.FloatTensor(hc_vector).to(self.device)
+
+        # 2. DINOv2 Features
+        dino_input = self.dinov2_transform(img_pil).unsqueeze(0).to(self.device)
+        with torch.no_grad():
+            dino_feats = self.dinov2(dino_input)
+            dino_feats = dino_feats / (dino_feats.norm(dim=-1, keepdim=True) + 1e-8)
+
+        # 3. BioCLIP Features
+        bioclip_logits = self.bioclip.predict_zero_shot_logits(img_pil)
+        bioclip_tensor = torch.FloatTensor(bioclip_logits).unsqueeze(0).to(self.device)
+
+        # 4. Ensemble Prediction
+        all_probs = []
+        if not self.models: return {"Error": "SPA Models not loaded"}
+
+        for model in self.models:
+            with torch.no_grad():
+                probs, _ = model(hc_tensor, dino_feats, bioclip_tensor)
+                probs = F.softmax(probs, dim=1).cpu().numpy()[0]
+                all_probs.append(probs)
+        
+        final_ens_probs = np.mean(all_probs, axis=0)
+        
+        # 5. Hybrid Routing
+        exp_logits = np.exp(bioclip_logits)
+        bioclip_probs = exp_logits / np.sum(exp_logits)
+        
+        ens_pred_idx = np.argmax(final_ens_probs)
+        ens_conf = final_ens_probs[ens_pred_idx]
+        
+        if ens_conf < CONFIDENCE_THRESHOLD:
+             final_probs = (final_ens_probs + bioclip_probs) / 2
+        else:
+             final_probs = final_ens_probs
+
+        # 6. Formatting
+        top_k = 5
+        top_indices = np.argsort(final_probs)[::-1][:top_k]
+        results = {}
+        for idx in top_indices:
+            class_id = self.idx_to_class[idx]
+            name = self.id_to_name.get(class_id, class_id)
+            score = float(final_probs[idx])
+            results[f"{name} ({class_id})"] = score
+            
+        return results
+
+# Initialize Singleton
+try:
+    spa_manager = ModelManager()
+except Exception as e:
+    print(f"CRITICAL ERROR initializing SPA: {e}")
+    spa_manager = None
+
+def predict_spa(image):
+    if spa_manager is None:
+        return {"Error": "SPA System failed to initialize."}
+    return spa_manager.predict(image)