<reponame>metal3d/imgnet-tools<gh_stars>0
#! /bin/env python3
""" Helper script to download image-net.org image corresponding
to a given ID.

Author: <NAME> <<EMAIL>>
Licence: MIT
"""
import requests
import os
import hashlib
import imghdr
import io
import queue
import threading
import multiprocessing
import sys

CPUs = multiprocessing.cpu_count()

# LIST_URL =
# 'http://image-net.org/api/text/imagenet.synset.geturls?wnid={imid}'

LIST_URL = ''.join([
    'http://www.image-net.org/',
    'api/text/imagenet.synset.geturls.getmapping?wnid={imid}'])

SEE_URL = 'http://www.image-net.org/synset?wnid={id}'

# not wanted images, mainly for flickr bad file, hugedomain logos...
BADIMG = [
    '880a7a58e05d3e83797f27573bb6d35c',  # flickr
    '596246739a83bb45e30e13437e0810d9',  # warning sign
    '969e595182a947c7fdaaef4a34401760',  # forbiden sign
    'af5db09e39ca35d8930b4e59962e09e5',  # hugedomain logo
]
TIMEOUT = 5
DEST = './'
DATAFILE = './data.csv'


q = None
locker = threading.Lock()


def init_queue(number: int = CPUs):
    """ Initiate Queue to that number of elements """
    global q
    if q is None:
        print("Initializing Queue to %d workers" % number)
        q = queue.Queue(number)
    else:
        print('Queue already initialized')


def get_list(imid: str) -> requests.Response:
    """ Return the requests.Response containing
        the list of images for a given image-net.org
        collection ID.
    """
    imlist = requests.get(LIST_URL.format(imid=imid))
    return imlist


def logthat(content, total=None, index=None):
    """ Helper function to write logs
        with progression status.
    """
    if total is not None and index is not None:
        print('%d/%d %s' % (index+1, total, content))
    else:
        print('%s' % content)


def save_data(imname: str, url: str, classname: str, nid: str):
    locker.acquire()

    if not os.path.exists(DATAFILE):
        try:
            # append header line
            with open(DATAFILE, 'a') as f:
                f.write('Sysnet Name,Base URL,Classname,Imagenet ID\n')
        except Exception as e:
            print(e)
            sys.exit(0)

    try:
        with open(DATAFILE, 'a') as f:
            f.write('"%s",%s,"%s",%s\n' % (imname, url, classname, nid))
    except Exception as e:
        print(e)
    finally:
        locker.release()


def is_in_db(imname: str):
    """ Fetch given url in CSV file and return boolean
        saying if the file is already downloaded.
    """
    locker.acquire()
    try:
        lines = open(DATAFILE).readlines()
    except Exception:
        return False
    finally:
        locker.release()

    for l in lines:
        if imname in l:
            return True
    return False


def dl_image(
        imname: str,
        imurl: str,
        classname='unknown',
        dest='./',
        total=None,
        index=None,
        nid=None):
    """ Download image from the given url, save it as classname
        and check if the image is correct.
    """
    # imname = os.path.basename(imurl)
    # imname = imname.decode()

    # prepare the desitnation directory
    os.makedirs(os.path.join(dest, classname), exist_ok=True)

    # deactivate, use md5
    # if os.path.exists(fileto):
    #    print('File %s already downloaded, skippiing' % fileto)
    #    return

    # check if image is in csv database
    if is_in_db(imname):
        logthat('Image already downloaded, skipping', total, index)
        return

    # download image, and check status_code
    try:
        im = requests.get(imurl, timeout=TIMEOUT)
        if im.status_code != 200:
            logthat(
                'Status code is not OK for %s, %d' % (
                    imurl, im.status_code),
                total,
                index
            )
            return
    except Exception as e:
        logthat(
            'Error, connot download %s %s' % (imurl, e),
            total,
            index
        )
        return

    # is it a valid image ?
    b = io.BytesIO()
    b.write(im.content)
    b.seek(0)
    ext = imghdr.what(b)
    if ext is None:
        logthat(
            "%s seems to not be a valid image file, skipping" % imname,
            total,
            index
        )
        return

    # is it a bad image ?
    md5 = hashlib.md5(im.content).hexdigest()
    if md5 in BADIMG:
        logthat(
            imname + ' md5 corresponds to a bad image, skipping',
            total,
            index)
        return

    fileto = os.path.join(dest, classname, imname + '.' + ext)
    # md5 is ok, downloaded file is not empty, we can save it
    with open(fileto, 'wb') as f:
        f.write(im.content)
        save_data(imname, imurl, classname, nid)
        logthat("%s file saved" % imname, total, index)


