fixing incorrect references

api.py

@@ -1,11 +1,8 @@
-# api.py
 from fastapi import FastAPI, UploadFile, File, Form
 from fastapi.middleware.cors import CORSMiddleware
 from PIL import Image
 from typing import List
 from pydantic import BaseModel
-from models.resnet_lstm_attention.loader import load_captioning_model
-from models.resnet_lstm_attention.cap_mod_defs import EncoderCNN
 
 from model_registry import get_model
 from models.resnet_lstm_attention.schemas import CaptionResult, ImageResult, TextQuery
@@ -23,7 +20,6 @@ class InferenceRequest(BaseModel):
     model_name: str
     top_k: int = 5
 
-#@app.post("/caption", response_model=CaptionResult)
 @app.post("/caption")
 async def caption_image(model_name: str = Form(...), file: UploadFile = File(...)):
     image = Image.open(file.file).convert("RGB")
@@ -31,7 +27,6 @@ async def caption_image(model_name: str = Form(...), file: UploadFile = File(...
     caption = model.generate_caption(image)
     return {"caption": caption}
 
-#@app.post("/search/text2img", response_model=List[ImageResult])
 @app.post("/search/text2img")
 async def text_to_image(model_name: str = Form(...), query: str = Form(...), top_k: int = Form(5)):
     model = get_model(model_name)
@@ -45,7 +40,6 @@ async def image_to_text(model_name: str = Form(...), file: UploadFile = File(...
     results = model.image_to_text(image, top_k)
     return results
 
-#@app.post("/search/img2img", response_model=List[ImageResult])
 @app.post("/search/img2img")
 async def image_to_image(model_name: str = Form(...), file: UploadFile = File(...), top_k: int = Form(5)):
     image = Image.open(file.file).convert("RGB")
@@ -61,4 +55,71 @@ async def text_to_text(model_name: str = Form(...), query: str = Form(...), top_
 
 @app.get("/health")
 def health_check():
-    return {"status": "healthy"}
+    return {"status": "healthy"}
+
+
+
+# # api.py
+# from fastapi import FastAPI, UploadFile, File, Form
+# from fastapi.middleware.cors import CORSMiddleware
+# from PIL import Image
+# from typing import List
+# from pydantic import BaseModel
+# from models.resnet_lstm_attention.loader import load_captioning_model
+# from models.resnet_lstm_attention.cap_mod_defs import EncoderCNN
+
+# from model_registry import get_model
+# from models.resnet_lstm_attention.schemas import CaptionResult, ImageResult, TextQuery
+
+# app = FastAPI(title="Multimodal Retrieval & Captioning API")
+
+# app.add_middleware(
+#     CORSMiddleware,
+#     allow_origins=["*"],
+#     allow_methods=["*"],
+#     allow_headers=["*"],
+# )
+
+# class InferenceRequest(BaseModel):
+#     model_name: str
+#     top_k: int = 5
+
+# #@app.post("/caption", response_model=CaptionResult)
+# @app.post("/caption")
+# async def caption_image(model_name: str = Form(...), file: UploadFile = File(...)):
+#     image = Image.open(file.file).convert("RGB")
+#     model = get_model(model_name)
+#     caption = model.generate_caption(image)
+#     return {"caption": caption}
+
+# #@app.post("/search/text2img", response_model=List[ImageResult])
+# @app.post("/search/text2img")
+# async def text_to_image(model_name: str = Form(...), query: str = Form(...), top_k: int = Form(5)):
+#     model = get_model(model_name)
+#     results = model.text_to_image(query, top_k)
+#     return results
+
+# @app.post("/search/img2text")
+# async def image_to_text(model_name: str = Form(...), file: UploadFile = File(...), top_k: int = Form(5)):
+#     image = Image.open(file.file).convert("RGB")
+#     model = get_model(model_name)
+#     results = model.image_to_text(image, top_k)
+#     return results
+
+# #@app.post("/search/img2img", response_model=List[ImageResult])
+# @app.post("/search/img2img")
+# async def image_to_image(model_name: str = Form(...), file: UploadFile = File(...), top_k: int = Form(5)):
+#     image = Image.open(file.file).convert("RGB")
+#     model = get_model(model_name)
+#     results = model.image_to_image(image, top_k)
+#     return results
+
+# @app.post("/search/text2text")
+# async def text_to_text(model_name: str = Form(...), query: str = Form(...), top_k: int = Form(5)):
+#     model = get_model(model_name)
+#     results = model.text_to_text(query, top_k)
+#     return results
+
+# @app.get("/health")
+# def health_check():
+#     return {"status": "healthy"}

app.py

@@ -1,4 +1,3 @@
-# app.py
 import streamlit as st
 import requests
 import subprocess
@@ -9,10 +8,6 @@ import base64  # For displaying retrieved images if needed
 import socket
 
 # Start FastAPI server in background
-# subprocess.Popen(["uvicorn", "api:app", "--host", "0.0.0.0", "--port", "8001"])
-# time.sleep(2)  # Wait for server to start
-
-# Check if port is free
 def is_port_free(port):
     with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
         return s.connect_ex(('localhost', port)) != 0
@@ -57,17 +52,6 @@ with tab_caption:
         else:
             st.error("Error: " + resp.text)
 
-# with tab_text2img:
-#     if text_input and st.button("Search Images"):
-#         data = {"model_name": model_name, "query": text_input, "top_k": top_k}
-#         resp = requests.post(f"{API_BASE}/search/text2img", data=data)
-#         if resp.status_code == 200:
-#             results = resp.json()
-#             for res in results:
-#                 st.image(res["image_path"], caption=f"Score: {res['score']:.3f}")
-#         else:
-#             st.error("Error: " + resp.text)
-
 with tab_text2img:
     if text_input and st.button("Search Images"):
         data = {"model_name": model_name, "query": text_input, "top_k": top_k}
@@ -82,10 +66,8 @@ with tab_text2img:
                         if res["image"] is not None:
                             st.image(res["image"], width=200)
                             st.caption(f"Score: {res['score']:.3f}")
-                            if "caption" in res:  # if you add caption to results later
-                                st.write(res["caption"])
                         else:
-                            st.caption(f"Score: {res['score']:.3f} (Image not found)")
+                            st.caption(f"Score: {res['score']:.3f} (Image not available)")
             else:
                 st.info("No results found.")
         else:
@@ -97,28 +79,21 @@ with tab_img2text:
         data = {"model_name": model_name, "top_k": top_k}
         resp = requests.post(f"{API_BASE}/search/img2text", files=files, data=data)
         if resp.status_code == 200:
-            st.write("Retrieved Texts:", resp.json())
+            results = resp.json()
+            if results:
+                st.subheader("Retrieved Texts:")
+                for idx, caption in enumerate(results, 1):
+                    st.markdown(f"**{idx}.** {caption}")
+            else:
+                st.info("No results found.")
         else:
-            st.error("Error: " + resp.text)
-
-# with tab_img2img:
-#     if image_file and st.button("Retrieve Similar Images"):
-#         files = {"file": image_file.getvalue()}
-#         data = {"model_name": model_name, "top_k": top_k}
-#         resp = requests.post(f"{API_BASE}/search/img2img", files=files, data=data)
-#         if resp.status_code == 200:
-#             results = resp.json()
-#             for res in results:
-#                 st.image(res["image_path"], caption=f"Score: {res['score']:.3f}")
-#         else:
-#             st.error("Error: " + resp.text)
+            st.error(f"Error: {resp.status_code} - {resp.text}")
 
 with tab_img2img:
     if image_file and st.button("Retrieve Similar Images"):
         files = {"file": image_file.getvalue()}
         data = {"model_name": model_name, "top_k": top_k}
         resp = requests.post(f"{API_BASE}/search/img2img", files=files, data=data)
-        
         if resp.status_code == 200:
             results = resp.json()
             if results:
@@ -126,18 +101,15 @@ with tab_img2img:
                 cols = st.columns(3)
                 for idx, res in enumerate(results):
                     with cols[idx % 3]:
-                        if "image" in res and res["image"] is not None:
-                            st.image(
-                                res["image"],
-                                width=200,        # recommended instead of use_column_width
-                                caption=f"Score: {res['score']:.3f}"
-                            )
+                        if res["image"] is not None:
+                            st.image(res["image"], width=200)
+                            st.caption(f"Score: {res['score']:.3f}")
                         else:
                             st.caption(f"Score: {res['score']:.3f} (Image not available)")
             else:
-                st.info("No similar images found in the dataset.")
+                st.info("No results found.")
         else:
-            st.error(f"Backend error: {resp.status_code} - {resp.text}")
+            st.error(f"Error: {resp.status_code} - {resp.text}")
 
 with tab_text2text:
     text_input_tt = st.text_input("Enter text to find similar captions", 
@@ -156,4 +128,167 @@ with tab_text2text:
             else:
                 st.info("No similar captions found.")
         else:
-            st.error(f"Error: {resp.status_code} - {resp.text}")
+            st.error(f"Error: {resp.status_code} - {resp.text}")
+
+
+# Old Code
+
+# # app.py
+# import streamlit as st
+# import requests
+# import subprocess
+# import time
+# from PIL import Image
+# import io
+# import base64  # For displaying retrieved images if needed
+# import socket
+
+# # Start FastAPI server in background
+# # subprocess.Popen(["uvicorn", "api:app", "--host", "0.0.0.0", "--port", "8001"])
+# # time.sleep(2)  # Wait for server to start
+
+# # Check if port is free
+# def is_port_free(port):
+#     with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+#         return s.connect_ex(('localhost', port)) != 0
+
+# if is_port_free(8001):
+#     subprocess.Popen(["uvicorn", "api:app", "--host", "0.0.0.0", "--port", "8001"])
+# else:
+#     print("Port 8001 in use - skipping backend startup")
+# time.sleep(5)  # longer wait
+
+# API_BASE = "http://localhost:8001"
+
+# st.set_page_config(page_title="Multimodal Retrieval & Captioning", layout="wide")
+
+# st.title("Multimodal Retrieval & Captioning System")
+
+# # Model selection (add more later)
+# model_name = st.sidebar.selectbox("Select Model", ["resnet_lstm_attention", "vit_lstm_attention", "vit_transformer"], index=0)
+
+# # Common inputs
+# input_method = st.sidebar.radio("Image Input", ["Upload", "Camera"])
+# image_file = st.file_uploader("Upload Image", type=["jpg", "jpeg", "png"]) if input_method == "Upload" else st.camera_input("Capture Image")
+# text_input = st.text_input("Text Input")
+# top_k = st.sidebar.slider("Top K", 1, 10, 5)
+
+# # Tabs for tasks
+# tab_caption, tab_text2img, tab_img2text, tab_img2img, tab_text2text = st.tabs([
+#     "Image → Caption",
+#     "Text → Image",
+#     "Image → Text",
+#     "Image → Image",
+#     "Text → Text"
+# ])
+
+# with tab_caption:
+#     if image_file and st.button("Generate Caption"):
+#         files = {"file": image_file.getvalue()}
+#         data = {"model_name": model_name}
+#         resp = requests.post(f"{API_BASE}/caption", files=files, data=data)
+#         if resp.status_code == 200:
+#             st.write("Caption:", resp.json()["caption"])
+#         else:
+#             st.error("Error: " + resp.text)
+
+# # with tab_text2img:
+# #     if text_input and st.button("Search Images"):
+# #         data = {"model_name": model_name, "query": text_input, "top_k": top_k}
+# #         resp = requests.post(f"{API_BASE}/search/text2img", data=data)
+# #         if resp.status_code == 200:
+# #             results = resp.json()
+# #             for res in results:
+# #                 st.image(res["image_path"], caption=f"Score: {res['score']:.3f}")
+# #         else:
+# #             st.error("Error: " + resp.text)
+
+# with tab_text2img:
+#     if text_input and st.button("Search Images"):
+#         data = {"model_name": model_name, "query": text_input, "top_k": top_k}
+#         resp = requests.post(f"{API_BASE}/search/text2img", data=data)
+#         if resp.status_code == 200:
+#             results = resp.json()
+#             if results:
+#                 st.subheader("Retrieved Images")
+#                 cols = st.columns(3)
+#                 for idx, res in enumerate(results):
+#                     with cols[idx % 3]:
+#                         if res["image"] is not None:
+#                             st.image(res["image"], width=200)
+#                             st.caption(f"Score: {res['score']:.3f}")
+#                             if "caption" in res:  # if you add caption to results later
+#                                 st.write(res["caption"])
+#                         else:
+#                             st.caption(f"Score: {res['score']:.3f} (Image not found)")
+#             else:
+#                 st.info("No results found.")
+#         else:
+#             st.error(f"Error: {resp.status_code} - {resp.text}")
+
+# with tab_img2text:
+#     if image_file and st.button("Retrieve Text"):
+#         files = {"file": image_file.getvalue()}
+#         data = {"model_name": model_name, "top_k": top_k}
+#         resp = requests.post(f"{API_BASE}/search/img2text", files=files, data=data)
+#         if resp.status_code == 200:
+#             st.write("Retrieved Texts:", resp.json())
+#         else:
+#             st.error("Error: " + resp.text)
+
+# # with tab_img2img:
+# #     if image_file and st.button("Retrieve Similar Images"):
+# #         files = {"file": image_file.getvalue()}
+# #         data = {"model_name": model_name, "top_k": top_k}
+# #         resp = requests.post(f"{API_BASE}/search/img2img", files=files, data=data)
+# #         if resp.status_code == 200:
+# #             results = resp.json()
+# #             for res in results:
+# #                 st.image(res["image_path"], caption=f"Score: {res['score']:.3f}")
+# #         else:
+# #             st.error("Error: " + resp.text)
+
+# with tab_img2img:
+#     if image_file and st.button("Retrieve Similar Images"):
+#         files = {"file": image_file.getvalue()}
+#         data = {"model_name": model_name, "top_k": top_k}
+#         resp = requests.post(f"{API_BASE}/search/img2img", files=files, data=data)
+        
+#         if resp.status_code == 200:
+#             results = resp.json()
+#             if results:
+#                 st.subheader("Retrieved Similar Images")
+#                 cols = st.columns(3)
+#                 for idx, res in enumerate(results):
+#                     with cols[idx % 3]:
+#                         if "image" in res and res["image"] is not None:
+#                             st.image(
+#                                 res["image"],
+#                                 width=200,        # recommended instead of use_column_width
+#                                 caption=f"Score: {res['score']:.3f}"
+#                             )
+#                         else:
+#                             st.caption(f"Score: {res['score']:.3f} (Image not available)")
+#             else:
+#                 st.info("No similar images found in the dataset.")
+#         else:
+#             st.error(f"Backend error: {resp.status_code} - {resp.text}")
+
+# with tab_text2text:
+#     text_input_tt = st.text_input("Enter text to find similar captions", 
+#                                  placeholder="A child playing with water in the garden")
+    
+#     if text_input_tt and st.button("Search Similar Captions"):
+#         data = {"model_name": model_name, "query": text_input_tt, "top_k": top_k}
+#         resp = requests.post(f"{API_BASE}/search/text2text", data=data)
+        
+#         if resp.status_code == 200:
+#             results = resp.json()
+#             if results:
+#                 st.subheader("Top similar captions:")
+#                 for idx, res in enumerate(results, 1):
+#                     st.markdown(f"**{idx}.** {res['caption']}  \nScore: `{res['score']:.4f}`")
+#             else:
+#                 st.info("No similar captions found.")
+#         else:
+#             st.error(f"Error: {resp.status_code} - {resp.text}")

models/resnet_lstm_attention/model.py

@@ -19,7 +19,6 @@ class ResNetLSTMAttentionModel(UnifiedModelInterface):
         self.retrieval_service = None
         self.device = torch.device("cpu")
         self.dataset = None
-        #self.model_repo = "skodan/resnet-lstm-attention-weights"
 
     def load(self) -> None:
         if self.caption_bundle is not None and self.retrieval_service is not None:
@@ -28,14 +27,14 @@ class ResNetLSTMAttentionModel(UnifiedModelInterface):
         MODEL_REPO = "skodan/resnet-lstm-attention-weights"
 
         files_to_download = [
-                "caption_model.pth",
-                "flickr8k_retrieval_model.pth",
-                "image_embeddings.faiss",
-                "text_embeddings.faiss",
-                "image_id_map.pkl",
-                "text_id_map.pkl",
-                "vocab.pkl"
-            ]
+            "caption_model.pth",
+            "flickr8k_retrieval_model.pth",
+            "image_embeddings.faiss",
+            "text_embeddings.faiss",
+            "image_id_map.pkl",
+            "text_id_map.pkl",
+            "vocab.pkl"  # only if large; otherwise commit it
+        ]
 
         downloaded_paths = {}
         for fname in files_to_download:
@@ -43,33 +42,33 @@ class ResNetLSTMAttentionModel(UnifiedModelInterface):
                 path = hf_hub_download(
                     repo_id=MODEL_REPO,
                     filename=fname,
-                    repo_type="model",
+                    repo_type="model"
                 )
                 downloaded_paths[fname] = path
             except Exception as e:
                 raise RuntimeError(f"Failed to download {fname} from {MODEL_REPO}: {e}")
 
-        # Download large files from HF Hub
         caption_pth = downloaded_paths["caption_model.pth"]
         retrieval_pth = downloaded_paths["flickr8k_retrieval_model.pth"]
         image_index_faiss = downloaded_paths["image_embeddings.faiss"]
         text_index_faiss = downloaded_paths["text_embeddings.faiss"]
         image_map_pkl = downloaded_paths["image_id_map.pkl"]
         text_map_pkl = downloaded_paths["text_id_map.pkl"]
-        vocab_pkl = downloaded_paths["vocab.pkl"]  
+        vocab_pkl = downloaded_paths["vocab.pkl"]
 
-        # Load configs (assume small, committed to repo)
         base_dir = os.path.dirname(os.path.dirname(os.path.dirname(__file__)))  # go up to project root
         config_path = os.path.join(base_dir, "configs", "caption_config.json")
         preprocess_cfg_path = os.path.join(base_dir, "configs", "preprocess_config.json")
 
+        if not os.path.exists(config_path):
+            raise FileNotFoundError(f"Config not found: {config_path}")
+
         with open(config_path, "r") as f:
             caption_config = json.load(f)
 
         with open(preprocess_cfg_path, "r") as f:
             preprocess_cfg = json.load(f)
 
-        # Load captioning
         self.caption_bundle = load_captioning_model(
             model_path=caption_pth,
             vocab_path=vocab_pkl,
@@ -77,7 +76,6 @@ class ResNetLSTMAttentionModel(UnifiedModelInterface):
             device=self.device
         )
 
-        # Load retrieval
         clip_model = load_clip_model(
             model_path=retrieval_pth,
             vocab=self.caption_bundle["vocab"],
@@ -94,19 +92,18 @@ class ResNetLSTMAttentionModel(UnifiedModelInterface):
         )
 
         if self.dataset is None:
-            print("Loading Flickr8k test split from Hugging Face...")
+            print("Loading full Flickr8k dataset from Hugging Face...")
             ds = load_dataset("jxie/flickr8k")
-            self.dataset = concatenate_datasets([
-                ds["train"],
-                ds["validation"],
-                ds["test"]
-            ])
+            self.dataset = concatenate_datasets([ds["train"], ds["validation"], ds["test"]])
             print(f"Loaded {len(self.dataset)} images/captions from full dataset.")
 
         print("Model components loaded successfully.")
 
     @torch.no_grad()
     def generate_caption(self, image: Image.Image) -> str:
+        if self.caption_bundle is None:
+            raise RuntimeError("Model not loaded. Call load() first.")
+
         encoder = self.caption_bundle["encoder"]
         decoder = self.caption_bundle["decoder"]
         vocab = self.caption_bundle["vocab"]
@@ -115,6 +112,7 @@ class ResNetLSTMAttentionModel(UnifiedModelInterface):
         transform = self.caption_bundle["transform"]
 
         image_tensor = transform(image).unsqueeze(0).to(self.device)
+
         features = encoder(image_tensor)
         tokens = decoder.generate(
             features,
@@ -124,125 +122,303 @@ class ResNetLSTMAttentionModel(UnifiedModelInterface):
         )
         return " ".join(tokens)
 
-    # def text_to_image(self, text: str, top_k: int = 5) -> List[Dict[str, Any]]:
-    #     return self.retrieval_service.text_to_image(text, top_k)
-
     def text_to_image(self, text: str, top_k: int = 5) -> List[Dict[str, Any]]:
         raw_results = self.retrieval_service.text_to_image(text, top_k)
-        return self._format_retrieval_results(raw_results)
-
-    # def text_to_image(self, text: str, top_k: int = 5) -> List[Dict[str, Any]]:
-    #     results = self.retrieval_service.text_to_image(text, top_k)
         
-    #     formatted_results = []
-    #     for res in results:
-    #         img_id_str = str(res["image_path"])  # this is likely the ID or filename without .jpg
-    #         img_file = f"{img_id_str}.jpg"
-    #         #img_file = f"{img_id}.jpg" if not img_id.endswith('.jpg') else img_id
-            
-    #         full_path = os.path.join("flickr8k_images", img_file)
+        formatted = []
+        for res in raw_results:
+            idx = int(res["image_path"])
             
-    #         # Only include if the file actually exists in the demo folder
-    #         if os.path.exists(full_path):
-    #             formatted_results.append({
-    #                 "image_path": full_path,
-    #                 "score": res["score"]
-    #             })
-    #         # Optional: skip or use placeholder if missing
-    #         else:
-    #             formatted_results.append({
-    #                 "image_path": "https://via.placeholder.com/300?text=Not+in+demo",
-    #                 "score": res["score"]
-    #             })
+            try:
+                pil_img = self.dataset[idx]["image"]
+                formatted.append({
+                    "image": pil_img,
+                    "score": float(res["score"])
+                })
+            except (IndexError, KeyError):
+                formatted.append({
+                    "image": None,
+                    "score": float(res["score"])
+                })
         
-    #     return formatted_results
-
-    # def image_to_text(self, image: Image.Image, top_k: int = 5) -> List[str]:
-    #     return self.retrieval_service.image_to_text(image, top_k)
+        return formatted
 
     def image_to_text(self, image: Image.Image, top_k: int = 5) -> List[str]:
         return self.retrieval_service.image_to_text(image, top_k)
 
-    # def image_to_text(self, image: Image.Image, top_k: int = 5) -> List[str]:
-    #     results = self.retrieval_service.image_to_text(image, top_k)  # assuming this returns list of dicts
-    
-    #     formatted_results = []
-    #     for res in results:
-    #         img_id = res["image_path"]  # same as above
-    #         img_file = f"{img_id}.jpg"# if not img_id.endswith('.jpg') else img_id
-            
-    #         full_path = os.path.join("flickr8k_images", img_file)
-            
-    #         if os.path.exists(full_path):
-    #             formatted_results.append({
-    #                 "image_path": full_path,
-    #                 "score": res["score"]
-    #             })
-    #         else:
-    #             # Optional fallback so UI doesn't crash
-    #             formatted.append({
-    #                 "image_path": "https://via.placeholder.com/300x200?text=Not+in+demo",
-    #                 "score": float(res["score"])
-    #             })
-        
-    #     return formatted_results
-
     def image_to_image(self, image: Image.Image, top_k: int = 5) -> List[Dict[str, Any]]:
-        raw_results = self.retrieval_service.image_to_image(image, top_k)  # new call
-        return self._format_retrieval_results(raw_results)
-
-
-    def _format_retrieval_results(self, raw_results: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+        raw_results = self.retrieval_service.image_to_image(image, top_k)
+        
         formatted = []
         for res in raw_results:
-            img_id = res["image_path"]  # this is integer ID
-            img_filename = f"{img_id}.jpg"  # always append .jpg, no .endswith needed
-            full_path = os.path.join("flickr8k_images", img_filename)
+            idx = int(res["image_path"])
             
-            if os.path.exists(full_path):
+            try:
+                pil_img = self.dataset[idx]["image"]
                 formatted.append({
-                    "image_path": full_path,
-                    "score": res["score"]
+                    "image": pil_img,
+                    "score": float(res["score"])
                 })
-            else:
+            except (IndexError, KeyError):
                 formatted.append({
-                    "image_path": "https://via.placeholder.com/300?text=Not+in+demo",
-                    "score": res["score"]
+                    "image": None,
+                    "score": float(res["score"])
                 })
+        
         return formatted
 
+    def text_to_text(self, text: str, top_k: int = 5) -> List[Dict[str, Any]]:
+        return self.retrieval_service.text_to_text(text, top_k)
+
+
+
+# Old code
+
+# import os
+# import json
+# import torch
+# from huggingface_hub import hf_hub_download
+# from PIL import Image
+# import numpy as np
+# from typing import List, Dict, Any
+# from datasets import load_dataset, concatenate_datasets
+
+# from models.resnet_lstm_attention.loader import load_captioning_model
+# from models.resnet_lstm_attention.retrieval import RetrievalService
+# from models.resnet_lstm_attention.clip_loader import load_clip_model
+# from models.resnet_lstm_attention.captioning import CaptioningService  # Not directly used, but for reference
+# from utils.interfaces import UnifiedModelInterface  # Adjust path if needed
+
+# class ResNetLSTMAttentionModel(UnifiedModelInterface):
+#     def __init__(self):
+#         self.caption_bundle = None
+#         self.retrieval_service = None
+#         self.device = torch.device("cpu")
+#         self.dataset = None
+#         #self.model_repo = "skodan/resnet-lstm-attention-weights"
+
+#     def load(self) -> None:
+#         if self.caption_bundle is not None and self.retrieval_service is not None:
+#             return
+        
+#         MODEL_REPO = "skodan/resnet-lstm-attention-weights"
 
-    # def image_to_image(self, image: Image.Image, top_k: int = 5) -> List[Dict[str, Any]]:
-    #     # image_tensor = self.retrieval_service.image_transform(image).unsqueeze(0).to(self.device)
-    #     # with torch.no_grad():
-    #     #     emb = self.retrieval_service.clip_model.encode_image(image_tensor).cpu().numpy()
-    #     # emb = self.retrieval_service._normalize(emb)
-    #     # scores, idxs = self.retrieval_service.image_index.search(emb, top_k)
-    #     # return [
-    #     #     {"image_path": self.retrieval_service.image_id_map[i], "score": float(scores[0][j])}
-    #     #     for j, i in enumerate(idxs[0])
-    #     # ]
-    #     raw_results = self.retrieval_service.image_to_image(image, top_k)
+#         files_to_download = [
+#                 "caption_model.pth",
+#                 "flickr8k_retrieval_model.pth",
+#                 "image_embeddings.faiss",
+#                 "text_embeddings.faiss",
+#                 "image_id_map.pkl",
+#                 "text_id_map.pkl",
+#                 "vocab.pkl"
+#             ]
+
+#         downloaded_paths = {}
+#         for fname in files_to_download:
+#             try:
+#                 path = hf_hub_download(
+#                     repo_id=MODEL_REPO,
+#                     filename=fname,
+#                     repo_type="model",
+#                 )
+#                 downloaded_paths[fname] = path
+#             except Exception as e:
+#                 raise RuntimeError(f"Failed to download {fname} from {MODEL_REPO}: {e}")
+
+#         # Download large files from HF Hub
+#         caption_pth = downloaded_paths["caption_model.pth"]
+#         retrieval_pth = downloaded_paths["flickr8k_retrieval_model.pth"]
+#         image_index_faiss = downloaded_paths["image_embeddings.faiss"]
+#         text_index_faiss = downloaded_paths["text_embeddings.faiss"]
+#         image_map_pkl = downloaded_paths["image_id_map.pkl"]
+#         text_map_pkl = downloaded_paths["text_id_map.pkl"]
+#         vocab_pkl = downloaded_paths["vocab.pkl"]  
+
+#         # Load configs (assume small, committed to repo)
+#         base_dir = os.path.dirname(os.path.dirname(os.path.dirname(__file__)))  # go up to project root
+#         config_path = os.path.join(base_dir, "configs", "caption_config.json")
+#         preprocess_cfg_path = os.path.join(base_dir, "configs", "preprocess_config.json")
+
+#         with open(config_path, "r") as f:
+#             caption_config = json.load(f)
+
+#         with open(preprocess_cfg_path, "r") as f:
+#             preprocess_cfg = json.load(f)
+
+#         # Load captioning
+#         self.caption_bundle = load_captioning_model(
+#             model_path=caption_pth,
+#             vocab_path=vocab_pkl,
+#             config_path=config_path,
+#             device=self.device
+#         )
+
+#         # Load retrieval
+#         clip_model = load_clip_model(
+#             model_path=retrieval_pth,
+#             vocab=self.caption_bundle["vocab"],
+#             device=self.device
+#         )
+
+#         self.retrieval_service = RetrievalService(
+#             clip_model=clip_model,
+#             image_index_path=image_index_faiss,
+#             text_index_path=text_index_faiss,
+#             image_map_path=image_map_pkl,
+#             text_map_path=text_map_pkl,
+#             preprocess=preprocess_cfg
+#         )
+
+#         if self.dataset is None:
+#             print("Loading Flickr8k test split from Hugging Face...")
+#             ds = load_dataset("jxie/flickr8k")
+#             self.dataset = concatenate_datasets([
+#                 ds["train"],
+#                 ds["validation"],
+#                 ds["test"]
+#             ])
+#             print(f"Loaded {len(self.dataset)} images/captions from full dataset.")
+
+#         print("Model components loaded successfully.")
+
+#     @torch.no_grad()
+#     def generate_caption(self, image: Image.Image) -> str:
+#         encoder = self.caption_bundle["encoder"]
+#         decoder = self.caption_bundle["decoder"]
+#         vocab = self.caption_bundle["vocab"]
+#         inv_vocab = self.caption_bundle["inv_vocab"]
+#         max_len = self.caption_bundle["max_len"]
+#         transform = self.caption_bundle["transform"]
+
+#         image_tensor = transform(image).unsqueeze(0).to(self.device)
+#         features = encoder(image_tensor)
+#         tokens = decoder.generate(
+#             features,
+#             vocab=vocab,
+#             inv_vocab=inv_vocab,
+#             max_len=max_len
+#         )
+#         return " ".join(tokens)
+
+#     # def text_to_image(self, text: str, top_k: int = 5) -> List[Dict[str, Any]]:
+#     #     return self.retrieval_service.text_to_image(text, top_k)
+
+#     def text_to_image(self, text: str, top_k: int = 5) -> List[Dict[str, Any]]:
+#         raw_results = self.retrieval_service.text_to_image(text, top_k)
+#         return self._format_retrieval_results(raw_results)
+
+#     # def text_to_image(self, text: str, top_k: int = 5) -> List[Dict[str, Any]]:
+#     #     results = self.retrieval_service.text_to_image(text, top_k)
+        
+#     #     formatted_results = []
+#     #     for res in results:
+#     #         img_id_str = str(res["image_path"])  # this is likely the ID or filename without .jpg
+#     #         img_file = f"{img_id_str}.jpg"
+#     #         #img_file = f"{img_id}.jpg" if not img_id.endswith('.jpg') else img_id
+            
+#     #         full_path = os.path.join("flickr8k_images", img_file)
+            
+#     #         # Only include if the file actually exists in the demo folder
+#     #         if os.path.exists(full_path):
+#     #             formatted_results.append({
+#     #                 "image_path": full_path,
+#     #                 "score": res["score"]
+#     #             })
+#     #         # Optional: skip or use placeholder if missing
+#     #         else:
+#     #             formatted_results.append({
+#     #                 "image_path": "https://via.placeholder.com/300?text=Not+in+demo",
+#     #                 "score": res["score"]
+#     #             })
+        
+#     #     return formatted_results
+
+#     # def image_to_text(self, image: Image.Image, top_k: int = 5) -> List[str]:
+#     #     return self.retrieval_service.image_to_text(image, top_k)
+
+#     def image_to_text(self, image: Image.Image, top_k: int = 5) -> List[str]:
+#         return self.retrieval_service.image_to_text(image, top_k)
+
+#     # def image_to_text(self, image: Image.Image, top_k: int = 5) -> List[str]:
+#     #     results = self.retrieval_service.image_to_text(image, top_k)  # assuming this returns list of dicts
     
-    #     formatted = []
-    #     for res in raw_results:
-    #         img_id_str = str(res["image_path"])
+#     #     formatted_results = []
+#     #     for res in results:
+#     #         img_id = res["image_path"]  # same as above
+#     #         img_file = f"{img_id}.jpg"# if not img_id.endswith('.jpg') else img_id
             
-    #         img_filename = f"{img_id_str}.jpg"
-    #         full_path = os.path.join("flickr8k_images", img_filename)
+#     #         full_path = os.path.join("flickr8k_images", img_file)
             
-    #         if os.path.exists(full_path):
-    #             formatted.append({
-    #                 "image_path": full_path,
-    #                 "score": float(res["score"])
-    #             })
-    #         else:
-    #             formatted.append({
-    #                 "image_path": "https://via.placeholder.com/300x200?text=Not+in+demo",
-    #                 "score": float(res["score"])
-    #             })
+#     #         if os.path.exists(full_path):
+#     #             formatted_results.append({
+#     #                 "image_path": full_path,
+#     #                 "score": res["score"]
+#     #             })
+#     #         else:
+#     #             # Optional fallback so UI doesn't crash
+#     #             formatted.append({
+#     #                 "image_path": "https://via.placeholder.com/300x200?text=Not+in+demo",
+#     #                 "score": float(res["score"])
+#     #             })
         
-    #     return formatted
+#     #     return formatted_results
 
-    def text_to_text(self, text: str, top_k: int = 5) -> List[Dict[str, Any]]:
-        return self.retrieval_service.text_to_text(text, top_k)
+#     def image_to_image(self, image: Image.Image, top_k: int = 5) -> List[Dict[str, Any]]:
+#         raw_results = self.retrieval_service.image_to_image(image, top_k)  # new call
+#         return self._format_retrieval_results(raw_results)
+
+
+#     def _format_retrieval_results(self, raw_results: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+#         formatted = []
+#         for res in raw_results:
+#             img_id = res["image_path"]  # this is integer ID
+#             img_filename = f"{img_id}.jpg"  # always append .jpg, no .endswith needed
+#             full_path = os.path.join("flickr8k_images", img_filename)
+            
+#             if os.path.exists(full_path):
+#                 formatted.append({
+#                     "image_path": full_path,
+#                     "score": res["score"]
+#                 })
+#             else:
+#                 formatted.append({
+#                     "image_path": "https://via.placeholder.com/300?text=Not+in+demo",
+#                     "score": res["score"]
+#                 })
+#         return formatted
+
+
+#     # def image_to_image(self, image: Image.Image, top_k: int = 5) -> List[Dict[str, Any]]:
+#     #     # image_tensor = self.retrieval_service.image_transform(image).unsqueeze(0).to(self.device)
+#     #     # with torch.no_grad():
+#     #     #     emb = self.retrieval_service.clip_model.encode_image(image_tensor).cpu().numpy()
+#     #     # emb = self.retrieval_service._normalize(emb)
+#     #     # scores, idxs = self.retrieval_service.image_index.search(emb, top_k)
+#     #     # return [
+#     #     #     {"image_path": self.retrieval_service.image_id_map[i], "score": float(scores[0][j])}
+#     #     #     for j, i in enumerate(idxs[0])
+#     #     # ]
+#     #     raw_results = self.retrieval_service.image_to_image(image, top_k)
+    
+#     #     formatted = []
+#     #     for res in raw_results:
+#     #         img_id_str = str(res["image_path"])
+            
+#     #         img_filename = f"{img_id_str}.jpg"
+#     #         full_path = os.path.join("flickr8k_images", img_filename)
+            
+#     #         if os.path.exists(full_path):
+#     #             formatted.append({
+#     #                 "image_path": full_path,
+#     #                 "score": float(res["score"])
+#     #             })
+#     #         else:
+#     #             formatted.append({
+#     #                 "image_path": "https://via.placeholder.com/300x200?text=Not+in+demo",
+#     #                 "score": float(res["score"])
+#     #             })
+        
+#     #     return formatted
+
+#     def text_to_text(self, text: str, top_k: int = 5) -> List[Dict[str, Any]]:
+#         return self.retrieval_service.text_to_text(text, top_k)

models/resnet_lstm_attention/retrieval.py

@@ -2,7 +2,6 @@ import faiss
 import pickle
 import torch
 import numpy as np
-import os
 from PIL import Image
 from torchvision import transforms
 from typing import List, Dict, Any
@@ -34,70 +33,23 @@ class RetrievalService:
 
     def _normalize(self, x):
         return x / np.linalg.norm(x, axis=1, keepdims=True)
-    
 
     def text_to_image(self, text: str, top_k: int = 5) -> List[Dict[str, Any]]:
-        raw_results = self.retrieval_service.text_to_image(text, top_k)
-        
-        formatted = []
-        for res in raw_results:
-            idx = int(res["image_path"])  # the FAISS index (integer)
-            
-            try:
-                pil_img = self.dataset[idx]["image"]  # directly get PIL.Image
-                formatted.append({
-                    "image": pil_img,  # ← pass PIL.Image to UI
-                    "score": float(res["score"])
-                })
-            except (IndexError, KeyError):
-                formatted.append({
-                    "image": None,
-                    "score": float(res["score"])
-                })
-        
-        return formatted
-
-
-    # def text_to_image(self, text: str, top_k: int = 5) -> List[Dict[str, Any]]:
-    #     raw_results = self.retrieval_service.text_to_image(text, top_k)
-        
-    #     formatted = []
-    #     for res in raw_results:
-    #         img_id = res["image_path"]  # int or str
-    #         img_id_str = str(img_id)
-    #         img_filename = f"{img_id_str}.jpg"  # always append .jpg, no .endswith
-    #         full_path = os.path.join("flickr8k_images", img_filename)
-            
-    #         if os.path.exists(full_path):
-    #             formatted.append({
-    #                 "image_path": full_path,
-    #                 "score": float(res["score"])
-    #             })
-    #         else:
-    #             formatted.append({
-    #                 "image_path": "https://via.placeholder.com/300?text=Not+in+demo",
-    #                 "score": float(res["score"])
-    #             })
-        
-    #     return formatted
-
-
-    # def text_to_image(self, text, top_k=5):
-    #     with torch.no_grad():
-    #         emb = self.clip_model.encode_text(text).cpu().numpy()
-    #     emb = self._normalize(emb)
+        with torch.no_grad():
+            emb = self.clip_model.encode_text(text).cpu().numpy()
+        emb = self._normalize(emb)
 
-    #     scores, idxs = self.image_index.search(emb, top_k)
-    #     return [
-    #         {
-    #             "image_path": self.image_id_map[i],
-    #             "score": float(scores[0][j])
-    #         }
-    #         for j, i in enumerate(idxs[0])
-    #     ]
+        scores, idxs = self.image_index.search(emb, top_k)
+        return [
+            {
+                "image_path": self.image_id_map[i],  # integer ID
+                "score": float(scores[0][j])
+            }
+            for j, i in enumerate(idxs[0])
+        ]
 
     def image_to_text(self, image: Image.Image, top_k=5):
-        image = self.image_transform(image).unsqueeze(0)
+        image = self.image_transform(image).unsqueeze(0).to(self.device)
         with torch.no_grad():
             emb = self.clip_model.encode_image(image).cpu().numpy()
         emb = self._normalize(emb)
@@ -107,7 +59,6 @@ class RetrievalService:
         print(f"DEBUG: Returning results: {results}")
         return results
 
-
     def text_to_text(self, text: str, top_k: int = 5):
         with torch.no_grad():
             emb = self.clip_model.encode_text(text).cpu().numpy()
@@ -125,67 +76,213 @@ class RetrievalService:
         
         print(f"DEBUG: Text-to-text results: {results}")
         return results
+
+    def image_to_image(self, image: Image.Image, top_k: int = 5):
+        image = self.image_transform(image).unsqueeze(0).to(self.device)
+        with torch.no_grad():
+            emb = self.clip_model.encode_image(image).cpu().numpy()
+        emb = self._normalize(emb)
+
+        scores, idxs = self.image_index.search(emb, top_k)
+        return [
+            {
+                "image_path": self.image_id_map[i],  # integer ID
+                "score": float(scores[0][j])
+            }
+            for j, i in enumerate(idxs[0])
+        ]
+
+
+# Old Code
+
+# import faiss
+# import pickle
+# import torch
+# import numpy as np
+# import os
+# from PIL import Image
+# from torchvision import transforms
+# from typing import List, Dict, Any
+
+# class RetrievalService:
+#     def __init__(self, clip_model, image_index_path, text_index_path,
+#                  image_map_path, text_map_path, preprocess):
+
+#         self.device = torch.device("cpu")
+#         self.clip_model = clip_model
+
+#         self.image_index = faiss.read_index(image_index_path)
+#         self.text_index = faiss.read_index(text_index_path)
+
+#         with open(image_map_path, "rb") as f:
+#             self.image_id_map = pickle.load(f)
+
+#         with open(text_map_path, "rb") as f:
+#             self.text_id_map = pickle.load(f)
+
+#         self.image_transform = transforms.Compose([
+#             transforms.Resize((224, 224)),
+#             transforms.ToTensor(),
+#             transforms.Normalize(
+#                 mean=preprocess["mean"],
+#                 std=preprocess["std"]
+#             )
+#         ])
+
+#     def _normalize(self, x):
+#         return x / np.linalg.norm(x, axis=1, keepdims=True)
+    
+
+#     def text_to_image(self, text: str, top_k: int = 5) -> List[Dict[str, Any]]:
+#         raw_results = self.retrieval_service.text_to_image(text, top_k)
+        
+#         formatted = []
+#         for res in raw_results:
+#             idx = int(res["image_path"])  # the FAISS index (integer)
+            
+#             try:
+#                 pil_img = self.dataset[idx]["image"]  # directly get PIL.Image
+#                 formatted.append({
+#                     "image": pil_img,  # ← pass PIL.Image to UI
+#                     "score": float(res["score"])
+#                 })
+#             except (IndexError, KeyError):
+#                 formatted.append({
+#                     "image": None,
+#                     "score": float(res["score"])
+#                 })
+        
+#         return formatted
+
+
+#     # def text_to_image(self, text: str, top_k: int = 5) -> List[Dict[str, Any]]:
+#     #     raw_results = self.retrieval_service.text_to_image(text, top_k)
+        
+#     #     formatted = []
+#     #     for res in raw_results:
+#     #         img_id = res["image_path"]  # int or str
+#     #         img_id_str = str(img_id)
+#     #         img_filename = f"{img_id_str}.jpg"  # always append .jpg, no .endswith
+#     #         full_path = os.path.join("flickr8k_images", img_filename)
+            
+#     #         if os.path.exists(full_path):
+#     #             formatted.append({
+#     #                 "image_path": full_path,
+#     #                 "score": float(res["score"])
+#     #             })
+#     #         else:
+#     #             formatted.append({
+#     #                 "image_path": "https://via.placeholder.com/300?text=Not+in+demo",
+#     #                 "score": float(res["score"])
+#     #             })
+        
+#     #     return formatted
+
+
+#     # def text_to_image(self, text, top_k=5):
+#     #     with torch.no_grad():
+#     #         emb = self.clip_model.encode_text(text).cpu().numpy()
+#     #     emb = self._normalize(emb)
+
+#     #     scores, idxs = self.image_index.search(emb, top_k)
+#     #     return [
+#     #         {
+#     #             "image_path": self.image_id_map[i],
+#     #             "score": float(scores[0][j])
+#     #         }
+#     #         for j, i in enumerate(idxs[0])
+#     #     ]
+
+#     def image_to_text(self, image: Image.Image, top_k=5):
+#         image = self.image_transform(image).unsqueeze(0)
+#         with torch.no_grad():
+#             emb = self.clip_model.encode_image(image).cpu().numpy()
+#         emb = self._normalize(emb)
+
+#         scores, idxs = self.text_index.search(emb, top_k)
+#         results = [self.text_id_map[i] for i in idxs[0]]
+#         print(f"DEBUG: Returning results: {results}")
+#         return results
+
+
+#     def text_to_text(self, text: str, top_k: int = 5):
+#         with torch.no_grad():
+#             emb = self.clip_model.encode_text(text).cpu().numpy()
+#         emb = self._normalize(emb)
+
+#         scores, idxs = self.text_index.search(emb, top_k)
+        
+#         results = []
+#         for j, i in enumerate(idxs[0]):
+#             caption = self.text_id_map[i]  # assuming text_id_map stores the actual caption string
+#             results.append({
+#                 "caption": caption,
+#                 "score": float(scores[0][j])
+#             })
+        
+#         print(f"DEBUG: Text-to-text results: {results}")
+#         return results
     
 
-    # def image_to_image(self, image: Image.Image, top_k=5):
-    #     """
-    #     Image → Image retrieval: encode input image, search image index, return image IDs and scores.
-    #     """
-    #     image = self.image_transform(image).unsqueeze(0).to(self.device)
-    #     with torch.no_grad():
-    #         emb = self.clip_model.encode_image(image).cpu().numpy()
-    #     emb = self._normalize(emb)
-
-    #     scores, idxs = self.image_index.search(emb, top_k)
-    #     return [
-    #         {
-    #             "image_path": self.image_id_map[i],  # integer ID
-    #             "score": float(scores[0][j])
-    #         }
-    #         for j, i in enumerate(idxs[0])
-    #     ]
+#     # def image_to_image(self, image: Image.Image, top_k=5):
+#     #     """
+#     #     Image → Image retrieval: encode input image, search image index, return image IDs and scores.
+#     #     """
+#     #     image = self.image_transform(image).unsqueeze(0).to(self.device)
+#     #     with torch.no_grad():
+#     #         emb = self.clip_model.encode_image(image).cpu().numpy()
+#     #     emb = self._normalize(emb)
+
+#     #     scores, idxs = self.image_index.search(emb, top_k)
+#     #     return [
+#     #         {
+#     #             "image_path": self.image_id_map[i],  # integer ID
+#     #             "score": float(scores[0][j])
+#     #         }
+#     #         for j, i in enumerate(idxs[0])
+#     #     ]
     
-    # def image_to_image(self, image: Image.Image, top_k: int = 5) -> List[Dict[str, Any]]:
-    #     raw_results = self.retrieval_service.image_to_image(image, top_k)  # now exists
-    #     # ... same logic as above ...
+#     # def image_to_image(self, image: Image.Image, top_k: int = 5) -> List[Dict[str, Any]]:
+#     #     raw_results = self.retrieval_service.image_to_image(image, top_k)  # now exists
+#     #     # ... same logic as above ...
         
-    #     formatted = []
-    #     for res in raw_results:
-    #         img_id = res["image_path"]
-    #         img_id_str = str(img_id)
-    #         img_filename = f"{img_id_str}.jpg"
-    #         full_path = os.path.join("flickr8k_images", img_filename)
+#     #     formatted = []
+#     #     for res in raw_results:
+#     #         img_id = res["image_path"]
+#     #         img_id_str = str(img_id)
+#     #         img_filename = f"{img_id_str}.jpg"
+#     #         full_path = os.path.join("flickr8k_images", img_filename)
             
-    #         if os.path.exists(full_path):
-    #             formatted.append({
-    #                 "image_path": full_path,
-    #                 "score": float(res["score"])
-    #             })
-    #         else:
-    #             formatted.append({
-    #                 "image_path": "https://via.placeholder.com/300?text=Not+in+demo",
-    #                 "score": float(res["score"])
-    #             })
+#     #         if os.path.exists(full_path):
+#     #             formatted.append({
+#     #                 "image_path": full_path,
+#     #                 "score": float(res["score"])
+#     #             })
+#     #         else:
+#     #             formatted.append({
+#     #                 "image_path": "https://via.placeholder.com/300?text=Not+in+demo",
+#     #                 "score": float(res["score"])
+#     #             })
         
-    #     return formatted
+#     #     return formatted
     
-    def image_to_image(self, image: Image.Image, top_k: int = 5) -> List[Dict[str, Any]]:
-        raw_results = self.retrieval_service.image_to_image(image, top_k)
+#     def image_to_image(self, image: Image.Image, top_k: int = 5) -> List[Dict[str, Any]]:
+#         raw_results = self.retrieval_service.image_to_image(image, top_k)
         
-        formatted = []
-        for res in raw_results:
-            idx = int(res["image_path"])
+#         formatted = []
+#         for res in raw_results:
+#             idx = int(res["image_path"])
             
-            try:
-                pil_img = self.dataset[idx]["image"]
-                formatted.append({
-                    "image": pil_img,
-                    "score": float(res["score"])
-                })
-            except (IndexError, KeyError):
-                formatted.append({
-                    "image": None,
-                    "score": float(res["score"])
-                })
+#             try:
+#                 pil_img = self.dataset[idx]["image"]
+#                 formatted.append({
+#                     "image": pil_img,
+#                     "score": float(res["score"])
+#                 })
+#             except (IndexError, KeyError):
+#                 formatted.append({
+#                     "image": None,
+#                     "score": float(res["score"])
+#                 })
         
-        return formatted
+#         return formatted