RSK World - Image Classification Dataset - Project Files Browser | RSK World

scripts/predict.py

"""
===================================================================================
    Project: Image Classification Dataset
    Description: Image prediction utilities with pre-trained models.
    
    Author: Molla Samser
    Email: help@rskworld.in
    Phone: +91 93305 39277
    Website: https://rskworld.in
    
    © 2025 RSK World. All rights reserved.
===================================================================================
"""

import os
import numpy as np
from PIL import Image
from typing import Tuple, List, Dict, Optional
import json


class ImagePredictor:
    """
    Image classification predictor with support for multiple frameworks.
    
    Author: Molla Samser
    Website: https://rskworld.in
    """
    
    def __init__(
        self,
        model_path: Optional[str] = None,
        class_names_path: Optional[str] = None,
        framework: str = 'tensorflow'
    ):
        """
        Initialize the predictor.
        
        Args:
            model_path: Path to the trained model
            class_names_path: Path to JSON file with class names
            framework: 'tensorflow' or 'pytorch'
        """
        self.model = None
        self.class_names = None
        self.framework = framework
        self.image_size = (224, 224)
        
        if model_path:
            self.load_model(model_path)
        
        if class_names_path:
            self.load_class_names(class_names_path)
    
    def load_model(self, model_path: str):
        """Load a trained model."""
        if self.framework == 'tensorflow':
            try:
                import tensorflow as tf
                self.model = tf.keras.models.load_model(model_path)
                print(f"Model loaded from: {model_path}")
            except Exception as e:
                print(f"Error loading TensorFlow model: {e}")
        elif self.framework == 'pytorch':
            try:
                import torch
                self.model = torch.load(model_path)
                self.model.eval()
                print(f"Model loaded from: {model_path}")
            except Exception as e:
                print(f"Error loading PyTorch model: {e}")
    
    def load_class_names(self, path: str):
        """Load class names from JSON file."""
        with open(path, 'r') as f:
            self.class_names = json.load(f)
        print(f"Loaded {len(self.class_names)} class names")
    
    def preprocess_image(self, image_path: str) -> np.ndarray:
        """
        Preprocess an image for prediction.
        
        Args:
            image_path: Path to the image file
            
        Returns:
            Preprocessed image array
        """
        img = Image.open(image_path).convert('RGB')
        img = img.resize(self.image_size, Image.Resampling.LANCZOS)
        img_array = np.array(img) / 255.0
        
        return img_array
    
    def predict(self, image_path: str, top_k: int = 5) -> List[Dict]:
        """
        Predict the class of an image.
        
        Args:
            image_path: Path to the image
            top_k: Number of top predictions to return
            
        Returns:
            List of dictionaries with class names and probabilities
        """
        if self.model is None:
            raise ValueError("No model loaded. Call load_model() first.")
        
        # Preprocess
        img = self.preprocess_image(image_path)
        img_batch = np.expand_dims(img, axis=0)
        
        # Predict
        if self.framework == 'tensorflow':
            predictions = self.model.predict(img_batch, verbose=0)[0]
        elif self.framework == 'pytorch':
            import torch
            with torch.no_grad():
                img_tensor = torch.tensor(img_batch).permute(0, 3, 1, 2).float()
                predictions = self.model(img_tensor).numpy()[0]
        
        # Get top-k predictions
        top_indices = np.argsort(predictions)[::-1][:top_k]
        
        results = []
        for idx in top_indices:
            class_name = self.class_names[idx] if self.class_names else f"Class {idx}"
            results.append({
                'class_id': int(idx),
                'class_name': class_name,
                'probability': float(predictions[idx]),
                'percentage': f"{predictions[idx] * 100:.2f}%"
            })
        
        return results
    
    def predict_batch(self, image_paths: List[str], top_k: int = 3) -> List[List[Dict]]:
        """
        Predict classes for multiple images.
        
        Args:
            image_paths: List of image paths
            top_k: Number of top predictions per image
            
        Returns:
            List of prediction results for each image
        """
        return [self.predict(path, top_k) for path in image_paths]
    
    def predict_from_array(self, image_array: np.ndarray, top_k: int = 5) -> List[Dict]:
        """
        Predict from a numpy array.
        
        Args:
            image_array: Image as numpy array (H, W, C)
            top_k: Number of top predictions
            
        Returns:
            List of predictions
        """
        if self.model is None:
            raise ValueError("No model loaded.")
        
        # Ensure correct shape
        if len(image_array.shape) == 3:
            image_array = np.expand_dims(image_array, axis=0)
        
        # Normalize if needed
        if image_array.max() > 1:
            image_array = image_array / 255.0
        
        predictions = self.model.predict(image_array, verbose=0)[0]
        top_indices = np.argsort(predictions)[::-1][:top_k]
        
        results = []
        for idx in top_indices:
            class_name = self.class_names[idx] if self.class_names else f"Class {idx}"
            results.append({
                'class_id': int(idx),
                'class_name': class_name,
                'probability': float(predictions[idx])
            })
        
        return results


