eizen-nsga 1.0.0__py3-none-any.whl → 1.0.1__py3-none-any.whl

eizen_nsga/__init__.py

@@ -0,0 +1,15 @@
+"""
+eizen-nsga: Simple inference package for NSGA-Net trained models
+
+Usage:
+    from eizen_nsga import NASModel
+
+    model = NASModel("trained_model.zip")
+    result = model.predict("image.jpg")
+"""
+
+__version__ = "1.0.0"
+
+from .model import NASModel
+
+__all__ = ['NASModel']

eizen_nsga/model.py

@@ -0,0 +1,372 @@
+"""
+NASModel - Universal model class for NSGA-Net trained models
+
+Simple interface like ultralytics YOLO:
+    model = NASModel("model.zip")
+    result = model("image.jpg")
+"""
+
+import os
+import re
+import ast
+import sys
+import zipfile
+import tempfile
+import shutil
+from typing import Union, Optional, Dict, Any, List
+from urllib.parse import urlparse
+from urllib.request import urlretrieve
+
+import torch
+import torch.nn as nn
+import numpy as np
+from PIL import Image
+import torchvision.transforms as transforms
+
+
+class NASModel:
+    """
+    Universal model class for NSGA-Net trained models.
+
+    Simple interface like YOLO:
+        model = NASModel("model.zip")
+        result = model.predict("image.jpg")
+
+    Supports:
+        - SOTA: Computer vision (YOLOv8, ResNet, EfficientNet)
+        - NN: Tabular data (coming soon)
+        - Transformer: LLMs (coming soon)
+    """
+
+    def __init__(self, path: str, device: Optional[str] = None):
+        """
+        Load a trained NSGA-Net model.
+
+        Args:
+            path: Path or URL to model ZIP file or directory
+                  Supports: local paths, HTTP/HTTPS URLs, S3, GCS, etc.
+            device: 'cuda', 'cpu', or None (auto-detect)
+        """
+        self.source_path = path
+        self._temp_dir = None
+
+        # Extract ZIP if needed
+        model_dir = self._extract_model(path)
+
+        # Find weights and config
+        weights_path, log_path = self._find_model_files(model_dir)
+
+        # Parse configuration
+        self.config = self._parse_log(log_path)
+        self.model_category = self.config.get('modelCategory', 'SOTA')
+        self.task = self.config.get('task', 'classification')
+
+        # Set device
+        if device is None:
+            device = 'cuda' if torch.cuda.is_available() else 'cpu'
+        self.device = device
+
+        # Build and load model
+        self.model = self._build_model()
+        self._load_weights(weights_path)
+        self.model.eval()
+
+        print(f"✅ Loaded {self.model_category} model for {self.task}")
+
+    def _extract_model(self, path: str) -> str:
+        """Extract ZIP if needed, return model directory. Supports local paths and URLs."""
+        # Check if path is a URL
+        parsed = urlparse(path)
+        is_url = parsed.scheme in ('http', 'https', 's3', 'gs', 'ftp')
+
+        if is_url:
+            # Download from URL to temporary file
+            print(f"📥 Downloading model from URL...")
+            self._temp_dir = tempfile.mkdtemp()
+            local_zip = os.path.join(self._temp_dir, 'model.zip')
+            urlretrieve(path, local_zip)
+            print(f"✅ Downloaded to {local_zip}")
+            path = local_zip
+
+        # Extract ZIP if it's a zip file
+        if path.endswith('.zip') and os.path.isfile(path):
+            if not self._temp_dir:
+                self._temp_dir = tempfile.mkdtemp()
+            with zipfile.ZipFile(path, 'r') as zf:
+                zf.extractall(self._temp_dir)
+            return self._temp_dir
+        return path
+
+    def _find_model_files(self, model_dir: str) -> tuple:
+        """Find weights.pt and log.txt in model directory."""
+        weights_path = None
+        log_path = None
+
+        # Check root first
+        if os.path.exists(os.path.join(model_dir, 'weights.pt')):
+            weights_path = os.path.join(model_dir, 'weights.pt')
+            log_path = os.path.join(model_dir, 'log.txt')
+        else:
+            # Search subdirectories
+            for root, dirs, files in os.walk(model_dir):
+                if 'weights.pt' in files:
+                    weights_path = os.path.join(root, 'weights.pt')
+                    log_path = os.path.join(root, 'log.txt')
+                    break
+
+        if not weights_path or not os.path.exists(weights_path):
+            raise FileNotFoundError(f"weights.pt not found in {model_dir}")
+
+        return weights_path, log_path
+
+    def _parse_log(self, log_path: str) -> Dict[str, Any]:
+        """Parse log.txt to extract configuration."""
+        config = {}
+
+        if not log_path or not os.path.exists(log_path):
+            return config
+
+        with open(log_path, 'r') as f:
+            content = f.read()
+
+        # Extract all key-value pairs
+        patterns = {
+            'modelCategory': r'modelCategory[:\s]*(\w+)',
+            'task': r'Task[:\s]*(\w+)',
+            'genome': r'Genome[:\s]*(.+)',
+            'backbone': r'backbone[:\s]*(\S+)',
+            'search_space': r'search_space[:\s]*(\w+)',
+            'num_cells': r'num_cells[:\s]*(\d+)',
+            'num_classes': r'num_classes[:\s]*(\d+)',
+            'image_size': r'image_size[:\s]*(\d+)',
+        }
+
+        for key, pattern in patterns.items():
+            match = re.search(pattern, content, re.IGNORECASE | re.MULTILINE)
+            if match:
+                value = match.group(1).strip()
+                try:
+                    config[key] = ast.literal_eval(value)
+                except:
+                    config[key] = value
+
+        return config
+
+    def _build_model(self) -> nn.Module:
+        """Build model architecture based on category."""
+        category = self.model_category.upper()
+
+        if category == 'SOTA':
+            return self._build_sota_model()
+        elif category == 'NN':
+            return self._build_nn_model()
+        elif category == 'TRANSFORMER':
+            return self._build_transformer_model()
+        else:
+            raise ValueError(f"Unknown model category: {category}")
+
+    def _build_sota_model(self) -> nn.Module:
+        """Build SOTA adapter model from genome."""
+        genome = self.config.get('genome')
+        backbone = self.config.get('backbone', 'yolov8n')
+        search_space = self.config.get('search_space', 'micro')
+        num_cells = int(self.config.get('num_cells', 3))
+        num_classes = int(self.config.get('num_classes', 10))
+
+        if genome is None:
+            raise ValueError("genome not found in log.txt")
+
+        # Import SOTA modules from package
+        from .sota.micro_models import GeneralAdapterNetwork
+        from .sota.macro_models import MacroAdapterNetwork
+        from .sota import micro_encoding, macro_encoding
+
+        # Decode and build
+        if search_space == 'macro':
+            genotype = macro_encoding.decode(genome)
+            model = MacroAdapterNetwork(
+                genotype=genotype,
+                backbone_name=backbone,
+                num_classes=num_classes,
+                num_cells=num_cells,
+                pretrained=True,
+                freeze_backbone=True
+            )
+        else:
+            genotype = micro_encoding.decode(genome)
+            model = GeneralAdapterNetwork(
+                genotype=genotype,
+                backbone_name=backbone,
+                num_classes=num_classes,
+                num_cells=num_cells,
+                pretrained=True,
+                freeze_backbone=True
+            )
+
+        return model.to(self.device)
+
+    def _build_nn_model(self) -> nn.Module:
+        """Build NN model (tabular)."""
+        raise NotImplementedError("NN model support coming soon")
+
+    def _build_transformer_model(self) -> nn.Module:
+        """Build Transformer model (LLM)."""
+        raise NotImplementedError("Transformer model support coming soon")
+
+    def _load_weights(self, weights_path: str):
+        """Load trained weights into model."""
+        state_dict = torch.load(weights_path, map_location=self.device)
+
+        # Filter FLOPs counting keys
+        filtered = {
+            k: v for k, v in state_dict.items()
+            if 'total_ops' not in k and 'total_params' not in k
+        }
+
+        # Load weights
+        missing, unexpected = self.model.load_state_dict(filtered, strict=False)
+
+        if missing:
+            print(f"⚠️  Missing keys: {len(missing)}")
+        if unexpected:
+            print(f"⚠️  Unexpected keys: {len(unexpected)}")
+
+    def predict(self, source: Union[str, np.ndarray, Image.Image], **kwargs) -> Dict[str, Any]:
+        """
+        Run inference on image(s).
+
+        Args:
+            source: Image path, PIL Image, or numpy array
+            **kwargs: Additional arguments (top_k, conf_threshold, etc.)
+
+        Returns:
+            Prediction results dictionary
+        """
+        if self.task == 'detection':
+            return self._detect(source, **kwargs)
+        else:
+            return self._classify(source, **kwargs)
+
+    def _classify(self, image: Union[str, np.ndarray, Image.Image],
+                  top_k: int = 5) -> Dict[str, Any]:
+        """Classification inference."""
+        # Preprocess
+        tensor = self._preprocess_image(image)
+
+        # Inference
+        with torch.no_grad():
+            outputs, _ = self.model(tensor)
+            probs = torch.softmax(outputs, dim=1)
+
+        # Get top-k
+        top_probs, top_indices = torch.topk(probs[0], min(top_k, probs.shape[1]))
+
+        predictions = [
+            {'class': idx.item(), 'confidence': prob.item()}
+            for idx, prob in zip(top_indices, top_probs)
+        ]
+
+        return {
+            'predictions': predictions,
+            'class': predictions[0]['class'],
+            'confidence': predictions[0]['confidence']
+        }
+
+    def _detect(self, image: Union[str, np.ndarray, Image.Image],
+                conf_threshold: float = 0.25,
+                iou_threshold: float = 0.45) -> Dict[str, Any]:
+        """Detection inference."""
+        # Preprocess with larger size for detection
+        tensor = self._preprocess_image(image, size=640)
+
+        # Inference
+        with torch.no_grad():
+            outputs, _ = self.model(tensor)
+
+        # Post-process detections
+        detections = self._postprocess_detections(outputs, conf_threshold, iou_threshold)
+
+        return {
+            'detections': detections,
+            'count': len(detections)
+        }
+
+    def _preprocess_image(self, image: Union[str, np.ndarray, Image.Image],
+                          size: int = None) -> torch.Tensor:
+        """Preprocess image for model input."""
+        if size is None:
+            size = self.config.get('image_size', 224)
+
+        # Load image
+        if isinstance(image, str):
+            image = Image.open(image).convert('RGB')
+        elif isinstance(image, np.ndarray):
+            image = Image.fromarray(image).convert('RGB')
+
+        # Transform
+        transform = transforms.Compose([
+            transforms.Resize((size, size)),
+            transforms.ToTensor(),
+            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+        ])
+
+        return transform(image).unsqueeze(0).to(self.device)
+
+    def _postprocess_detections(self, outputs: torch.Tensor,
+                                 conf_threshold: float,
+                                 iou_threshold: float) -> List[Dict]:
+        """Post-process detection outputs."""
+        detections = []
+
+        try:
+            from .sota.detection_heads import decode_predictions, non_max_suppression
+
+            grid_size = outputs.size(2) if outputs.dim() >= 3 else 7
+            pred_boxes, pred_scores, pred_labels = decode_predictions(
+                outputs, conf_threshold=conf_threshold, grid_size=grid_size
+            )
+
+            for i in range(len(pred_boxes)):
+                if len(pred_boxes[i]) > 0:
+                    keep = non_max_suppression(
+                        pred_boxes[i], pred_scores[i], pred_labels[i],
+                        iou_threshold=iou_threshold
+                    )
+                    for j in keep:
+                        detections.append({
+                            'box': pred_boxes[i][j].cpu().numpy().tolist(),
+                            'confidence': pred_scores[i][j].item(),
+                            'class': pred_labels[i][j].item()
+                        })
+        except ImportError:
+            pass
+
+        return detections
+
+    def __call__(self, source: Union[str, np.ndarray, Image.Image], **kwargs):
+        """Call model like a function (similar to YOLO)."""
+        return self.predict(source, **kwargs)
+
+    def to(self, device: str):
+        """Move model to device."""
+        self.device = device
+        self.model = self.model.to(device)
+        return self
+
+    def info(self) -> Dict[str, Any]:
+        """Get model information."""
+        return {
+            'category': self.model_category,
+            'task': self.task,
+            'backbone': self.config.get('backbone'),
+            'num_classes': self.config.get('num_classes'),
+            'device': self.device
+        }
+
+    def __del__(self):
+        """Cleanup temp directory."""
+        if self._temp_dir and os.path.exists(self._temp_dir):
+            shutil.rmtree(self._temp_dir, ignore_errors=True)
+
+    def __repr__(self):
+        return f"NASModel(category={self.model_category}, task={self.task}, device={self.device})"

