Update app.py

app.py

@@ -14,7 +14,6 @@ from sklearn.preprocessing import LabelEncoder
 import joblib
 from PIL import UnidentifiedImageError, Image
 import gradio as gr
-import html
 
 # Optimized Constants
 MAX_TEXT_LENGTH = 100
@@ -22,36 +21,28 @@ EMBEDDING_DIM = 50
 IMAGE_SIZE = 160
 BATCH_SIZE = 64
 
-# Store model examples
-model_examples = {}
-
 def load_and_preprocess_data(subset_size=10000):
+    # Load dataset
     dataset = load_dataset("thefcraft/civitai-stable-diffusion-337k")
     dataset_subset = dataset['train'].shuffle(seed=42).select(range(subset_size))
-    dataset_subset = dataset_subset.filter(lambda x: not x['nsfw'])
     
-    for item in dataset_subset:
-        if item['Model'] not in model_examples:
-            model_examples[item['Model']] = item['url']
+    # Filter out NSFW content
+    dataset_subset = dataset_subset.filter(lambda x: not x['nsfw'])
     
     return dataset_subset
 
 def process_text_data(dataset_subset):
-    text_data = ["default prompt" for _ in dataset_subset]
+    # Combine prompt and negative prompt
+    text_data = [f"{sample['prompt']} {sample['negativePrompt']}" for sample in dataset_subset]
+    
+    # Tokenize text
     tokenizer = Tokenizer(num_words=10000)
     tokenizer.fit_on_texts(text_data)
     sequences = tokenizer.texts_to_sequences(text_data)
     text_data_padded = pad_sequences(sequences, maxlen=MAX_TEXT_LENGTH)
+    
     return text_data_padded, tokenizer
 
-def download_image(url):
-    try:
-        response = requests.get(url, timeout=5)
-        response.raise_for_status()
-        return Image.open(requests.get(url, stream=True).raw)
-    except:
-        return None
-
 def process_image_data(dataset_subset):
     image_dir = 'civitai_images'
     os.makedirs(image_dir, exist_ok=True)
@@ -64,6 +55,7 @@ def process_image_data(dataset_subset):
         img_path = os.path.join(image_dir, os.path.basename(img_url))
         
         try:
+            # Download and save image
             response = requests.get(img_url, timeout=5)
             response.raise_for_status()
             
@@ -73,6 +65,7 @@ def process_image_data(dataset_subset):
             with open(img_path, 'wb') as f:
                 f.write(response.content)
             
+            # Load and preprocess image
             img = image.load_img(img_path, target_size=(IMAGE_SIZE, IMAGE_SIZE))
             img_array = image.img_to_array(img)
             img_array = preprocess_input(img_array)
@@ -86,21 +79,26 @@ def process_image_data(dataset_subset):
     return np.array(image_data), valid_indices
 
 def create_multimodal_model(num_words, num_classes):
-    image_input = Input(shape=(IMAGE_SIZE, IMAGE_SIZE, 3), name='image_input')
+    # Image input branch (CNN)
+    image_input = Input(shape=(IMAGE_SIZE, IMAGE_SIZE, 3))
     cnn_base = ResNet50(weights='imagenet', include_top=False, pooling='avg')
     
+    # Freeze most of the ResNet50 layers
     for layer in cnn_base.layers[:-10]:
         layer.trainable = False
         
     cnn_features = cnn_base(image_input)
     
-    text_input = Input(shape=(MAX_TEXT_LENGTH,), name='text_input')
+    # Text input branch
+    text_input = Input(shape=(MAX_TEXT_LENGTH,))
     embedding_layer = Embedding(num_words, EMBEDDING_DIM)(text_input)
     flatten_text = Flatten()(embedding_layer)
     text_features = Dense(128, activation='relu')(flatten_text)
     
+    # Combine features
     combined = Concatenate()([cnn_features, text_features])
     
+    # Simplified fully connected layers
     x = Dense(256, activation='relu')(combined)
     output = Dense(num_classes, activation='softmax')(x)
     
@@ -108,16 +106,24 @@ def create_multimodal_model(num_words, num_classes):
     return model
 
 def train_model():
+    # Load and preprocess data
     dataset_subset = load_and_preprocess_data()
+    
+    # Process text data
     text_data_padded, tokenizer = process_text_data(dataset_subset)
+    
+    # Process image data
     image_data, valid_indices = process_image_data(dataset_subset)
     
+    # Get valid text data and labels
     text_data_padded = text_data_padded[valid_indices]
     model_names = [dataset_subset[i]['Model'] for i in valid_indices]
     
+    # Encode labels
     label_encoder = LabelEncoder()
     encoded_labels = label_encoder.fit_transform(model_names)
     