def predict_image(image_path: str, model_path: str, class_names: List[str]) -> Dict:
    """
    Simple function to predict a single image.
    
    Author: Molla Samser | https://rskworld.in
    
    Args:
        image_path: Path to the image
        model_path: Path to the model
        class_names: List of class names
        
    Returns:
        Prediction result dictionary
    """
    predictor = ImagePredictor(framework='tensorflow')
    predictor.load_model(model_path)
    predictor.class_names = class_names
    
    results = predictor.predict(image_path, top_k=1)
    
    return {
        'image': image_path,
        'predicted_class': results[0]['class_name'],
        'confidence': results[0]['percentage']
    }


class EnsemblePredictor:
    """
    Ensemble prediction using multiple models.
    
    Author: Molla Samser
    Website: https://rskworld.in
    """
    
    def __init__(self, model_paths: List[str], weights: Optional[List[float]] = None):
        """
        Initialize ensemble predictor.
        
        Args:
            model_paths: List of paths to trained models
            weights: Optional weights for each model
        """
        self.predictors = []
        self.weights = weights or [1.0 / len(model_paths)] * len(model_paths)
        
        for path in model_paths:
            predictor = ImagePredictor(model_path=path)
            self.predictors.append(predictor)
    
    def predict(self, image_path: str, top_k: int = 5) -> List[Dict]:
        """
        Make ensemble prediction.
        
        Args:
            image_path: Path to the image
            top_k: Number of top predictions
            
        Returns:
            Ensemble prediction results
        """
        all_predictions = []
        
        for predictor in self.predictors:
            img = predictor.preprocess_image(image_path)
            img_batch = np.expand_dims(img, axis=0)
            pred = predictor.model.predict(img_batch, verbose=0)[0]
            all_predictions.append(pred)
        
        # Weighted average
        ensemble_pred = np.zeros_like(all_predictions[0])
        for pred, weight in zip(all_predictions, self.weights):
            ensemble_pred += pred * weight
        
        # Get top-k
        top_indices = np.argsort(ensemble_pred)[::-1][:top_k]
        
        results = []
        class_names = self.predictors[0].class_names
        
        for idx in top_indices:
            class_name = class_names[idx] if class_names else f"Class {idx}"
            results.append({
                'class_id': int(idx),
                'class_name': class_name,
                'probability': float(ensemble_pred[idx])
            })
        
        return results


if __name__ == "__main__":
    # Example usage
    print("=" * 60)
    print("Image Classification Dataset - Prediction Utilities")
    print("=" * 60)
    print("Author: Molla Samser")
    print("Website: https://rskworld.in")
    print("Email: help@rskworld.in")
    print("=" * 60)
    
    print("\nExample usage:")
    print("""
    # Single model prediction
    predictor = ImagePredictor(
        model_path='models/image_classifier.h5',
        class_names_path='models/class_names.json'
    )
    
    results = predictor.predict('test_image.jpg', top_k=5)
    for r in results:
        print(f"{r['class_name']}: {r['percentage']}")
    
    # Ensemble prediction
    ensemble = EnsemblePredictor([
        'models/model1.h5',
        'models/model2.h5',
        'models/model3.h5'
    ])
    
    results = ensemble.predict('test_image.jpg')
    """)

310 lines•9.7 KB

python

scripts/train_model.py

Raw Download

"""
===================================================================================
    Project: Image Classification Dataset
    Description: CNN model training script for image classification.
    
    Author: Molla Samser
    Email: help@rskworld.in
    Phone: +91 93305 39277
    Website: https://rskworld.in
    
    © 2025 RSK World. All rights reserved.
===================================================================================
"""

import os
import numpy as np
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers, models, callbacks
from tensorflow.keras.preprocessing.image import ImageDataGenerator
from typing import Tuple, Optional
import matplotlib.pyplot as plt
from datetime import datetime


