Create app.py

app.py

@@ -0,0 +1,134 @@
+import gradio as gr
+import numpy as np
+import pickle
+from keras.models import load_model
+from tensorflow.keras.applications.resnet50 import ResNet50, preprocess_input
+from tensorflow.keras.preprocessing import image
+from tensorflow.keras.preprocessing.sequence import pad_sequences
+
+# Load the trained model
+model = load_model('image_captioning_model.h5')
+
+# Load ResNet50 for feature extraction
+resnet_model = ResNet50(include_top=False, weights='imagenet', pooling='avg', input_shape=(224, 224, 3))
+
+# Load word mappings and config
+with open('words_to_indices.pkl', 'rb') as f:
+    words_to_indices = pickle.load(f)
+
+with open('indices_to_words.pkl', 'rb') as f:
+    indices_to_words = pickle.load(f)
+
+with open('config.pkl', 'rb') as f:
+    config = pickle.load(f)
+    vocab_size = config['vocab_size']
+    max_length = config['max_length']
+
+def extract_features(img):
+    """Extract features from image using ResNet50"""
+    img = img.resize((224, 224))
+    x = image.img_to_array(img)
+    x = np.expand_dims(x, axis=0)
+    x = preprocess_input(x)
+    features = resnet_model.predict(x, verbose=0)
+    return features.squeeze()
+
+def greedy_search(photo):
+    """Generate caption using greedy search"""
+    photo = photo.reshape(1, 2048)
+    in_text = '<start>'
+    
+    for i in range(max_length):
+        sequence = [words_to_indices.get(s, 0) for s in in_text.split(" ")]
+        sequence = pad_sequences([sequence], maxlen=max_length, padding='post')
+        y_pred = model.predict([photo, sequence], verbose=0)
+        y_pred = np.argmax(y_pred[0])
+        word = indices_to_words.get(y_pred, 'Unk')
+        in_text += ' ' + word
+        if word == '<end>':
+            break
+    
+    final = in_text.split()
+    final = final[1:-1]
+    return " ".join(final)
+
+def beam_search(photo, k=3):
+    """Generate caption using beam search"""
+    photo = photo.reshape(1, 2048)
+    in_text = '<start>'
+    sequence = [words_to_indices.get(s, 0) for s in in_text.split(" ")]
+    sequence = pad_sequences([sequence], maxlen=max_length, padding='post')
+    y_pred = model.predict([photo, sequence], verbose=0)
+    
+    predicted = []
+    y_pred = y_pred.reshape(-1)
+    for i in range(y_pred.shape[0]):
+        predicted.append((i, y_pred[i]))
+    predicted = sorted(predicted, key=lambda x: x[1], reverse=True)
+    
+    b_search = []
+    for i in range(k):
+        word = indices_to_words.get(predicted[i][0], 'Unk')
+        b_search.append((in_text + ' ' + word, predicted[i][1]))
+    
+    for idx in range(max_length):
+        b_search_square = []
+        for text in b_search:
+            if text[0].split(" ")[-1] == "<end>":
+                break
+            sequence = [words_to_indices.get(s, 0) for s in text[0].split(" ")]
+            sequence = pad_sequences([sequence], maxlen=max_length, padding='post')
+            y_pred = model.predict([photo, sequence], verbose=0)
+            predicted = []
+            y_pred = y_pred.reshape(-1)
+            for i in range(y_pred.shape[0]):
+                predicted.append((i, y_pred[i]))
+            predicted = sorted(predicted, key=lambda x: x[1], reverse=True)
+            for i in range(k):
+                word = indices_to_words.get(predicted[i][0], 'Unk')
+                b_search_square.append((text[0] + ' ' + word, predicted[i][1] * text[1]))
+        if len(b_search_square) > 0:
+            b_search = sorted(b_search_square, key=lambda x: x[1], reverse=True)[:k]
+    
+    final = b_search[0][0].split()
+    final = final[1:-1]
+    return " ".join(final)
+
+def generate_caption(img, search_method="Greedy Search"):
+    """Main function to generate caption from image"""
+    try:
+        # Extract features
+        features = extract_features(img)
+        
+        # Generate caption based on selected method
+        if search_method == "Greedy Search":
+            caption = greedy_search(features)
+        else:  # Beam Search
+            caption = beam_search(features, k=3)
+        
+        return caption
+    except Exception as e:
+        return f"Error generating caption: {str(e)}"
+
+# Create Gradio interface
+demo = gr.Interface(
+    fn=generate_caption,
+    inputs=[
+        gr.Image(type="pil", label="Upload Image"),
+        gr.Radio(
+            choices=["Greedy Search", "Beam Search (k=3)"],
+            value="Greedy Search",
+            label="Caption Generation Method"
+        )
+    ],
+    outputs=gr.Textbox(label="Generated Caption"),
+    title="Image Captioning with LSTM",
+    description="Upload an image to generate a descriptive caption. Choose between Greedy Search or Beam Search for caption generation.",
+    examples=[
+        # Add example images if you have them
+    ],
+    theme=gr.themes.Soft()
+)
+
+if __name__ == "__main__":
+    demo.launch()