+    # Create and compile model
     model = create_multimodal_model(
         num_words=10000,
         num_classes=len(label_encoder.classes_)
@@ -129,6 +135,7 @@ def train_model():
         metrics=['accuracy']
     )
     
+    # Train model
     history = model.fit(
         [image_data, text_data_padded],
         encoded_labels,
@@ -137,77 +144,68 @@ def train_model():
         validation_split=0.2
     )
     
-    model.save('multimodal_model.keras')
+    # Save models and encoders with correct extensions
+    model.save('multimodal_model.keras')  # Changed from 'multimodal_model'
     joblib.dump(tokenizer, 'tokenizer.pkl')
     joblib.dump(label_encoder, 'label_encoder.pkl')
-    joblib.dump(model_examples, 'model_examples.pkl')
     
     return model, tokenizer, label_encoder
 
-def get_recommendations(image_input, model, tokenizer, label_encoder, top_k=5):
+def get_recommendations(image_input, text_input, model, tokenizer, label_encoder, top_k=5):
+    # Preprocess image
     img_array = image.img_to_array(image_input)
     img_array = tf.image.resize(img_array, (IMAGE_SIZE, IMAGE_SIZE))
     img_array = preprocess_input(img_array)
     img_array = np.expand_dims(img_array, axis=0)
     
-    text_sequence = tokenizer.texts_to_sequences(["default prompt"])
+    # Preprocess text
+    text_sequence = tokenizer.texts_to_sequences([text_input])
     text_padded = pad_sequences(text_sequence, maxlen=MAX_TEXT_LENGTH)
     
-    predictions = model.predict([img_array, text_padded], verbose=0)
+    # Get predictions
+    predictions = model.predict([img_array, text_padded])
     top_indices = np.argsort(predictions[0])[-top_k:][::-1]
     
-    recommendations = []
-    for idx in top_indices:
-        model_name = label_encoder.inverse_transform([idx])[0]
-        if model_name in model_examples:
-            example_image = download_image(model_examples[model_name])
-            if example_image:
-                recommendations.append((model_name, example_image))
+    # Get recommended model names and confidence scores
+    recommendations = [
+        (label_encoder.inverse_transform([idx])[0], predictions[0][idx])
+        for idx in top_indices
+    ]
     
     return recommendations
 
-def create_gallery_html(recommendations):
-    html_content = '<div style="display: grid; grid-template-columns: repeat(auto-fill, minmax(200px, 1fr)); gap: 20px; padding: 20px;">'
-    
-    for model_name, image in recommendations:
-        # Konversi gambar ke base64 jika diperlukan
-        html_content += f'''
-            <div style="text-align: center; border: 1px solid #ddd; padding: 10px; border-radius: 8px;">
-                <img src="{html.escape(image.url) if hasattr(image, 'url') else ''}" 
-                     style="width: 100%; height: auto; border-radius: 4px;" />
-                <p style="margin-top: 10px; font-weight: bold;">{html.escape(model_name)}</p>
-            </div>
-        '''
-    
-    html_content += '</div>'
-    return html_content
-
 def create_gradio_interface():
-    model = tf.keras.models.load_model('multimodal_model.keras')
+    # Load saved models with correct path
+    model = tf.keras.models.load_model('multimodal_model.keras')  # Changed from 'multimodal_model'
     tokenizer = joblib.load('tokenizer.pkl')
     label_encoder = joblib.load('label_encoder.pkl')
     
-    def predict(img):
-        recommendations = get_recommendations(img, model, tokenizer, label_encoder)
-        return gr.HTML(create_gallery_html(recommendations))
+    def predict(img, text):
+        recommendations = get_recommendations(img, text, model, tokenizer, label_encoder)
+        return "\n".join([f"Model: {name}, Confidence: {conf:.2f}" for name, conf in recommendations])
     
     interface = gr.Interface(
         fn=predict,
-        inputs=gr.Image(type="pil", label="Upload Image"),
-        outputs=gr.HTML(),
-        title="AI Model Recommendation System",
-        description="Upload an image to get similar model recommendations"
+        inputs=[
+            gr.Image(type="pil", label="Upload Image"),
+            gr.Textbox(label="Enter Prompt")
+        ],
+        outputs=gr.Textbox(label="Recommended Models"),
+        title="Multimodal Model Recommendation System",
+        description="Upload an image and enter a prompt to get model recommendations"
     )
     
     return interface
 
 if __name__ == "__main__":
-    if not os.path.exists('multimodal_model.keras'):
+    # Train model if not already trained
+    if not os.path.exists('multimodal_model.keras'):  # Changed from 'multimodal_model'
         print("Training new model...")
         model, tokenizer, label_encoder = train_model()
         print("Training completed!")
     else:
         print("Loading existing model...")
     
+    # Launch Gradio interface
     interface = create_gradio_interface()
     interface.launch()