class ImageClassificationModel:
    """
    CNN model for image classification tasks.
    
    Author: Molla Samser
    Website: https://rskworld.in
    Email: help@rskworld.in
    """
    
    def __init__(
        self,
        input_shape: Tuple[int, int, int] = (224, 224, 3),
        num_classes: int = 10,
        model_name: str = 'cnn_classifier'
    ):
        """
        Initialize the model.
        
        Args:
            input_shape: Shape of input images (height, width, channels)
            num_classes: Number of classification categories
            model_name: Name for the model
        """
        self.input_shape = input_shape
        self.num_classes = num_classes
        self.model_name = model_name
        self.model = None
        self.history = None
        
    def build_cnn_model(self) -> keras.Model:
        """
        Build a custom CNN architecture.
        
        Returns:
            Compiled Keras model
        """
        model = models.Sequential([
            # First Convolutional Block
            layers.Conv2D(32, (3, 3), activation='relu', input_shape=self.input_shape),
            layers.BatchNormalization(),
            layers.MaxPooling2D((2, 2)),
            
            # Second Convolutional Block
            layers.Conv2D(64, (3, 3), activation='relu'),
            layers.BatchNormalization(),
            layers.MaxPooling2D((2, 2)),
            
            # Third Convolutional Block
            layers.Conv2D(128, (3, 3), activation='relu'),
            layers.BatchNormalization(),
            layers.MaxPooling2D((2, 2)),
            
            # Fourth Convolutional Block
            layers.Conv2D(256, (3, 3), activation='relu'),
            layers.BatchNormalization(),
            layers.MaxPooling2D((2, 2)),
            
            # Flatten and Dense Layers
            layers.Flatten(),
            layers.Dense(512, activation='relu'),
            layers.Dropout(0.5),
            layers.Dense(256, activation='relu'),
            layers.Dropout(0.3),
            layers.Dense(self.num_classes, activation='softmax')
        ])
        
        model.compile(
            optimizer=keras.optimizers.Adam(learning_rate=0.001),
            loss='categorical_crossentropy',
            metrics=['accuracy']
        )
        
        self.model = model
        return model
    
    def build_transfer_learning_model(
        self,
        base_model: str = 'resnet50',
        trainable_layers: int = 20
    ) -> keras.Model:
        """
        Build a transfer learning model using pre-trained weights.
        
        Args:
            base_model: Name of the pre-trained model ('resnet50', 'vgg16', 'efficientnet')
            trainable_layers: Number of layers to make trainable
            
        Returns:
            Compiled Keras model
        """
        # Select base model
        if base_model == 'resnet50':
            base = keras.applications.ResNet50(
                weights='imagenet',
                include_top=False,
                input_shape=self.input_shape
            )
        elif base_model == 'vgg16':
            base = keras.applications.VGG16(
                weights='imagenet',
                include_top=False,
                input_shape=self.input_shape
            )
        elif base_model == 'efficientnet':
            base = keras.applications.EfficientNetB0(
                weights='imagenet',
                include_top=False,
                input_shape=self.input_shape
            )
        else:
            raise ValueError(f"Unknown base model: {base_model}")
        
        # Freeze early layers
        for layer in base.layers[:-trainable_layers]:
            layer.trainable = False
        
        # Build complete model
        model = models.Sequential([
            base,
            layers.GlobalAveragePooling2D(),
            layers.Dense(512, activation='relu'),
            layers.Dropout(0.5),
            layers.Dense(256, activation='relu'),
            layers.Dropout(0.3),
            layers.Dense(self.num_classes, activation='softmax')
        ])
        
        model.compile(
            optimizer=keras.optimizers.Adam(learning_rate=0.0001),
            loss='categorical_crossentropy',
            metrics=['accuracy']
        )
        
        self.model = model
        return model
    
    def create_data_generator(
        self,
        augment: bool = True
    ) -> ImageDataGenerator:
        """
        Create an image data generator with optional augmentation.
        
        Args:
            augment: Whether to apply data augmentation
            
        Returns:
            ImageDataGenerator instance
        """
        if augment:
            datagen = ImageDataGenerator(
                rescale=1./255,
                rotation_range=20,
                width_shift_range=0.2,
                height_shift_range=0.2,
                horizontal_flip=True,
                zoom_range=0.2,
                shear_range=0.2,
                fill_mode='nearest'
            )
        else:
            datagen = ImageDataGenerator(rescale=1./255)
            
        return datagen
    
    def train(
        self,
        train_data: Tuple[np.ndarray, np.ndarray],
        val_data: Optional[Tuple[np.ndarray, np.ndarray]] = None,
        epochs: int = 50,
        batch_size: int = 32,
        augment: bool = True,
        early_stopping: bool = True,
        save_best: bool = True
    ) -> keras.callbacks.History:
        """
        Train the model.
        
        Args:
            train_data: Tuple of (images, labels)
            val_data: Optional validation data
            epochs: Number of training epochs
            batch_size: Batch size
            augment: Whether to use data augmentation
            early_stopping: Whether to use early stopping
            save_best: Whether to save the best model
            
        Returns:
            Training history
        """
        X_train, y_train = train_data
        
        # Convert labels to one-hot encoding
        y_train_onehot = keras.utils.to_categorical(y_train, self.num_classes)
        
        if val_data:
            X_val, y_val = val_data
            y_val_onehot = keras.utils.to_categorical(y_val, self.num_classes)
            validation_data = (X_val, y_val_onehot)
        else:
            validation_data = None
        
        # Create callbacks
        callback_list = []
        
        if early_stopping:
            callback_list.append(callbacks.EarlyStopping(
                monitor='val_loss' if val_data else 'loss',
                patience=10,
                restore_best_weights=True
            ))
        
        if save_best:
            timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
            model_path = f'models/{self.model_name}_{timestamp}.h5'
            os.makedirs('models', exist_ok=True)
            callback_list.append(callbacks.ModelCheckpoint(
                model_path,
                monitor='val_accuracy' if val_data else 'accuracy',
                save_best_only=True
            ))
        
        # Add learning rate scheduler
        callback_list.append(callbacks.ReduceLROnPlateau(
            monitor='val_loss' if val_data else 'loss',
            factor=0.5,
            patience=5,
            min_lr=1e-7
        ))
        
        # Train with or without augmentation
        if augment:
            datagen = self.create_data_generator(augment=True)
            
            self.history = self.model.fit(
                datagen.flow(X_train, y_train_onehot, batch_size=batch_size),
                epochs=epochs,
                validation_data=validation_data,
                callbacks=callback_list,
                steps_per_epoch=len(X_train) // batch_size
            )
        else:
            self.history = self.model.fit(
                X_train, y_train_onehot,
                batch_size=batch_size,
                epochs=epochs,
                validation_data=validation_data,
                callbacks=callback_list
            )
        
        return self.history
    
    def evaluate(
        self,
        test_data: Tuple[np.ndarray, np.ndarray]
    ) -> Tuple[float, float]:
        """
        Evaluate the model on test data.
        
        Args:
            test_data: Tuple of (images, labels)
            
        Returns:
            Tuple of (loss, accuracy)
        """
        X_test, y_test = test_data
        y_test_onehot = keras.utils.to_categorical(y_test, self.num_classes)
        
        loss, accuracy = self.model.evaluate(X_test, y_test_onehot)
        
        print(f"\nTest Results:")
        print(f"  Loss: {loss:.4f}")
        print(f"  Accuracy: {accuracy:.4f} ({accuracy*100:.2f}%)")
        
        return loss, accuracy
    
    def predict(self, images: np.ndarray) -> np.ndarray:
        """
        Make predictions on new images.
        
        Args:
            images: Array of images
            
        Returns:
            Predicted class indices
        """
        predictions = self.model.predict(images)
        return np.argmax(predictions, axis=1)
    
    def plot_training_history(self, save_path: Optional[str] = None):
        """
        Plot training history.
        
        Args:
            save_path: Optional path to save the plot
        """
        if self.history is None:
            print("No training history available.")
            return
        
        fig, axes = plt.subplots(1, 2, figsize=(14, 5))
        
        # Plot accuracy
        axes[0].plot(self.history.history['accuracy'], label='Train')
        if 'val_accuracy' in self.history.history:
            axes[0].plot(self.history.history['val_accuracy'], label='Validation')
        axes[0].set_title('Model Accuracy')
        axes[0].set_xlabel('Epoch')
        axes[0].set_ylabel('Accuracy')
        axes[0].legend()
        axes[0].grid(True)
        
        # Plot loss
        axes[1].plot(self.history.history['loss'], label='Train')
        if 'val_loss' in self.history.history:
            axes[1].plot(self.history.history['val_loss'], label='Validation')
        axes[1].set_title('Model Loss')
        axes[1].set_xlabel('Epoch')
        axes[1].set_ylabel('Loss')
        axes[1].legend()
        axes[1].grid(True)
        
        plt.tight_layout()
        
        if save_path:
            plt.savefig(save_path, dpi=150)
            print(f"Plot saved to {save_path}")
        
        plt.show()
    
    def save_model(self, path: str):
        """Save the model to disk."""
        self.model.save(path)
        print(f"Model saved to {path}")
    
    def load_model(self, path: str):
        """Load a model from disk."""
        self.model = keras.models.load_model(path)
        print(f"Model loaded from {path}")


if __name__ == "__main__":
    # Example usage
    # Author: Molla Samser | https://rskworld.in
    
    print("=" * 60)
    print("Image Classification Dataset - Model Training")
    print("=" * 60)
    print("Author: Molla Samser")
    print("Website: https://rskworld.in")
    print("Email: help@rskworld.in")
    print("=" * 60)
    
    # Example: Create and train a model
    classifier = ImageClassificationModel(
        input_shape=(224, 224, 3),
        num_classes=12,
        model_name='image_classifier'
    )
    
    # Build CNN model
    model = classifier.build_cnn_model()
    model.summary()
    
    print("\nModel built successfully!")
    print("To train the model, use the train() method with your data.")
    print("\nExample:")
    print("  classifier.train(train_data=(X_train, y_train), epochs=50)")

395 lines•12.6 KB

python

Theme Settings

Color Scheme

Display Options

Font Size