Upload 5 files

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+faiss_index.idx filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,3 +1,13 @@
----
-license: apache-2.0
----
+
+# 🏛️ RAG Image Captioning with Landmark Location
+
+This model generates captions for monument/landmark images using a retrieval-augmented generation approach.
+
+## How it works:
+- Uses CLIP to extract image embeddings.
+- Retrieves top-k similar captions via FAISS.
+- Generates a detailed caption with name and location using T5.
+
+## Example
+Input: 🏰 Image of the Taj Mahal  
+Output: _"The place might be: Agra. The Taj Mahal is a white marble mausoleum located in Agra, India."_

faiss_index.idx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a43f9410efe919810cd35354d77d7396cdee594a5d3998aadb6bc03606274332
+size 3446829

inference.py

@@ -0,0 +1,110 @@
+
+from PIL import Image
+import torch
+from transformers import CLIPProcessor, CLIPModel, T5Tokenizer, T5ForConditionalGeneration
+from sentence_transformers import SentenceTransformer
+import faiss
+import numpy as np
+import json
+import spacy
+
+# Load models and resources
+clip_model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
+clip_processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
+text_encoder = SentenceTransformer('all-MiniLM-L6-v2')
+tokenizer = T5Tokenizer.from_pretrained("t5-small")
+generator = T5ForConditionalGeneration.from_pretrained("t5-small")
+nlp = spacy.load("en_core_web_sm")
+
+# Load FAISS index and captions
+faiss_index = faiss.read_index("./faiss_index.idx")
+with open("./captions.json", "r", encoding="utf-8") as f:
+    captions = json.load(f)
+
+def extract_image_features(image):
+    """
+    Extract image features using CLIP model.
+    Input: PIL Image or image path (str).
+    Output: Normalized image embedding (numpy array).
+    """
+    try:
+        # Handle both PIL Image and file path
+        if isinstance(image, str):
+            image = Image.open(image).convert("RGB")
+        else:
+            image = image.convert("RGB")
+        inputs = clip_processor(images=image, return_tensors="pt")
+        with torch.no_grad():
+            features = clip_model.get_image_features(**inputs)
+        features = torch.nn.functional.normalize(features, p=2, dim=-1)
+        return features.squeeze(0).cpu().numpy().astype("float32")
+    except Exception as e:
+        print(f"Error extracting features: {e}")
+        return None
+
+def retrieve_similar_captions(image_embedding, k=5):
+    """
+    Retrieve k most similar captions using FAISS index.
+    Input: Image embedding (numpy array).
+    Output: List of captions.
+    """
+    if image_embedding.ndim == 1:
+        image_embedding = image_embedding.reshape(1, -1)
+    D, I = faiss_index.search(image_embedding, k)
+    return [captions[i] for i in I[0]]
+
+def extract_location_names(texts):
+    """
+    Extract location names from captions using spaCy.
+    Input: List of captions.
+    Output: List of unique location names.
+    """
+    names = []
+    for text in texts:
+        doc = nlp(text)
+        for ent in doc.ents:
+            if ent.label_ in ["GPE", "LOC", "FAC"]:
+                names.append(ent.text)
+    return list(set(names))
+
+def generate_caption_from_retrieved(retrieved_captions):
+    """
+    Generate a caption from retrieved captions using T5.
+    Input: List of retrieved captions.
+    Output: Generated caption (str).
+    """
+    locations = extract_location_names(retrieved_captions)
+    location_hint = f"The place might be: {', '.join(locations)}. " if locations else ""
+    prompt = location_hint + " ".join(retrieved_captions) + " Generate a caption with the landmark name:"
+    inputs = tokenizer(prompt, return_tensors="pt", truncation=True)
+    outputs = generator.generate(
+        input_ids=inputs.input_ids,
+        attention_mask=inputs.attention_mask,
+        max_length=300,
+        num_beams=5,
+        early_stopping=True
+    )
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+def generate_rag_caption(image):
+    """
+    Generate a RAG-based caption for an image.
+    Input: PIL Image or image path (str).
+    Output: Caption (str).
+    """
+    embedding = extract_image_features(image)
+    if embedding is None:
+        return "Failed to process image."
+    retrieved = retrieve_similar_captions(embedding, k=5)
+    if not retrieved:
+        return "No similar captions found."
+    return generate_caption_from_retrieved(retrieved)
+
+def predict(image):
+    """
+    API-compatible function for inference.
+    Input: PIL Image or image file path.
+    Output: Dictionary with caption.
+    """
+    caption = generate_rag_caption(image)
+    return {"caption": caption}

requirements.txt

@@ -0,0 +1,11 @@
+
+transformers>=4.30.0 
+sentence-transformers>=2.2.0 
+faiss-cpu>=1.7.0 
+torch>=2.0.0 
+torchvision>=0.15.0 
+torchaudio>=2.0.0 
+pillow>=9.0.0 
+spacy>=3.5.0 
+langchain>=0.0.200 
+huggingface_hub>=0.15.0