retrieval-ai / app.py
joy1515's picture
Update app.py
500b4bf verified
import gradio as gr
import torch
from transformers import CLIPProcessor, CLIPModel
import numpy as np
import kagglehub
from PIL import Image
import os
from pathlib import Path
import logging
import faiss
from tqdm import tqdm
import speech_recognition as sr
from gtts import gTTS
import tempfile
import torch.nn.utils.prune as prune
import random
# Configure logging
logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
logger = logging.getLogger(__name__)
class ImageSearchSystem:
def __init__(self):
self.device = "cuda" if torch.cuda.is_available() else "cpu"
logger.info(f"Using device: {self.device}")
# Load CLIP model
self.processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch16")
self.model = CLIPModel.from_pretrained("openai/clip-vit-base-patch16").to(self.device)
# Prune the model (optimize memory usage)
for name, module in self.model.named_modules():
if isinstance(module, torch.nn.Linear):
prune.l1_unstructured(module, name='weight', amount=0.2)
# Initialize dataset
self.image_paths = []
self.index = None
self.initialized = False
def initialize_dataset(self) -> None:
"""Automatically download and process the dataset with a 500-sample limit."""
try:
logger.info("Downloading dataset from KaggleHub...")
dataset_path = kagglehub.dataset_download("alessandrasala79/ai-vs-human-generated-dataset")
image_folder = os.path.join(dataset_path, 'test_data_v2') # Adjust if needed
# Validate dataset
if not os.path.exists(image_folder):
raise FileNotFoundError(f"Expected dataset folder not found: {image_folder}")
# Load images dynamically
all_images = [f for f in Path(image_folder).glob("**/*") if f.suffix.lower() in ['.jpg', '.jpeg', '.png']]
if not all_images:
raise ValueError("No images found in the dataset!")
# Limit dataset to 500 randomly selected samples
self.image_paths = random.sample(all_images, min(500, len(all_images)))
logger.info(f"Loaded {len(self.image_paths)} images (limited to 500 samples).")
# Create image index
self._create_image_index()
self.initialized = True
except Exception as e:
logger.error(f"Dataset initialization failed: {str(e)}")
raise
def _create_image_index(self, batch_size: int = 32) -> None:
"""Create FAISS index for fast image retrieval."""
try:
all_features = []
for i in tqdm(range(0, len(self.image_paths), batch_size), desc="Indexing images"):
batch_paths = self.image_paths[i:i + batch_size]
batch_images = [Image.open(img).convert("RGB") for img in batch_paths]
if batch_images:
inputs = self.processor(images=batch_images, return_tensors="pt", padding=True)
inputs = {k: v.to(self.device) for k, v in inputs.items()}
with torch.no_grad():
features = self.model.get_image_features(**inputs)
features = features / features.norm(dim=-1, keepdim=True)
all_features.append(features.cpu().numpy())
all_features = np.concatenate(all_features, axis=0)
self.index = faiss.IndexFlatIP(all_features.shape[1])
self.index.add(all_features)
logger.info("Image index created successfully")
except Exception as e:
logger.error(f"Failed to create image index: {str(e)}")
raise
def search(self, query: str, audio_path: str = None, k: int = 5):
"""Search for images using text or speech."""
try:
if not self.initialized:
raise RuntimeError("System not initialized. Call initialize_dataset() first.")
# Convert speech to text if audio input is provided
if audio_path:
recognizer = sr.Recognizer()
with sr.AudioFile(audio_path) as source:
audio_data = recognizer.record(source)
try:
query = recognizer.recognize_google(audio_data)
except sr.UnknownValueError:
return [], "Could not understand the spoken query.", None
# Process text query
inputs = self.processor(text=[query], return_tensors="pt", padding=True)
inputs = {k: v.to(self.device) for k, v in inputs.items()}
with torch.no_grad():
text_features = self.model.get_text_features(**inputs)
text_features = text_features / text_features.norm(dim=-1, keepdim=True)
# Search FAISS index
scores, indices = self.index.search(text_features.cpu().numpy(), k)
results = [Image.open(self.image_paths[idx]) for idx in indices[0]]
# Generate Text-to-Speech
tts = gTTS(f"Showing results for {query}")
temp_audio = tempfile.NamedTemporaryFile(delete=False, suffix=".mp3")
tts.save(temp_audio.name)
return results, query, temp_audio.name
except Exception as e:
logger.error(f"Search failed: {str(e)}")
return [], "Error during search.", None
def create_demo_interface() -> gr.Interface:
"""Create Gradio interface with dark mode & speech support."""
system = ImageSearchSystem()
try:
system.initialize_dataset()
except Exception as e:
logger.error(f"Failed to initialize system: {str(e)}")
raise
examples = [
["a beautiful landscape with mountains"],
["people working in an office"],
["a cute dog playing"],
["a modern city skyline at night"],
["a delicious-looking meal"]
]
return gr.Interface(
fn=system.search,
inputs=[
gr.Textbox(label="Enter your search query:", placeholder="Describe the image...", lines=2),
gr.Audio(sources=["microphone"], type="filepath", label="Speak Your Query (Optional)")
],
outputs=[
gr.Gallery(label="Search Results", show_label=True, columns=5, height="auto"),
gr.Textbox(label="Spoken Query", interactive=False),
gr.Audio(label="Results Spoken Out Loud")
],
title="Multi-Modal Image Search",
description="Use text or voice to search for images.",
theme="dark",
examples=examples,
cache_examples=True,
css=".gradio-container {background-color: #121212; color: #ffffff;}"
)
if __name__ == "__main__":
try:
demo = create_demo_interface()
demo.launch(share=True, max_threads=40)
except Exception as e:
logger.error(f"Failed to launch app: {str(e)}")
raise