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

# Load pre-trained GPT-2 model and tokenizer
gpt2_model_name = "diamantrsd/copywriting-otomatis-v1"
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 = ['Blouse', 'Blazer', 'Kemeja', 'Bomber', 'Anorak', 
                  'Tee', 'Tank', 'Top', 'Sweater', 'Flannel', 'Hoodie', 
                  'Cardigan', 'Jaket', 'Henley', 'Poncho', 'Jersey', 'Turtleneck', 'Parka', 
                  'Peacoat', 'Halter', 'Skirt', 'Shorts', 'Jeans', 'Joggers', 'Sweatpants', 'Jeggings', 
                  'Cutoffs', 'Sweatshorts', 'Leggings', 'Culottes', 'Chinos', 'Culottes', 'Sarung', 'Gauchos', 'Jodhpurs', 
                  'Capris', 'Dress', 'Romper', 'Coat', 'Kimono', 'Jumpsuit', 
                  'Robe', 'Caftan', 'Kaftan', 'Coverup', 'Onesie']
    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()