import gradio as gr
import tensorflow as tf
from transformers import GPT2LMHeadModel, GPT2Tokenizer
import numpy as np
from tensorflow.image import resize

# Load the image classification model
image_classification_model = tf.keras.models.load_model("mobilenetfashion_v2.h5")

# Load pre-trained GPT-2 model and tokenizer
gpt2_model_name = "diamantrsd/copywriting-otomatis"
gpt2_model = GPT2LMHeadModel.from_pretrained(gpt2_model_name)
gpt2_tokenizer = GPT2Tokenizer.from_pretrained(gpt2_model_name)

def classify_and_generate_text(image, keywords=""):
    try:
        # Convert Gradio Image interface output to a NumPy array
        img_array = image.astype('float32') / 255.0

        # Resize the image to the expected shape (224, 224)
        img_array_resized = resize(img_array, (224, 224))

        # Classify the resized image using the image classification model
        class_label = image_classification_model.predict(np.expand_dims(img_array_resized, axis=0))

        # Map class label to corresponding category (adjust as needed)
        category = map_class_label_to_category(class_label)

        # Generate text based on the category and keywords using the GPT-2 model
        generated_text = generate_text_with_gpt2(category, keywords)

        return generated_text
    except Exception as e:
        return f"Error: {str(e)}"

def map_class_label_to_category(class_label):
    # Map the class label to a category (replace with your own mapping)
    categories = [ 'Backpack','Celana Panjang','Celana Pendek','Dompet',
        'Dress','Kacamata','Kaos', 'Kaos Kaki','Kemeja', 'Outerwear','Sandal', 'Sepatu',
        'Sepatu Flat','Tas','Topi']
    return categories[np.argmax(class_label, axis=-1)[0]]

def generate_text_with_gpt2(product_category, keywords):
    prompt = f"Produk: {product_category}, Keywords: {keywords}, Copywriting:"
    input_ids = gpt2_tokenizer.encode(prompt, return_tensors="pt")

    # Adjust parameters as needed
    max_length = 50
    no_repeat_ngram_size = 3
    top_k = 50
    top_p = 0.95

    output = gpt2_model.generate(
        input_ids,
        max_length=max_length,
        no_repeat_ngram_size=no_repeat_ngram_size,
        top_k=top_k,
        top_p=top_p,
        pad_token_id=gpt2_tokenizer.eos_token_id,
        num_return_sequences=1
    )

    generated_text = gpt2_tokenizer.decode(output[0], skip_special_tokens=True)
    return generated_text

# Create Gradio Interface
iface = gr.Interface(
    fn=classify_and_generate_text,
    inputs=[gr.Image(image_mode="RGB"), gr.Textbox(placeholder="", label="Keywords")],
    outputs="text",
    live=True
)

# Launch the Gradio Interface
iface.launch()