fastai--fastai
wehub-resource-sync
107 行
4.4 KiB
Python
107 行
4.4 KiB
Python
"""Some utils function to quickly download a bunch of images, check them and pre-resize them
|
|
|
|
Docs: https://docs.fast.ai/vision.utils.html.md"""
|
|
|
|
# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/09b_vision.utils.ipynb.
|
|
|
|
# %% auto #0
|
|
__all__ = ['download_images', 'resize_to', 'verify_image', 'verify_images', 'resize_image', 'resize_images']
|
|
|
|
# %% ../../nbs/09b_vision.utils.ipynb #a4ba4284
|
|
import uuid
|
|
from ..torch_basics import *
|
|
from ..data.all import *
|
|
from .core import *
|
|
from fastdownload import download_url
|
|
from pathlib import Path
|
|
|
|
# %% ../../nbs/09b_vision.utils.ipynb #f534f4c7
|
|
def _get_downloaded_image_filename(dest, name, suffix):
|
|
start_index = 1
|
|
candidate_name = name
|
|
|
|
while (dest/f"{candidate_name}{suffix}").is_file():
|
|
candidate_name = f"{candidate_name}{start_index}"
|
|
start_index += 1
|
|
|
|
return candidate_name
|
|
|
|
# %% ../../nbs/09b_vision.utils.ipynb #e94c0aa9
|
|
def _download_image_inner(dest, inp, timeout=4, preserve_filename=False):
|
|
i,url = inp
|
|
url = url.split("?")[0]
|
|
url_path = Path(url)
|
|
suffix = url_path.suffix if url_path.suffix else '.jpg'
|
|
name = _get_downloaded_image_filename(dest, url_path.stem, suffix) if preserve_filename else str(uuid.uuid4())
|
|
try: download_url(url, dest/f"{name}{suffix}", show_progress=False, timeout=timeout)
|
|
except Exception as e: f"Couldn't download {url}."
|
|
|
|
# %% ../../nbs/09b_vision.utils.ipynb #d0593ffc
|
|
def download_images(dest, url_file=None, urls=None, max_pics=1000, n_workers=8, timeout=4, preserve_filename=False):
|
|
"Download images listed in text file `url_file` to path `dest`, at most `max_pics`"
|
|
if urls is None: urls = url_file.read_text().strip().split("\n")[:max_pics]
|
|
dest = Path(dest)
|
|
dest.mkdir(exist_ok=True)
|
|
parallel(partial(_download_image_inner, dest, timeout=timeout, preserve_filename=preserve_filename),
|
|
list(enumerate(urls)), n_workers=n_workers, threadpool=True)
|
|
|
|
# %% ../../nbs/09b_vision.utils.ipynb #41c04937
|
|
def resize_to(img, targ_sz, use_min=False):
|
|
"Size to resize to, to hit `targ_sz` at same aspect ratio, in PIL coords (i.e w*h)"
|
|
w,h = img.size
|
|
min_sz = (min if use_min else max)(w,h)
|
|
ratio = targ_sz/min_sz
|
|
return int(w*ratio),int(h*ratio)
|
|
|
|
# %% ../../nbs/09b_vision.utils.ipynb #3bff393f
|
|
def verify_image(fn):
|
|
"Confirm that `fn` can be opened"
|
|
try:
|
|
im = Image.open(fn)
|
|
im.draft(im.mode, (32,32))
|
|
im.load()
|
|
return True
|
|
except: return False
|
|
|
|
# %% ../../nbs/09b_vision.utils.ipynb #f418bded
|
|
def verify_images(fns):
|
|
"Find images in `fns` that can't be opened"
|
|
return L(fns[i] for i,o in enumerate(parallel(verify_image, fns)) if not o)
|
|
|
|
# %% ../../nbs/09b_vision.utils.ipynb #aad88d0a
|
|
def resize_image(file, dest, src='.', max_size=None, n_channels=3, ext=None,
|
|
img_format=None, resample=BILINEAR, resume=False, **kwargs ):
|
|
"Resize file to dest to max_size"
|
|
dest = Path(dest)
|
|
|
|
dest_fname = dest/file
|
|
dest_fname.parent.mkdir(exist_ok=True, parents=True)
|
|
file = Path(src)/file
|
|
if resume and dest_fname.exists(): return
|
|
if not verify_image(file): return
|
|
|
|
img = Image.open(file)
|
|
imgarr = np.array(img)
|
|
img_channels = 1 if len(imgarr.shape) == 2 else imgarr.shape[2]
|
|
if ext is not None: dest_fname=dest_fname.with_suffix(ext)
|
|
if (max_size is not None and (img.height > max_size or img.width > max_size)) or img_channels != n_channels:
|
|
if max_size is not None:
|
|
new_sz = resize_to(img, max_size)
|
|
img = img.resize(new_sz, resample=resample)
|
|
if n_channels == 3: img = img.convert("RGB")
|
|
img.save(dest_fname, img_format, **kwargs)
|
|
elif file != dest_fname : shutil.copy2(file, dest_fname)
|
|
|
|
# %% ../../nbs/09b_vision.utils.ipynb #e9103784
|
|
def resize_images(path, max_workers=defaults.cpus, max_size=None, recurse=False,
|
|
dest=Path('.'), n_channels=3, ext=None, img_format=None, resample=BILINEAR,
|
|
resume=None, **kwargs):
|
|
"Resize files on path recursively to dest to max_size"
|
|
path = Path(path)
|
|
if resume is None and dest != Path('.'): resume=False
|
|
os.makedirs(dest, exist_ok=True)
|
|
files = get_image_files(path, recurse=recurse)
|
|
files = [o.relative_to(path) for o in files]
|
|
parallel(resize_image, files, src=path, n_workers=max_workers, max_size=max_size, dest=dest, n_channels=n_channels, ext=ext,
|
|
img_format=img_format, resample=resample, resume=resume, **kwargs)
|