download_landsat.py

# -*- coding: UTF-8 -*-
# Python
"""
16-12-2022
@author: jeremy auclair

Original script is from USGS Machine to Machine python exampes https://m2m.cr.usgs.gov/api/docs/example/download_data-py

A few modifications have been added to adapt the script to the ModSpa project.

Download LandSat data pre-modspa
"""

import json  # to open and read json files
import requests  # to send url requests
import sys  # for path management
import os  # for path exploration
import time  # for wait time between downloads
import ruamel.yaml  # to open yaml config file
import geopandas as gpd  # to read shapefiles
import csv  # to save download
from datetime import datetime  # manage dates
from dateutil.relativedelta import relativedelta  # date math
from fnmatch import fnmatch  # to compare character strings
import tarfile as tr  # to extract file for tar archive
from tqdm import tqdm  # to print progress bars
from p_tqdm import p_map  # for multiprocessing with progress bars
from psutil import cpu_count  # to get number of physical cores available
from typing import List  # to declare variables


# send http request
def sendRequest(url: str, data: dict, apiKey = None) -> str:  
    """
    Send http request

    Arguments
    =========
    
    1. url: ``str``
        download url
    2. data: ``dict``
        payload for request
    3. apiKey: ``int``
        apiKey for USGS m2m service

    Returns
    =======
    
    1. output['data']: ``str``
        result of the http request
    """

    # Format data with json
    json_data = json.dumps(data)
    
    if apiKey == None:
        response = requests.post(url, json_data)
    else:
        headers = {'X-Auth-Token': apiKey}              
        response = requests.post(url, json_data, headers = headers)
    
    # Try request
    try:
      httpStatusCode = response.status_code 
      if response == None:
          print("No output from service")
          sys.exit()
      output = json.loads(response.text)
    #   print(output)	
      if output['errorCode'] != None:
          print(output['errorCode'], "- ", output['errorMessage'])
          sys.exit()
      if  httpStatusCode == 404:
          print("404 Not Found")
          sys.exit()
      elif httpStatusCode == 401: 
          print("401 Unauthorized")
          sys.exit()
      elif httpStatusCode == 400:
          print("Error Code", httpStatusCode)
          sys.exit()
    except Exception as e: 
          response.close()
          print(e)
          sys.exit()
    response.close()
    
    return output['data']


def download_file_multiprocess(args: tuple) -> str:
    """
    Code adapted from [github](https://gist.github.com/yanqd0/c13ed29e29432e3cf3e7c38467f42f51).

    Downloads tar file from given url, unpacks the red, nir and pixel quality bands and deletes
    the tar archive. Is suited for multiprocessing.

    Arguments (packed in args)
    ==========================
    
    1. url: ``str``
        download url
    2. file_path: ``str``
        path to save file
    3. chunk_size: ``int``
        download chunk size

    Returns
    =======
    
    1. file_path: ``str``
        path to saved file
    """

    # Unpack args
    url, file_path, chunk_size = args

    # Get file name and path
    if url.count("gen-bundle?"):
        file_name = url.split('=')[1].split('&')[0] + '.tar'
    else:
        file_name = url.split('/')[-1]
    file_path = file_path + os.sep + file_name
    extract_path = file_path[:-4] #+ os.sep
    
    # Check if file exists
    if os.path.exists(file_path):
        if os.path.exists(extract_path):
            count=0
            for path in os.listdir(extract_path):
                # check if current path is a file
                if os.path.isfile(os.path.join(extract_path, path)): count += 1
            
            if count == 3:
                print(f'File {file_name} already exists !')
                return extract_path

    # Prepare download
    resp = requests.get(url, stream=True)
    total = int(resp.headers.get('content-length', 0))

    # Download file
    with open(file_path, 'wb') as file, tqdm(
        desc = file_name,
        total = total,
        unit = 'iB',
        unit_scale = True,
        unit_divisor = 1024,
    ) as bar:
        for data in resp.iter_content(chunk_size = chunk_size):
            size = file.write(data)
            bar.update(size)

    # Extract red, nir and pixel quality files from tar file
    with tr.open(file_path, mode = 'r') as mytar:
        file_list = (mytar.getmembers())
        for f in file_list:
            if fnmatch(f.get_info()['name'], '*_SR_B4.TIF'): # red band
                mytar.extract(f, path = extract_path)
            elif fnmatch(f.get_info()['name'], '*_SR_B5.TIF'): # nir band
                mytar.extract(f, path = extract_path)
            elif fnmatch(f.get_info()['name'], '*_QA_PIXEL.TIF'): # pixel quality
                mytar.extract(f, path = extract_path)
    
    return extract_path


