iBrokeTheCode/Multimodal_Product_Classification

1from json import load2from typing import Any, Dict, Optional3 4from numpy import array, expand_dims, float32, ndarray, transpose, zeros5from PIL import Image6from sentence_transformers import SentenceTransformer7from tensorflow import constant8from tensorflow.keras.models import load_model9from transformers import TFConvNextV2Model10 11# 📌 GLOBAL VARIABLES (categories)12CATEGORY_MAP: Dict[str, str] = {}13CLASS_LABELS = []14 15 16def build_category_map(categories_json_path: str):17    """18    Builds a flat dictionary and a list of category labels by traversing the hierarchical categories.json file.19    """20    global CATEGORY_MAP, CLASS_LABELS21 22    try:23        with open(categories_json_path, "r") as f:24            categories_data = load(f)25    except FileNotFoundError:26        print(27            f"❌ Error: {categories_json_path} not found. Using hardcoded labels as fallback."28        )29        return30 31    category_map = {}32 33    model_trained_ids = [34        "abcat0100000",35        "abcat0200000",36        "abcat0207000",37        "abcat0300000",38        "abcat0400000",39        "abcat0500000",40        "abcat0700000",41        "abcat0800000",42        "abcat0900000",43        "cat09000",44        "pcmcat128500050004",45        "pcmcat139900050002",46        "pcmcat242800050021",47        "pcmcat252700050006",48        "pcmcat312300050015",49        "pcmcat332000050000",50    ]51 52    def traverse_categories(categories):53        for category in categories:54            category_map[category["id"]] = category["name"]55            if "subCategories" in category and category["subCategories"]:56                traverse_categories(category["subCategories"])57            if "path" in category and category["path"]:58                for path_item in category["path"]:59                    category_map[path_item["id"]] = path_item["name"]60 61    traverse_categories(categories_data)62 63    CATEGORY_MAP = category_map64    CLASS_LABELS = model_trained_ids65 66 67# 📌 LOAD MODELS68print("💬 Loading embedding models...")69try:70    text_embedding_model = SentenceTransformer("all-MiniLM-L6-v2")71    image_feature_extractor = TFConvNextV2Model.from_pretrained(72        "facebook/convnextv2-tiny-22k-224"73    )74    print("✅ Embedding models loaded successfully!")75except Exception as e:76    print(f"❌ Error loading embedding models: {e}")77    text_embedding_model, image_feature_extractor = None, None78 79# Load the final classification models (MLP heads)80print("💬 Loading classification models...")81try:82    text_model = load_model("./models/text_model")83    image_model = load_model("./models/image_model")84    multimodal_model = load_model("./models/multimodal_model")85    print("✅ Classification models loaded successfully!")86except Exception as e:87    print(f"❌ Error loading classification models: {e}")88    text_model, image_model, multimodal_model = None, None, None89 90# Generate category map and class labels list91build_category_map("./data/raw/categories.json")92 93 94# 📌 EMBEDDING FUNCTIONS95def get_text_embeddings(text: Optional[str]) -> ndarray:96    """97    Generates a dense embedding vector from a text string.98 99    Args:100        text (Optional[str]): The input text. Can be None or an empty string.101 102    Returns:103        np.ndarray: A NumPy array of shape (1, 384) representing the text104                    embedding. Returns a zero vector if the input is empty.105    """106    # Handle cases where no text is provided107    if not text or not text.strip():108        # Returns a zero vector with the correct dimension (384)109        return zeros(110            (1, text_embedding_model.get_sentence_embedding_dimension()), dtype=float32111        )112 113    # Use the pre-trained SentenceTransformer to encode the text114    embeddings = text_embedding_model.encode([text])115    return array(embeddings, dtype=float32)116 117 118def get_image_embeddings(image_path: Optional[str]) -> ndarray:119    """120    Preprocesses an image and generates an embedding vector using a pre-trained model.121 122    Args:123        image_path (Optional[str]): The file path to the image.124 125    Returns:126        np.ndarray: A NumPy array of shape (1, 768) representing the image127                    embedding. Returns a zero vector if no image is provided.128    """129    # Handle cases where no image is provided130    if image_path is None:131        return zeros((1, 768), dtype=float32)132 133    # Load the image and convert to RGB format134    image = Image.open(image_path).convert("RGB")135 136    # Resize the image to the model's expected input size (224x224)137    image = image.resize((224, 224), Image.Resampling.LANCZOS)138 139    # Convert to NumPy array and add a batch dimension (1, H, W, C)140    image_array = array(image, dtype=float32)141    image_array = expand_dims(image_array, axis=0)142 143    # Transpose the array to match the model's channel order (1, C, H, W)144    image_array = transpose(image_array, (0, 3, 1, 2))145 146    # Normalize the pixel values (not strictly necessary for this model, but good practice)147    image_array = image_array / 255.0148 149    # Pass the preprocessed image through the feature extractor model150    embeddings_output = image_feature_extractor(constant(image_array))151 152    # Extract the final embedding from the pooler_output153    embeddings = embeddings_output.pooler_output154 155    return embeddings.numpy()156 157 158# 📌 MAIN PREDICTION FUNCTION159def predict(160    mode: str, text: Optional[str], image_path: Optional[str]161) -> Dict[str, Any]:162    """163    Predicts the category of a product based on the selected mode.164 165    Args:166        mode (str): The prediction mode ("Multimodal", "Text Only", "Image Only").167        text (Optional[str]): The product description text.168        image_path (Optional[str]): The file path to the product image.169 170    Returns:171        Dict[str, Any]: A dictionary of class labels and their corresponding172                        prediction probabilities. Returns an empty dictionary173                        if the mode is invalid.174    """175    # Generate embeddings for both inputs176    text_emb = get_text_embeddings(text)177    image_emb = get_image_embeddings(image_path)178 179    # Get predictions based on the selected mode180    if mode == "Multimodal":181        predictions = multimodal_model.predict([text_emb, image_emb])182    elif mode == "Text Only":183        predictions = text_model.predict(text_emb)184    elif mode == "Image Only":185        predictions = image_model.predict(image_emb)186    else:187        # Return an empty dictionary if the mode is not recognized188        return {}189 190    # Format the output into a dictionary with labels and probabilities191    # The model's output is a 2D array, so we take the first row (index 0)192    prediction_dict_raw = dict(zip(CLASS_LABELS, predictions[0]))193 194    #  Map the raw IDs to human-readable names195    prediction_dict_mapped = {}196    for class_id, probability in prediction_dict_raw.items():197        # Get the human-readable name, defaulting to the raw ID if not found198        category_name = CATEGORY_MAP.get(class_id, class_id)199        prediction_dict_mapped[category_name] = probability200 201    # Sort the dictionary by probability in descending order for a cleaner display202    sorted_predictions = dict(203        sorted(prediction_dict_mapped.items(), key=lambda item: item[1], reverse=True)204    )205 206    return sorted_predictions207