Update app.py

app.py

@@ -1,167 +1,175 @@
-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
-
-# 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)
-        
-        # Initialize dataset
-        self.image_paths = []
-        self.index = None
-        self.initialized = False
-        
-    def initialize_dataset(self) -> None:
-        """Download and process dataset"""
-        try:
-            path = kagglehub.dataset_download("alessandrasala79/ai-vs-human-generated-dataset")
-            image_folder = os.path.join(path, 'test_data_v2')
-            
-            self.image_paths = [
-                f for f in Path(image_folder).glob("**/*") 
-                if f.suffix.lower() in ['.jpg', '.jpeg', '.png']
-            ]
-            
-            if not self.image_paths:
-                raise ValueError(f"No images found in {image_folder}")
-            
-            logger.info(f"Found {len(self.image_paths)} images")
-            
-            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"""
-        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(source="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, enable_queue=True, max_threads=40)
-    except Exception as e:
-        logger.error(f"Failed to launch app: {str(e)}")
-        raise
+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
+
+# 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
+        pruned_model = self.model.prune(
+            pruning_method="l1_unstructured",
+            num_heads_to_prune=10,
+            num_layers_to_prune=2,
+        )
+        self.model = pruned_model
+        
+        # Initialize dataset
+        self.image_paths = []
+        self.index = None
+        self.initialized = False
+        
+    def initialize_dataset(self) -> None:
+        """Download and process dataset"""
+        try:
+            path = kagglehub.dataset_download("alessandrasala79/ai-vs-human-generated-dataset")
+            image_folder = os.path.join(path, 'test_data_v2')
+            
+            self.image_paths = [
+                f for f in Path(image_folder).glob("**/*") 
+                if f.suffix.lower() in ['.jpg', '.jpeg', '.png']
+            ]
+            
+            if not self.image_paths:
+                raise ValueError(f"No images found in {image_folder}")
+            
+            logger.info(f"Found {len(self.image_paths)} images")
+            
+            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"""
+        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(source="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, enable_queue=True, max_threads=40)
+    except Exception as e:
+        logger.error(f"Failed to launch app: {str(e)}")
+        raise