Spaces:
Sleeping
Sleeping
File size: 2,807 Bytes
bbae066 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 |
"""Utility functions for image_to_fen module."""
import base64
import contextlib
import hashlib
from io import BytesIO
import os
from pathlib import Path
from typing import Union
from urllib.request import urlretrieve
import numpy as np
from PIL import Image
import smart_open
from tqdm import tqdm
def to_categorical(y, num_classes):
"""1-hot encode a tensor."""
return np.eye(num_classes, dtype="uint8")[y]
def read_image_pil(image_uri: Union[Path, str], grayscale=False) -> Image:
with smart_open.open(image_uri, "rb") as image_file:
return read_image_pil_file(image_file, grayscale)
def read_image_pil_file(image_file, grayscale=False) -> Image:
with Image.open(image_file) as image:
if grayscale:
image = image.convert(mode="L")
else:
image = image.convert(mode=image.mode)
return image
@contextlib.contextmanager
def temporary_working_directory(working_dir: Union[str, Path]):
"""Temporarily switches to a directory, then returns to the original directory on exit."""
curdir = os.getcwd()
os.chdir(working_dir)
try:
yield
finally:
os.chdir(curdir)
def encode_b64_image(image, format="png"):
"""Encode a PIL image as a base64 string."""
_buffer = BytesIO() # bytes that live in memory
image.save(_buffer, format=format) # but which we write to like a file
encoded_image = base64.b64encode(_buffer.getvalue()).decode("utf8")
return encoded_image
def compute_sha256(filename: Union[Path, str]):
"""Return SHA256 checksum of a file."""
with open(filename, "rb") as f:
return hashlib.sha256(f.read()).hexdigest()
class TqdmUpTo(tqdm):
"""From https://github.com/tqdm/tqdm/blob/master/examples/tqdm_wget.py"""
def update_to(self, blocks=1, bsize=1, tsize=None):
"""
Parameters
----------
blocks: int, optional
Number of blocks transferred so far [default: 1].
bsize: int, optional
Size of each block (in tqdm units) [default: 1].
tsize: int, optional
Total size (in tqdm units). If [default: None] remains unchanged.
"""
if tsize is not None:
self.total = tsize
self.update(blocks * bsize - self.n) # will also set self.n = b * bsize
def download_url(url, filename):
"""Download a file from url to filename, with a progress bar."""
with TqdmUpTo(unit="B", unit_scale=True, unit_divisor=1024, miniters=1) as t:
urlretrieve(url, filename, reporthook=t.update_to, data=None) # noqa: S310
# the function takes the original prediction and the iou threshold.
# function to convert a torchtensor back to PIL image
def torch_to_pil(img):
return torchvision.transforms.ToPILImage()(img).convert('RGB') |