{eizen_nsga-1.0.0.dist-info → eizen_nsga-1.0.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: eizen-nsga
-Version: 1.0.0
+Version: 1.0.1
 Summary: Simple inference package for NSGA-Net trained models
 Home-page: https://github.com/eizen-ai/nsga-net
 Author: Eizen.ai Team

eizen_nsga-1.0.1.dist-info/RECORD

@@ -0,0 +1,15 @@
+eizen_nsga/__init__.py,sha256=SYzEu-7pzFy8Ul-bkjnNzuRQtvVQcfbTLBG1Dapun5g,291
+eizen_nsga/model.py,sha256=bIasvlnrVmaXnaYza7J9dMDoFIK-X1LC2PcPNIzdaBI,13224
+eizen_nsga/sota/__init__.py,sha256=3hyX_CR86TvV2RYQaES0FhZdtklUBu3DB7HAv1Z0yGo,525
+eizen_nsga/sota/detection_heads.py,sha256=Tx7qgqc0_BBtaIenxNHg_dIwbDHyOQz3QsamNEa9gX0,11206
+eizen_nsga/sota/macro_encoding.py,sha256=lwV8Nptwt014UWW-2rAmivYqS-5Jfs3g8DCmzhRzHsA,3037
+eizen_nsga/sota/macro_models.py,sha256=okitzfXyaRPjGb5lAh9BNwqUVjTaUqaj7jwypKKVVm0,12254
+eizen_nsga/sota/micro_encoding.py,sha256=LrhgzhgnGMP2Up7Uqy0zT5nedzefqDXl2HJ8YL2TOVM,4820
+eizen_nsga/sota/micro_models.py,sha256=piLM6jt4LEiiQuox6AL6I-K0W0XcGLcL2qHWZgPbpvA,22062
+eizen_nsga/sota/micro_operations.py,sha256=3E2JL1eSzWYmtI4X9omBkocWcAoZslXdkIj8CS1h4dQ,8702
+eizen_nsga/sota/model_registry.py,sha256=lQQKUKCVWEQEql5mDZrlwRJEIVaqdokNWEEN4SWrnA0,17750
+eizen_nsga-1.0.1.dist-info/licenses/LICENSE,sha256=_HF_TY-jv6lFXR4QcG1iTnMYgDC3MfslE5S0OvUGlQA,1091
+eizen_nsga-1.0.1.dist-info/METADATA,sha256=D-hn9vHIdz88l9C2MgzxLHZU2nH8gj3DVfSKipCiEdQ,5268
+eizen_nsga-1.0.1.dist-info/WHEEL,sha256=wUyA8OaulRlbfwMtmQsvNngGrxQHAvkKcvRmdizlJi0,92
+eizen_nsga-1.0.1.dist-info/top_level.txt,sha256=BPCkm-TWk4CpW-H-eKYfaa6KWJuepOHlKR3M5UhB4-4,11
+eizen_nsga-1.0.1.dist-info/RECORD,,

eizen_nsga-1.0.1.dist-info/top_level.txt

@@ -0,0 +1 @@
+eizen_nsga

eizen_nsga-1.0.0.dist-info/RECORD

@@ -1,13 +0,0 @@
-eizen_nsga-1.0.0.dist-info/licenses/LICENSE,sha256=_HF_TY-jv6lFXR4QcG1iTnMYgDC3MfslE5S0OvUGlQA,1091
-sota/__init__.py,sha256=3hyX_CR86TvV2RYQaES0FhZdtklUBu3DB7HAv1Z0yGo,525
-sota/detection_heads.py,sha256=Tx7qgqc0_BBtaIenxNHg_dIwbDHyOQz3QsamNEa9gX0,11206
-sota/macro_encoding.py,sha256=lwV8Nptwt014UWW-2rAmivYqS-5Jfs3g8DCmzhRzHsA,3037
-sota/macro_models.py,sha256=okitzfXyaRPjGb5lAh9BNwqUVjTaUqaj7jwypKKVVm0,12254
-sota/micro_encoding.py,sha256=LrhgzhgnGMP2Up7Uqy0zT5nedzefqDXl2HJ8YL2TOVM,4820
-sota/micro_models.py,sha256=piLM6jt4LEiiQuox6AL6I-K0W0XcGLcL2qHWZgPbpvA,22062
-sota/micro_operations.py,sha256=3E2JL1eSzWYmtI4X9omBkocWcAoZslXdkIj8CS1h4dQ,8702
-sota/model_registry.py,sha256=lQQKUKCVWEQEql5mDZrlwRJEIVaqdokNWEEN4SWrnA0,17750
-eizen_nsga-1.0.0.dist-info/METADATA,sha256=uQP38LPaimMrKbFtdFhUXuBaOsIMaQntO1oJRvjDLK4,5268
-eizen_nsga-1.0.0.dist-info/WHEEL,sha256=wUyA8OaulRlbfwMtmQsvNngGrxQHAvkKcvRmdizlJi0,92
-eizen_nsga-1.0.0.dist-info/top_level.txt,sha256=pPnDGDIcVIcY4aYyBjID0_x2rcMP9_Q3G-ZHKSZGwQI,5
-eizen_nsga-1.0.0.dist-info/RECORD,,

eizen_nsga-1.0.0.dist-info/top_level.txt

@@ -1 +0,0 @@
-sota