feat: reimplement the 3 batch methods

gee-community · Oct 7, 2024 · a465bc8 · a465bc8
1 parent 718e453
commit a465bc8
Showing 1 changed file with 129 additions and 1 deletion.
diff --git a/geetools/batch/Export/__init__.py b/geetools/batch/Export/__init__.py
@@ -46,7 +46,7 @@ def toAsset(
                 index_property: The property of the image to use as name. Default is "system:id".
                 description: The description of the task.
                 assetId: The asset id where to export the image collection.
-                **kwargs: every paramet that you would use for a vanilla ee.batch.Export.image.toAsset
+                **kwargs: every parameter that you would use for a vanilla ee.batch.Export.image.toAsset
 
             Returns:
                 The task created.
@@ -92,3 +92,131 @@ def toAsset(
                 task_list.append(ee.batch.Export.image.toAsset(**kwargs))
 
             return task_list
+
+        @staticmethod
+        def toDrive(
+            imagecollection: ee.ImageCollection,
+            index_property: str = "system:id",
+            description: Optional[str] = None,
+            folder: Optional[str] = None,
+            **kwargs,
+        ) -> List[ee.batch.Task]:
+            """Creates a list of tasks to export an EE ImageCollection to Google Drive.
+
+            The method will create a folder in Google Drive with the description value and populate
+            it with the exported images. Each image in the Collection Will be named using the
+            index_property value of the image. If no Folder is provided the folder will be created at the root
+            of the current drive and use the description name.
+
+            Parameters:
+                imagecollection: The image collection to export.
+                index_property: The property of the image to use as name. Default is "system:id".
+                description: The description of the task.
+                folder: The folder id where to export the image collection. It will be stored at the root of the drive.
+                **kwargs: every parameter that you would use for a vanilla ee.batch.Export.image.toDrive
+
+            Returns:
+                The list of created tasks
+
+            Examples:
+                .. code-block:: python
+
+                    import ee
+                    import geetools
+
+                    ee.Initialize()
+
+                    # create a test image collection
+                    collection = ee.ImageCollection("COPERNICUS/S2").limit(5)
+
+                    # export the collection
+                    tasks = geetools.batch.Export.imagecollection.toDrive(collection, "system:index", "test export")
+            """
+            # sanity check on parameters
+            # renaming them for mypy type reassignment and compactness
+            desc = description or folder
+            folder or description
+
+            # loop over the collection and export each image
+            nb_images = imagecollection.size().getInfo()
+            imageList = imagecollection.toList(nb_images)
+            task_list = []
+            for i in range(nb_images):
+                # extract image information
+                locImage = ee.Image(imageList.get(i))
+                loc_id = locImage.get(index_property).getInfo()
+
+                # override the parameters related to the image itself
+                # the folder will be created by the first task
+                kwargs["image"] = locImage
+                kwargs["folder"] = utils.format_asset_id(folder)
+                kwargs["description"] = utils.format_description(f"{desc}_{loc_id}")
+
+                # create the task
+                task_list.append(ee.batch.Export.image.toDrive(**kwargs))
+
+            return task_list
+
+        @staticmethod
+        def toCloudStorage(
+            imagecollection: ee.ImageCollection,
+            index_property: str = "system:id",
+            description: Optional[str] = None,
+            folder: Optional[str] = None,
+            **kwargs,
+        ) -> List[ee.batch.Task]:
+            """Creates a list of tasks to export an EE ImageCollection to Google cloud.
+
+            The method will create a folder in Google cloud bucket with the description value and populate
+            it with the exported images. Each image in the Collection Will be named using the
+            index_property value of the image. If no Folder is provided the folder will be created at the root
+            of the bucket and use the description name.
+
+            Parameters:
+                imagecollection: The image collection to export.
+                index_property: The property of the image to use as name. Default is "system:id".
+                description: The description of the task.
+                folder: The folder id where to export the image collection. It will be stored at the root of the drive.
+                **kwargs: every parameter that you would use for a vanilla ee.batch.Export.image.toCloudStorage
+
+            Returns:
+                The list of created tasks
+
+            Examples:
+                .. code-block:: python
+
+                    import ee
+                    import geetools
+
+                    ee.Initialize()
+
+                    # create a test image collection
+                    collection = ee.ImageCollection("COPERNICUS/S2").limit(5)
+
+                    # export the collection
+                    tasks = geetools.batch.Export.imagecollection.toDrive(collection, "system:index", "test export")
+            """
+            # sanity check on parameters
+            # renaming them for mypy type reassignment and compactness
+            desc = description or folder
+            folder or description
+
+            # loop over the collection and export each image
+            nb_images = imagecollection.size().getInfo()
+            imageList = imagecollection.toList(nb_images)
+            task_list = []
+            for i in range(nb_images):
+                # extract image information
+                locImage = ee.Image(imageList.get(i))
+                loc_id = locImage.get(index_property).getInfo()
+
+                # override the parameters related to the image itself
+                # the folder will be created by the first task
+                kwargs["image"] = locImage
+                kwargs["fileNamePrefix"] = f"{utils.format_asset_id(folder)}/"
+                kwargs["description"] = utils.format_description(f"{desc}_{loc_id}")
+
+                # create the task
+                task_list.append(ee.batch.Export.image.toCloudStorage(**kwargs))
+
+            return task_list