Update app.py

app.py

@@ -13,6 +13,7 @@ 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")
@@ -38,28 +39,28 @@ class ImageSearchSystem:
         self.initialized = False
     
     def initialize_dataset(self) -> None:
-        """Automatically download and process the dataset."""
+        """Automatically download and process the dataset with a 500-sample limit."""
         try:
-            dataset_path = os.path.expanduser("~/.kagglehub/datasets/alessandrasala79/ai-vs-human-generated-dataset")
-            image_folder = os.path.join(dataset_path, 'test_data_v2')
-            
-            # Download dataset if not already present
-            if not os.path.exists(dataset_path):
-                logger.info("Downloading dataset from Kaggle...")
-                dataset_path = kagglehub.dataset_download("alessandrasala79/ai-vs-human-generated-dataset")
+            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
-            self.image_paths = [f for f in Path(image_folder).glob("**/*") if f.suffix.lower() in ['.jpg', '.jpeg', '.png']]
+            all_images = [f for f in Path(image_folder).glob("**/*") if f.suffix.lower() in ['.jpg', '.jpeg', '.png']]
             
-            if not self.image_paths:
+            if not all_images:
                 raise ValueError("No images found in the dataset!")
 
-            logger.info(f"Successfully loaded {len(self.image_paths)} images.")
-            
+            # 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