def task_download():
    """ Helper function that is launched by threading.Thread.
        It reads Queue and call dl_image() in parallel.
    """
    init_queue()
    while True:
        imname, item, classname, total, index, nid = q.get()
        if item is None:
            break
        dl_image(imname, item, classname, DEST, total, index, nid)
        q.task_done()


if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser(
            formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument(
        'nid',
        help='Id of the imagenet group to download. '
             'It can be a coma separated of ids. It\'s an optionnal '
             'argument, if leaved empty, so a search is made on imagenet '
             'with the class name argument and nothing is downloaded.',
        nargs='?'
    )
    parser.add_argument(
        'name',
        help='Real classname of the images to '
             'download (it will create the corresponding directory name)'
    )
    parser.add_argument(
        '-d',
        '--dest',
        help='destination root (without classname) directory '
             'where images will be downloaded',
        default=DEST
    )
    parser.add_argument(
        '-t',
        '--timeout',
        type=int,
        help='Timeout in seconds for requests before to abandon download',
        default=TIMEOUT
    )
    parser.add_argument(
        '-n',
        '--num-worker',
        type=int,
        help='Number of worker for parallel '
             'download, use number of cpus by default',
        default=CPUs
    )
    parser.add_argument(
        '-c',
        '--csv',
        type=str,
        help='CSV file where to keep downloaded file information',
        default=DATAFILE
    )

    args = parser.parse_args()

    nids = args.nid
    classname = args.name
    NUM_WORKERS = args.num_worker
    TIMEOUT = args.timeout
    DEST = args.dest
    DATAFILE = args.csv

    if args.nid is None:
        # get words
        nids = []
        words = requests.get('http://image-net.org/archive/words.txt')
        print('You didn\'t provide ID, choose in the list')
        for line in words.iter_lines():
            nid, terms = line.decode().split('\t')
            for n in classname.split(','):
                if n in terms:
                    print(
                        nid,
                        terms,
                        "See results at: %s" % SEE_URL.format(id=nid)
                    )
        sys.exit(0)

    init_queue(NUM_WORKERS)

    threads = []
    for i in range(NUM_WORKERS):
        t = threading.Thread(target=task_download)
        t.start()
        threads.append(t)

    lines_count = 1
    results = {}
    for nid in nids.split(','):
        nid = nid.strip()
        res = get_list(nid)
        # list is not very well formatted, we need
        # to iterate to count, and not find "\n" or whatever
        # the line delimiter
        for _ in res.iter_lines():
            lines_count += 1

        results[nid] = res

    lastidx = 0
    for nid, res in results.items():
        for u in res.iter_lines():
            u = u.decode().split(' ')
            if len(u) == 2:
                imname, u = list(u)
            else:
                continue
            q.put((imname, u, classname, lines_count, lastidx, nid))
            lastidx += 1

    q.join()

    for i in range(NUM_WORKERS):
        q.put((None, None, None, None, None, None))

    for t in threads:
        t.join()

#! /bin/env python3
""" Helper script to download image-net.org image corresponding
to a given ID.

Author: <NAME> <<EMAIL>>
Licence: MIT
"""
import requests
import os
import hashlib
import imghdr
import io
import queue
import threading
import multiprocessing
import sys

CPUs = multiprocessing.cpu_count()

# LIST_URL =
# 'http://image-net.org/api/text/imagenet.synset.geturls?wnid={imid}'

LIST_URL = ''.join([
    'http://www.image-net.org/',
    'api/text/imagenet.synset.geturls.getmapping?wnid={imid}'])

SEE_URL = 'http://www.image-net.org/synset?wnid={id}'

# not wanted images, mainly for flickr bad file, hugedomain logos...
BADIMG = [
    '880a7a58e05d3e83797f27573bb6d35c',  # flickr
    '596246739a83bb45e30e13437e0810d9',  # warning sign
    '969e595182a947c7fdaaef4a34401760',  # forbiden sign
    'af5db09e39ca35d8930b4e59962e09e5',  # hugedomain logo
]
TIMEOUT = 5
DEST = './'
DATAFILE = './data.csv'


q = None
locker = threading.Lock()


def init_queue(number: int = CPUs):
    """ Initiate Queue to that number of elements """
    global q
    if q is None:
        print("Initializing Queue to %d workers" % number)
        q = queue.Queue(number)
    else:
        print('Queue already initialized')


def get_list(imid: str) -> requests.Response:
    """ Return the requests.Response containing
        the list of images for a given image-net.org
        collection ID.
    """
    imlist = requests.get(LIST_URL.format(imid=imid))
    return imlist


def logthat(content, total=None, index=None):
    """ Helper function to write logs
        with progression status.
    """
    if total is not None and index is not None:
        print('%d/%d %s' % (index+1, total, content))
    else:
        print('%s' % content)


def save_data(imname: str, url: str, classname: str, nid: str):
    locker.acquire()

    if not os.path.exists(DATAFILE):
        try:
            # append header line
            with open(DATAFILE, 'a') as f:
                f.write('Sysnet Name,Base URL,Classname,Imagenet ID\n')
        except Exception as e:
            print(e)
            sys.exit(0)

    try:
        with open(DATAFILE, 'a') as f:
            f.write('"%s",%s,"%s",%s\n' % (imname, url, classname, nid))
    except Exception as e:
        print(e)
    finally:
        locker.release()


def is_in_db(imname: str):
    """ Fetch given url in CSV file and return boolean
        saying if the file is already downloaded.
    """
    locker.acquire()
    try:
        lines = open(DATAFILE).readlines()
    except Exception:
        return False
    finally:
        locker.release()

    for l in lines:
        if imname in l:
            return True
    return False


def dl_image(
        imname: str,
        imurl: str,
        classname='unknown',
        dest='./',
        total=None,
        index=None,
        nid=None):
    """ Download image from the given url, save it as classname
        and check if the image is correct.
    """
    # imname = os.path.basename(imurl)
    # imname = imname.decode()

    # prepare the desitnation directory
    os.makedirs(os.path.join(dest, classname), exist_ok=True)

    # deactivate, use md5
    # if os.path.exists(fileto):
    #    print('File %s already downloaded, skippiing' % fileto)
    #    return

    # check if image is in csv database
    if is_in_db(imname):
        logthat('Image already downloaded, skipping', total, index)
        return

    # download image, and check status_code
    try:
        im = requests.get(imurl, timeout=TIMEOUT)
        if im.status_code != 200:
            logthat(
                'Status code is not OK for %s, %d' % (
                    imurl, im.status_code),
                total,
                index
            )
            return
    except Exception as e:
        logthat(
            'Error, connot download %s %s' % (imurl, e),
            total,
            index
        )
        return

    # is it a valid image ?
    b = io.BytesIO()
    b.write(im.content)
    b.seek(0)
    ext = imghdr.what(b)
    if ext is None:
        logthat(
            "%s seems to not be a valid image file, skipping" % imname,
            total,
            index
        )
        return

    # is it a bad image ?
    md5 = hashlib.md5(im.content).hexdigest()
    if md5 in BADIMG:
        logthat(
            imname + ' md5 corresponds to a bad image, skipping',
            total,
            index)
        return

    fileto = os.path.join(dest, classname, imname + '.' + ext)
    # md5 is ok, downloaded file is not empty, we can save it
    with open(fileto, 'wb') as f:
        f.write(im.content)
        save_data(imname, imurl, classname, nid)
        logthat("%s file saved" % imname, total, index)


def task_download():
    """ Helper function that is launched by threading.Thread.
        It reads Queue and call dl_image() in parallel.
    """
    init_queue()
    while True:
        imname, item, classname, total, index, nid = q.get()
        if item is None:
            break
        dl_image(imname, item, classname, DEST, total, index, nid)
        q.task_done()


if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser(
            formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument(
        'nid',
        help='Id of the imagenet group to download. '
             'It can be a coma separated of ids. It\'s an optionnal '
             'argument, if leaved empty, so a search is made on imagenet '
             'with the class name argument and nothing is downloaded.',
        nargs='?'
    )
    parser.add_argument(
        'name',
        help='Real classname of the images to '
             'download (it will create the corresponding directory name)'
    )
    parser.add_argument(
        '-d',
        '--dest',
        help='destination root (without classname) directory '
             'where images will be downloaded',
        default=DEST
    )
    parser.add_argument(
        '-t',
        '--timeout',
        type=int,
        help='Timeout in seconds for requests before to abandon download',
        default=TIMEOUT
    )
    parser.add_argument(
        '-n',
        '--num-worker',
        type=int,
        help='Number of worker for parallel '
             'download, use number of cpus by default',
        default=CPUs
    )
    parser.add_argument(
        '-c',
        '--csv',
        type=str,
        help='CSV file where to keep downloaded file information',
        default=DATAFILE
    )

    args = parser.parse_args()

    nids = args.nid
    classname = args.name
    NUM_WORKERS = args.num_worker
    TIMEOUT = args.timeout
    DEST = args.dest
    DATAFILE = args.csv

    if args.nid is None:
        # get words
        nids = []
        words = requests.get('http://image-net.org/archive/words.txt')
        print('You didn\'t provide ID, choose in the list')
        for line in words.iter_lines():
            nid, terms = line.decode().split('\t')
            for n in classname.split(','):
                if n in terms:
                    print(
                        nid,
                        terms,
                        "See results at: %s" % SEE_URL.format(id=nid)
                    )
        sys.exit(0)

    init_queue(NUM_WORKERS)

    threads = []
    for i in range(NUM_WORKERS):
        t = threading.Thread(target=task_download)
        t.start()
        threads.append(t)

    lines_count = 1
    results = {}
    for nid in nids.split(','):
        nid = nid.strip()
        res = get_list(nid)
        # list is not very well formatted, we need
        # to iterate to count, and not find "\n" or whatever
        # the line delimiter
        for _ in res.iter_lines():
            lines_count += 1

        results[nid] = res

    lastidx = 0
    for nid, res in results.items():
        for u in res.iter_lines():
            u = u.decode().split(' ')
            if len(u) == 2:
                imname, u = list(u)
            else:
                continue
            q.put((imname, u, classname, lines_count, lastidx, nid))
            lastidx += 1

    q.join()

    for i in range(NUM_WORKERS):
        q.put((None, None, None, None, None, None))

    for t in threads:
        t.join()
#! /bin/env python3 Helper script to download image-net.org image corresponding to a given ID.  Author: <NAME> <<EMAIL>> Licence: MIT # LIST_URL = # 'http://image-net.org/api/text/imagenet.synset.geturls?wnid={imid}' # not wanted images, mainly for flickr bad file, hugedomain logos... # flickr # warning sign # forbiden sign # hugedomain logo Initiate Queue to that number of elements Return the requests.Response containing         the list of images for a given image-net.org         collection ID. Helper function to write logs         with progression status. # append header line Fetch given url in CSV file and return boolean         saying if the file is already downloaded. Download image from the given url, save it as classname         and check if the image is correct. # imname = os.path.basename(imurl) # imname = imname.decode() # prepare the desitnation directory # deactivate, use md5 # if os.path.exists(fileto): #    print('File %s already downloaded, skippiing' % fileto) #    return # check if image is in csv database # download image, and check status_code # is it a valid image ? # is it a bad image ? # md5 is ok, downloaded file is not empty, we can save it Helper function that is launched by threading.Thread.         It reads Queue and call dl_image() in parallel. # get words # list is not very well formatted, we need # to iterate to count, and not find "\n" or whatever # the line delimiter