import gradio as gr from PIL import Image import torch import torch.nn as nn import torchvision.transforms as transforms import torchvision.models as models import os from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, pipeline import torch import gc # Set device device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') # Load the main classifier (Main_Classifier_best_model.pth) main_model = models.resnet18(weights=None) # Updated: weights=None num_ftrs = main_model.fc.in_features main_model.fc = nn.Linear(num_ftrs, 3) # 3 classes: Soda drinks, Clothing, Mobile Phones main_model.load_state_dict(torch.load('Main_Classifier_best_model.pth', map_location=device, weights_only=True)) # Updated: weights_only=True main_model = main_model.to(device) main_model.eval() # Define class names for the main classifier based on folder structure main_class_names = ['Clothing', 'Mobile Phones', 'Soda drinks'] # Sub-classifier models def load_soda_drinks_model(): model = models.resnet18(weights=None) # Updated: weights=None num_ftrs = model.fc.in_features model.fc = nn.Linear(num_ftrs, 3) # 3 classes: Miranda, Pepsi, Seven Up model.load_state_dict(torch.load('Soda_drinks_best_model.pth', map_location=device, weights_only=True)) # Updated model = model.to(device) model.eval() return model def load_clothing_model(): model = models.resnet18(weights=None) # Updated: weights=None num_ftrs = model.fc.in_features model.fc = nn.Linear(num_ftrs, 3) # 2 classes: Pants, T-Shirt model.load_state_dict(torch.load('Clothes_best_model.pth', map_location=device, weights_only=True)) # Updated model = model.to(device) model.eval() return model def load_mobile_phones_model(): model = models.resnet18(weights=None) # Updated: weights=None num_ftrs = model.fc.in_features model.fc = nn.Linear(num_ftrs, 2) # 2 classes: Apple, Samsung model.load_state_dict(torch.load('Phone_best_model.pth', map_location=device, weights_only=True)) # Updated model = model.to(device) model.eval() return model def convert_to_rgb(image): """ Converts 'P' mode images with transparency to 'RGBA', and then to 'RGB'. This is to avoid transparency issues during model training. """ if image.mode in ('P', 'RGBA'): return image.convert('RGB') return image # Define preprocessing transformations (same used during training) preprocess = transforms.Compose([ transforms.Lambda(convert_to_rgb), transforms.Resize((224, 224)), # Resize here, no need for shape argument in gr.Image transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) # ImageNet normalization ]) # Load Meta's LLaMA model for generating product descriptions def load_llama(): model_name = "meta-llama/Meta-Llama-3-8B-Instruct" token = os.getenv("HUGGINGFACE_TOKEN") tokenizer = AutoTokenizer.from_pretrained(model_name, use_auth_token=token) # Define the quantization configuration to optimize model performance quantization_config = BitsAndBytesConfig( activation_quantization_bits=16, weight_quantization_bits=16, bnb_4bit_quant_type="nf4", bnb_4bit_use_double_quant=True, bnb_4bit_compute_dtype=torch.bfloat16 ) # Load the model with the specified configurations model = AutoModelForCausalLM.from_pretrained( model_name, quantization_config=quantization_config, use_auth_token=token ).to(device) # Initialize the text generation pipeline with the prepared model text_generation = pipeline( "text-generation", model=model, tokenizer=tokenizer ) return tokenizer, model # # Load Meta's LLaMA model for generating product descriptions # def load_llama(): # model_name = "meta-llama/Llama-3.2-1B-Instruct" # token = os.getenv("HUGGINGFACE_TOKEN") # tokenizer = AutoTokenizer.from_pretrained(model_name, use_auth_token=token) # model = AutoModelForCausalLM.from_pretrained(model_name, use_auth_token=token).to(device) # return tokenizer, model llama_tokenizer, llama_model = load_llama() # Generate product description using LLaMA def generate_description(category, subclass): prompt = f"Generate a detailed and engaging product description for a {category} of type {subclass}." inputs = llama_tokenizer.encode(prompt, return_tensors="pt").to(device) outputs = llama_model.generate(inputs, max_length=100, do_sample=True, temperature=0.7, top_k=50, top_p=0.95) description = llama_tokenizer.decode(outputs[0], skip_special_tokens=True) return description def classify_image(image): # Open the image using PIL image = Image.fromarray(image) # Preprocess the image input_image = preprocess(image).unsqueeze(0).to(device) # Perform inference with the main classifier with torch.no_grad(): output = main_model(input_image) probabilities = torch.nn.functional.softmax(output[0], dim=0) confidence, predicted_class = torch.max(probabilities, 0) # Main classifier result main_prediction = main_class_names[predicted_class] main_confidence = confidence.item() if main_confidence <=0.90: main_prediction = 'Others' main_confidence = 100-main_confidence sub_prediction = "Undefined" sub_confidence = -100 description = None # Load and apply the sub-classifier based on the main classification if main_prediction in ['Clothing', 'Mobile Phones', 'Soda drinks']: if main_prediction == 'Soda drinks': soda_model = load_soda_drinks_model() sub_class_names = ['Miranda', 'Pepsi', 'Seven Up'] with torch.no_grad(): sub_output = soda_model(input_image) elif main_prediction == 'Clothing': clothing_model = load_clothing_model() sub_class_names = ['Pants', 'T-Shirt','others'] with torch.no_grad(): sub_output = clothing_model(input_image) elif main_prediction == 'Mobile Phones': phones_model = load_mobile_phones_model() sub_class_names = ['Apple', 'Samsung'] with torch.no_grad(): sub_output = phones_model(input_image) # Perform inference with the sub-classifier sub_probabilities = torch.nn.functional.softmax(sub_output[0], dim=0) sub_confidence, sub_predicted_class = torch.max(sub_probabilities, 0) sub_prediction = sub_class_names[sub_predicted_class] sub_confidence = sub_confidence.item() if sub_confidence < 0.90 : sub_prediction = "Others" sub_confidence = 100- sub_confidence description=None else: # Generate product description description = generate_description(main_prediction, sub_prediction) return f"Main Predicted Class: {main_prediction} (Confidence: {main_confidence:.4f})", \ f"Sub Predicted Class: {sub_prediction} (Confidence: {sub_confidence:.4f})", \ f"Product Description: {description}" # Gradio interface (updated) image_input = gr.Image(image_mode="RGB") # Removed shape argument output_text = gr.Textbox() gr.Interface(fn=classify_image, inputs=image_input, outputs=[output_text], title="Main and Sub-Classifier System product description ", description="Upload an image to classify whether it belongs to Clothing, Mobile Phones, or Soda Drinks. Based on the prediction, it will further classify within the subcategory and generate a detailed product description .", theme="default").launch()