Max limit

run_google_lens.py

@@ -31,6 +31,7 @@ class GoogleLensPipeline:
         accept_button_xpath=None,
         file_input_xpath=None,
         root_directory=None,
+        max_images=100000,  # Added max_images parameter
     ):
         self.link = link or "https://lens.google.com/"
         self.accept_button_xpath = accept_button_xpath or '//span[text()="Accept all"]'
@@ -41,6 +42,7 @@ class GoogleLensPipeline:
         self.tmp_directory_two = os.path.join(self.root_directory, "tmp_two")
         self.final_directory = os.path.join(self.root_directory, "images")
 
+        self.max_images = max_images  # Initialize max_images
         self.setup_directories()
 
     def start_driver(self):
@@ -87,7 +89,7 @@ class GoogleLensPipeline:
         except Exception as e:
             logger.error(f"Failed to open page: {str(e)}")
 
-    def accept_terms(self, timeout=30):
+    def accept_terms(self, timeout=20):
         try:
             wait = WebDriverWait(self.driver, timeout)
             accept_button = wait.until(
@@ -98,7 +100,7 @@ class GoogleLensPipeline:
         except Exception as e:
             logger.error(f"Failed to accept terms: {str(e)}")
 
-    def upload_image(self, image_path, timeout=30, sleep=10):
+    def upload_image(self, image_path, timeout=20, sleep=10):
         try:
             wait = WebDriverWait(self.driver, timeout)
             file_input = wait.until(
@@ -164,6 +166,14 @@ class GoogleLensPipeline:
         try:
             await self.run_pipeline(image_path, save_path=self.tmp_directory_one)
 
+            # Check if the maximum number of images has been reached
+            total_images = len(os.listdir(self.final_directory))
+            if total_images >= self.max_images:
+                logger.info(
+                    f"Reached maximum number of images ({self.max_images}). Stopping."
+                )
+                return
+
             for val in os.listdir(self.tmp_directory_one):
                 try:
                     await self.run_pipeline(
@@ -174,6 +184,15 @@ class GoogleLensPipeline:
                         os.path.join(self.tmp_directory_one, val),
                         os.path.join(self.final_directory, val),
                     )
+
+                    # Check after each image
+                    total_images = len(os.listdir(self.final_directory))
+                    if total_images >= self.max_images:
+                        logger.info(
+                            f"Reached maximum number of images ({self.max_images}). Stopping."
+                        )
+                        return
+
                 except Exception as e:
                     logger.error(
                         f"Error processing file {val} in tmp_directory_one: {str(e)}"
@@ -189,6 +208,15 @@ class GoogleLensPipeline:
                         os.path.join(self.tmp_directory_two, val),
                         os.path.join(self.final_directory, val),
                     )
+
+                    # Check after each image
+                    total_images = len(os.listdir(self.final_directory))
+                    if total_images >= self.max_images:
+                        logger.info(
+                            f"Reached maximum number of images ({self.max_images}). Stopping."
+                        )
+                        return
+
                 except Exception as e:
                     logger.error(
                         f"Error processing file {val} in tmp_directory_two: {str(e)}"
@@ -220,6 +248,15 @@ if __name__ == "__main__":
 
     path = "/home/user/app/images/"
 
+    pipeline = GoogleLensPipeline(
+        root_directory=path, max_images=100000
+    )  # Set max_images
+
     for image_path in get_image_paths(path):
-        pipeline = GoogleLensPipeline(root_directory=path)
+        total_images = len(os.listdir(pipeline.final_directory))
+        if total_images >= pipeline.max_images:
+            logger.info(
+                f"Reached maximum number of images ({pipeline.max_images}). Stopping."
+            )
+            break
         asyncio.run(pipeline.loop_run_pipeline(image_path))