def download_landsat_data(path_to_config_file: str, start_date: str, end_date: str, shapefile: str, download_path: str, save_path: str, collection: str = 'landsat_ot_c2_l2', cloud_cover_limit: int = 80, max_downloads: int = 200, chunk_size: int = 20480, max_cpu: int = 4) -> List[str]:
    """
    Code adapted from the Machine to Machine USGS website [download_data.py](https://m2m.cr.usgs.gov/api/docs/example/download_data-py).

    download_landsat_data looks for all products in the usgs earth explorer database during a specific time window and covering the 
    whole shapefile enveloppe (several LandSat tiles might be needed). It then downloads that data into the download path defined
    in the config file. Paths to the downloaded data are returned and saved as a ``csv`` file.

    Arguments
    =========
    
    1. path_to_config_file: ``str``
        path to the eodag configuration file that contains user credentials
    2. start_date: ``str``
        beginning of the time window to download (format: `yyyy-mm-dd`)
    3. end_date: ``str``
        end of the time window to download (format: `yyyy-mm-dd`)
    4. shapefile: ``str``
        path to the shapefile (`.shp`) for which the data is downloaded
    5. download_path: ``str``
        path to save directory for LandSat images
    6. save_path: ``str``
        path where a csv file containing the product paths will be saved
    7. collection: ``str``
        name of the LandSat collection/dataset in which to search products
    8. cloud_cover_limit: ``int`` ``default = 80``
        maximum percentage to pass the filter before download (between 0 and 100)
    9. max_downloads: ``int`` ``default = 200``
        maximum number of images to download in a single run
    10. chunk_size: ``int`` ``default = 20480``
        download chunk size in bytes
    11. max_cpu: ``int`` ``default = 4``
        max number of cpu cores to use

    Returns
    =======
    
    1. product_paths: ``List[str]``
        a list of the paths to the downloaded data
    """
    
    # user input    
    # find credentials in eodag config file    
    parameter_list, _, _ = ruamel.yaml.util.load_yaml_guess_indent(open(path_to_config_file))
    username = parameter_list['usgs']['api']['credentials']['username']
    password = parameter_list['usgs']['api']['credentials']['password']

    # Open shapefile containing geometry
    geopandas_shape = gpd.read_file(shapefile)
    geopandas_shape = geopandas_shape.to_crs(epsg='4326')  # Force WGS84 projection
    bounds = list(geopandas_shape.geometry.total_bounds)  # In WGS84 projection
    
    # Change start and end date to better cover the chosen period
    start_date = (datetime.strptime(start_date, '%Y-%m-%d') - relativedelta(months = 1)).strftime('%Y-%m-%d')
    end_date = (datetime.strptime(end_date, '%Y-%m-%d') + relativedelta(months = 1)).strftime('%Y-%m-%d')

    print("Sending request...\n")
    
    serviceUrl = "https://m2m.cr.usgs.gov/api/api/json/stable/"
    
    # login
    credentials = {'username' : username, 'password' : password}
    
    apiKey = sendRequest(serviceUrl + "login", credentials)
    
    print("API Key: " + apiKey + "\n")
    
    datasetName = collection
    
    spatialFilter =  {'filterType' : "mbr",
                    'lowerLeft' : {'latitude' : bounds[1], 'longitude' : bounds[0]},
                    'upperRight' : { 'latitude' : bounds[3], 'longitude' : bounds[2]}}
                     
    temporalFilter = {'start' : start_date, 'end' : end_date}

    cloudCoverFilter = {"max": cloud_cover_limit,
                        "min": 0,
                        "includeUnknown": True}
    
    payload = {'datasetName' : datasetName,
                'spatialFilter' : spatialFilter,
                'temporalFilter' : temporalFilter,
                'cloudCoverFilter' : cloudCoverFilter}                     
    
    print("Searching datasets...\n")

    datasets = sendRequest(serviceUrl + "dataset-search", payload, apiKey)
    
    print("Found ", len(datasets), " datasets")

    # List of paths to downloaded products
    landsat_products = []
    
    # download datasets
    for dataset in datasets:
        if dataset['datasetAlias'] != datasetName:
            continue
            
        # I don't want to limit my results, but using the dataset-filters request, you can
        # find additional filters
        
        acquisitionFilter = {"end": end_date,
                             "start": start_date}        
            
        payload = {'datasetName' : dataset['datasetAlias'],
                    'maxResults' : max_downloads,
                    'startingNumber' : 1, 
                    'sceneFilter' : {
                            'spatialFilter' : spatialFilter,
                            'acquisitionFilter' : acquisitionFilter,
                            'cloudCoverFilter' : cloudCoverFilter},
                    'username' : username, 
                    'password' : password}
        
        # Now I need to run a scene search to find data to download
        print('\nSearching scenes for dataset ' + dataset['collectionName'] + '...\n')   
        
        scenes = sendRequest(serviceUrl + "scene-search", payload, apiKey)
    
        # Did we find anything?
        if scenes['recordsReturned'] > 0:
            # Aggregate a list of scene ids
            sceneIds = []
            for result in scenes['results']:
                # Add this scene to the list I would like to download
                sceneIds.append(result['entityId'])
            
            # Find the download options for these scenes
            # NOTE :: Remember the scene list cannot exceed 50,000 items!
            payload = {'datasetName' : dataset['datasetAlias'], 'entityIds' : sceneIds}
                                
            downloadOptions = sendRequest(serviceUrl + "download-options", payload, apiKey)
        
            # Aggregate a list of available products
            downloads = []
            for product in downloadOptions:
                    # Make sure the product is available for this scene
                    if product['available'] == True:
                         downloads.append({'entityId' : product['entityId'],
                                           'productId' : product['id']})
            
            print(len(downloads), 'products available with current parameters.\n')

            # Did we find products?
            if downloads:
                requestedDownloadsCount = len(downloads)
                # set a label for the download request
                label = "download-sample"
                payload = {'downloads' : downloads,
                                             'label' : label}
                # Call the download to get the direct download urls
                requestResults = sendRequest(serviceUrl + "download-request", payload, apiKey)          
                              
                # PreparingDownloads has a valid link that can be used but data may not be immediately available
                # Call the download-retrieve method to get download that is available for immediate download
                if requestResults['preparingDownloads'] != None and len(requestResults['preparingDownloads']) > 0:
                    payload = {'label' : label}
                    moreDownloadUrls = sendRequest(serviceUrl + "download-retrieve", payload, apiKey)
                    
                    downloadIds = []  
                    
                    for download in moreDownloadUrls['available']:
                        downloadIds.append(download['downloadId'])
                        print("DOWNLOAD: " + download['url'])
                        
                    for download in moreDownloadUrls['requested']:   
                        downloadIds.append(download['downloadId'])
                        print("DOWNLOAD: " + download['url'])
                     
                    # Didn't get all of the reuested downloads, call the download-retrieve method again probably after 30 seconds
                    while len(downloadIds) < requestedDownloadsCount: 
                        preparingDownloads = requestedDownloadsCount - len(downloadIds)
                        print("\n", preparingDownloads, "downloads are not available. Waiting for 30 seconds.\n")
                        time.sleep(30)
                        print("Trying to retrieve data\n")
                        moreDownloadUrls = sendRequest(serviceUrl + "download-retrieve", payload, apiKey)
                        for download in moreDownloadUrls['available']:                            
                            if download['downloadId'] not in downloadIds:
                                downloadIds.append(download['downloadId'])
                                print("DOWNLOAD: " + download['url']) 
                            
                else:
                    # TODO: manage cases where multiple tiles are downloaded for pixel mode (add mode argument)
                    # Get all available downloads and pack them in args
                    save_product_path = download_path + os.sep + 'USGS'
                    args = [(download['url'], save_product_path, chunk_size) for download in requestResults['availableDownloads']]

                    # Get cpu count
                    nb_cores = min([max_cpu, cpu_count(logical = False), len(os.sched_getaffinity(0))])

                    print('Starting LandSat download with %d cores for %d images...\n' %(nb_cores, len(downloads)))

                    # Start pool and get results
                    results = p_map(download_file_multiprocess, args, **{"num_cpus": nb_cores})
                    
                    # Collect results and sort them
                    landsat_products = []
                    for result in results:
                        landsat_products.append(result)
                    landsat_products.sort()

                print("\nAll products are available to download.\n")
        else:
            print("Search found no results.\n")
                
    # Logout so the API Key cannot be used anymore
    endpoint = "logout"  
    if sendRequest(serviceUrl + endpoint, None, apiKey) == None:        
        print("\nLogged Out\n")
    else:
        print("\nLogout Failed\n")
    
    # Save list of paths as a csv file for later use
    with open(save_path, 'w', newline='') as f:
        # using csv.writer method from CSV package
        write = csv.writer(f)

        for product in landsat_products:
            write.writerow([product])

    return landsat_products