Spaces:
Runtime error
Runtime error
File size: 6,684 Bytes
eb179f1 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 |
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
from transformers import AutoTokenizer, AutoModelForCausalLM
import os
# 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/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() |