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("klasifikasi_pt1.h5")

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



def classify_and_generate_text(image):
    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 using the GPT-2 model
        generated_text = generate_text_with_gpt2(category)

        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 = ['Blazer', 'Blouse', 'Cardigan', 'Dress', 'Jacket',
                  'Jeans', 'Jumpsuit', 'Romper', 'Shorts', 'Skirts',
                  'Sweater', 'Sweatpants', 'Tank', 'Tee', 'Top']
    return categories[np.argmax(class_label, axis=-1)[0]]

def generate_text_with_gpt2(product_category):
    prompt = f"Produk: {product_category}, 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"),
    outputs="text",
    live=True
)

# Launch the Gradio Interface
iface.launch()