load_data.py

import os
import requests

# Downloads the data for paper "Automated Structure Discovery for Scanning Tunneling Microscopy"
band_url = 'https://zenodo.org/records/10687786/files/band.h5?download=1'
rot_url = 'https://zenodo.org/records/10687786/files/rotations_210611.pickle?download=1'
model_url = 'https://zenodo.org/records/11035026/files/disks.pt?download=1'

def download_data():

    if not os.path.exists('data'):
        os.makedirs('data')

    band_path = 'data/band.h5'
    rot_path = 'data/rotations_210611.pickle'
    model_path = 'pretrained/disks.pt'

    if not os.path.exists(band_path):
        print('Downloading band data...')
        r = requests.get(band_url, allow_redirects=True)
        open(band_path, 'wb').write(r.content)

    if not os.path.exists(rot_path):
        print('Downloading rotations...')
        r = requests.get(rot_url, allow_redirects=True)
        open(rot_path, 'wb').write(r.content)

    if not os.path.exists(model_path):
        print('Downloading model...')
        r = requests.get(model_url, allow_redirects=True)
        open(model_path, 'wb').write(r.content)

    print('Data and model downloaded.')

    return band_path, rot_path, model_path

if __name__=='__main__':
    band_path, rot_path, model_path = download_data()