Spaces:
Running on Zero
Running on Zero
| # ref https://github.com/scardine/image_size/blob/master/get_image_size.py | |
| import atexit | |
| import collections | |
| import json | |
| import os | |
| import io | |
| import struct | |
| import threading | |
| from typing import TYPE_CHECKING | |
| import cv2 | |
| import numpy as np | |
| import torch | |
| from diffusers import AutoencoderTiny | |
| from PIL import Image as PILImage | |
| FILE_UNKNOWN = "Sorry, don't know how to get size for this file." | |
| class UnknownImageFormat(Exception): | |
| pass | |
| types = collections.OrderedDict() | |
| BMP = types['BMP'] = 'BMP' | |
| GIF = types['GIF'] = 'GIF' | |
| ICO = types['ICO'] = 'ICO' | |
| JPEG = types['JPEG'] = 'JPEG' | |
| PNG = types['PNG'] = 'PNG' | |
| TIFF = types['TIFF'] = 'TIFF' | |
| WEBP = types['WEBP'] = 'WEBP' | |
| # EXIF orientations 5-8 are transposed: the displayed image swaps width/height | |
| _EXIF_SWAP_ORIENTATIONS = (5, 6, 7, 8) | |
| # SOF0-15 minus DHT (C4), JPG (C8), DAC (CC) | |
| _JPEG_SOF_MARKERS = frozenset(range(0xC0, 0xD0)) - {0xC4, 0xC8, 0xCC} | |
| # markers with no length field: TEM, RST0-7, SOI | |
| _JPEG_STANDALONE_MARKERS = frozenset(range(0xD0, 0xD8)) | {0x01, 0xD8} | |
| def _exif_orientation(exif): | |
| """Orientation tag (0x0112) from a TIFF-format EXIF blob, or None.""" | |
| if exif.startswith(b"Exif\x00\x00"): | |
| exif = exif[6:] | |
| if len(exif) < 8 or exif[:2] not in (b"II", b"MM"): | |
| return None | |
| bo = "<" if exif[:2] == b"II" else ">" | |
| try: | |
| ifd = struct.unpack(bo + "I", exif[4:8])[0] | |
| count = struct.unpack(bo + "H", exif[ifd:ifd + 2])[0] | |
| for i in range(count): | |
| off = ifd + 2 + i * 12 | |
| tag, typ, n = struct.unpack(bo + "HHI", exif[off:off + 8]) | |
| if tag == 0x0112: | |
| if typ != 3 or n != 1: | |
| return None | |
| return struct.unpack(bo + "H", exif[off + 8:off + 10])[0] | |
| except struct.error: | |
| return None | |
| return None | |
| def _png_exif_orientation(input): | |
| """Walk PNG chunks after IHDR for eXIf; it must precede IDAT, so stop there.""" | |
| pos = 33 # signature + IHDR chunk | |
| try: | |
| while True: | |
| input.seek(pos) | |
| head = input.read(8) | |
| if len(head) < 8: | |
| return None | |
| length, ctype = struct.unpack(">I4s", head) | |
| if ctype in (b"IDAT", b"IEND"): | |
| return None | |
| if ctype == b"eXIf": | |
| return _exif_orientation(input.read(length)) | |
| pos += 12 + length | |
| except struct.error: | |
| return None | |
| def _webp_exif_orientation(input, size, pos): | |
| """Walk RIFF chunks from pos for EXIF; a few seeks, never the bitstream.""" | |
| try: | |
| while pos + 8 <= size: | |
| input.seek(pos) | |
| fourcc, length = struct.unpack("<4sI", input.read(8)) | |
| if fourcc == b"EXIF": | |
| return _exif_orientation(input.read(length)) | |
| pos += 8 + length + (length & 1) | |
| except struct.error: | |
| return None | |
| return None | |
| image_fields = ['path', 'type', 'file_size', 'width', 'height'] | |
| class Image(collections.namedtuple('Image', image_fields)): | |
| def to_str_row(self): | |
| return ("%d\t%d\t%d\t%s\t%s" % ( | |
| self.width, | |
| self.height, | |
| self.file_size, | |
| self.type, | |
| self.path.replace('\t', '\\t'), | |
| )) | |
| def to_str_row_verbose(self): | |
| return ("%d\t%d\t%d\t%s\t%s\t##%s" % ( | |
| self.width, | |
| self.height, | |
| self.file_size, | |
| self.type, | |
| self.path.replace('\t', '\\t'), | |
| self)) | |
| def to_str_json(self, indent=None): | |
| return json.dumps(self._asdict(), indent=indent) | |
| def get_image_size(file_path): | |
| """ | |
| Return (width, height) for a given img file content - no external | |
| dependencies except the os and struct builtin modules | |
| """ | |
| img = get_image_metadata(file_path) | |
| return (img.width, img.height) | |
| def get_image_size_from_bytesio(input, size): | |
| """ | |
| Return (width, height) for a given img file content - no external | |
| dependencies except the os and struct builtin modules | |
| Args: | |
| input (io.IOBase): io object support read & seek | |
| size (int): size of buffer in byte | |
| """ | |
| img = get_image_metadata_from_bytesio(input, size) | |
| return (img.width, img.height) | |
| def get_image_metadata(file_path): | |
| """ | |
| Return an `Image` object for a given img file content - no external | |
| dependencies except the os and struct builtin modules | |
| Args: | |
| file_path (str): path to an image file | |
| Returns: | |
| Image: (path, type, file_size, width, height) | |
| """ | |
| size = os.path.getsize(file_path) | |
| # be explicit with open arguments - we need binary mode | |
| with io.open(file_path, "rb") as input: | |
| return get_image_metadata_from_bytesio(input, size, file_path) | |
| def get_image_metadata_from_bytesio(input, size, file_path=None): | |
| """ | |
| Return an `Image` object for a given img file content - no external | |
| dependencies except the os and struct builtin modules | |
| Width/height honor the EXIF orientation tag (JPEG APP1, PNG eXIf, WebP | |
| EXIF chunk), so they match `PIL.ImageOps.exif_transpose(...).size`. | |
| Args: | |
| input (io.IOBase): io object support read & seek | |
| size (int): size of buffer in byte | |
| file_path (str): path to an image file | |
| Returns: | |
| Image: (path, type, file_size, width, height) | |
| """ | |
| height = -1 | |
| width = -1 | |
| orientation = None | |
| data = input.read(30) | |
| msg = " raised while trying to decode as JPEG." | |
| if (size >= 10) and data[:6] in (b'GIF87a', b'GIF89a'): | |
| # GIFs | |
| imgtype = GIF | |
| w, h = struct.unpack("<HH", data[6:10]) | |
| width = int(w) | |
| height = int(h) | |
| elif ((size >= 24) and data.startswith(b'\211PNG\r\n\032\n') | |
| and (data[12:16] == b'IHDR')): | |
| # PNGs | |
| imgtype = PNG | |
| w, h = struct.unpack(">LL", data[16:24]) | |
| width = int(w) | |
| height = int(h) | |
| orientation = _png_exif_orientation(input) | |
| elif (size >= 2) and data.startswith(b'\377\330'): | |
| # JPEG | |
| imgtype = JPEG | |
| input.seek(2) | |
| try: | |
| while True: | |
| b = input.read(1) | |
| if not b: | |
| raise UnknownImageFormat("no SOF marker" + msg) | |
| if b != b'\xff': | |
| continue | |
| marker = input.read(1) | |
| while marker == b'\xff': # fill bytes | |
| marker = input.read(1) | |
| if not marker: | |
| raise UnknownImageFormat("no SOF marker" + msg) | |
| m = ord(marker) | |
| if m in _JPEG_STANDALONE_MARKERS: | |
| continue | |
| if m in (0xD9, 0xDA): # EOI / SOS before any SOF | |
| raise UnknownImageFormat("no SOF marker" + msg) | |
| seg_len = struct.unpack(">H", input.read(2))[0] - 2 | |
| if m in _JPEG_SOF_MARKERS: | |
| input.read(1) # sample precision | |
| h, w = struct.unpack(">HH", input.read(4)) | |
| break | |
| if m == 0xE1 and orientation is None: | |
| # first Exif APP1 wins, same as PIL | |
| orientation = _exif_orientation(input.read(seg_len)) | |
| else: | |
| input.seek(seg_len, 1) | |
| width = int(w) | |
| height = int(h) | |
| except UnknownImageFormat: | |
| raise | |
| except struct.error: | |
| raise UnknownImageFormat("StructError" + msg) | |
| except ValueError: | |
| raise UnknownImageFormat("ValueError" + msg) | |
| except Exception as e: | |
| raise UnknownImageFormat(e.__class__.__name__ + msg) | |
| elif (size >= 30) and data[:4] == b'RIFF' and data[8:12] == b'WEBP': | |
| # WebP: RIFF header, then first chunk fourcc at 12, chunk size at 16, payload at 20 | |
| imgtype = WEBP | |
| chunk = data[12:16] | |
| if chunk == b'VP8 ': | |
| # lossy: 3-byte frame tag, 3-byte start code, then 14-bit width/height | |
| if data[23:26] != b'\x9d\x01\x2a': | |
| raise UnknownImageFormat("bad VP8 start code") | |
| w, h = struct.unpack("<HH", data[26:30]) | |
| width = int(w & 0x3FFF) | |
| height = int(h & 0x3FFF) | |
| elif chunk == b'VP8L': | |
| # lossless: signature byte, then 14-bit width-1 / height-1 packed LE | |
| if data[20] != 0x2F: | |
| raise UnknownImageFormat("bad VP8L signature") | |
| bits = struct.unpack("<I", data[21:25])[0] | |
| width = (bits & 0x3FFF) + 1 | |
| height = ((bits >> 14) & 0x3FFF) + 1 | |
| elif chunk == b'VP8X': | |
| # extended: flags byte, 3 reserved, then 24-bit canvas width-1 / height-1 | |
| width = int.from_bytes(data[24:27], 'little') + 1 | |
| height = int.from_bytes(data[27:30], 'little') + 1 | |
| if data[20] & 0x08: # EXIF flag | |
| chunk_size = struct.unpack("<I", data[16:20])[0] | |
| orientation = _webp_exif_orientation(input, size, 20 + chunk_size + (chunk_size & 1)) | |
| else: | |
| raise UnknownImageFormat("Unknown WebP chunk: " + repr(chunk)) | |
| elif (size >= 26) and data.startswith(b'BM'): | |
| # BMP | |
| imgtype = 'BMP' | |
| headersize = struct.unpack("<I", data[14:18])[0] | |
| if headersize == 12: | |
| w, h = struct.unpack("<HH", data[18:22]) | |
| width = int(w) | |
| height = int(h) | |
| elif headersize >= 40: | |
| w, h = struct.unpack("<ii", data[18:26]) | |
| width = int(w) | |
| # as h is negative when stored upside down | |
| height = abs(int(h)) | |
| else: | |
| raise UnknownImageFormat( | |
| "Unkown DIB header size:" + | |
| str(headersize)) | |
| elif (size >= 8) and data[:4] in (b"II\052\000", b"MM\000\052"): | |
| # Standard TIFF, big- or little-endian | |
| # BigTIFF and other different but TIFF-like formats are not | |
| # supported currently | |
| imgtype = TIFF | |
| byteOrder = data[:2] | |
| boChar = ">" if byteOrder == "MM" else "<" | |
| # maps TIFF type id to size (in bytes) | |
| # and python format char for struct | |
| tiffTypes = { | |
| 1: (1, boChar + "B"), # BYTE | |
| 2: (1, boChar + "c"), # ASCII | |
| 3: (2, boChar + "H"), # SHORT | |
| 4: (4, boChar + "L"), # LONG | |
| 5: (8, boChar + "LL"), # RATIONAL | |
| 6: (1, boChar + "b"), # SBYTE | |
| 7: (1, boChar + "c"), # UNDEFINED | |
| 8: (2, boChar + "h"), # SSHORT | |
| 9: (4, boChar + "l"), # SLONG | |
| 10: (8, boChar + "ll"), # SRATIONAL | |
| 11: (4, boChar + "f"), # FLOAT | |
| 12: (8, boChar + "d") # DOUBLE | |
| } | |
| ifdOffset = struct.unpack(boChar + "L", data[4:8])[0] | |
| try: | |
| countSize = 2 | |
| input.seek(ifdOffset) | |
| ec = input.read(countSize) | |
| ifdEntryCount = struct.unpack(boChar + "H", ec)[0] | |
| # 2 bytes: TagId + 2 bytes: type + 4 bytes: count of values + 4 | |
| # bytes: value offset | |
| ifdEntrySize = 12 | |
| for i in range(ifdEntryCount): | |
| entryOffset = ifdOffset + countSize + i * ifdEntrySize | |
| input.seek(entryOffset) | |
| tag = input.read(2) | |
| tag = struct.unpack(boChar + "H", tag)[0] | |
| if (tag == 256 or tag == 257): | |
| # if type indicates that value fits into 4 bytes, value | |
| # offset is not an offset but value itself | |
| type = input.read(2) | |
| type = struct.unpack(boChar + "H", type)[0] | |
| if type not in tiffTypes: | |
| raise UnknownImageFormat( | |
| "Unkown TIFF field type:" + | |
| str(type)) | |
| typeSize = tiffTypes[type][0] | |
| typeChar = tiffTypes[type][1] | |
| input.seek(entryOffset + 8) | |
| value = input.read(typeSize) | |
| value = int(struct.unpack(typeChar, value)[0]) | |
| if tag == 256: | |
| width = value | |
| else: | |
| height = value | |
| if width > -1 and height > -1: | |
| break | |
| except Exception as e: | |
| raise UnknownImageFormat(str(e)) | |
| elif size >= 2: | |
| # see http://en.wikipedia.org/wiki/ICO_(file_format) | |
| imgtype = 'ICO' | |
| input.seek(0) | |
| reserved = input.read(2) | |
| if 0 != struct.unpack("<H", reserved)[0]: | |
| raise UnknownImageFormat(FILE_UNKNOWN) | |
| format = input.read(2) | |
| if 1 != struct.unpack("<H", format)[0]: | |
| raise UnknownImageFormat(FILE_UNKNOWN) | |
| num = input.read(2) | |
| num = struct.unpack("<H", num)[0] | |
| if num > 1: | |
| import warnings | |
| warnings.warn("ICO File contains more than one image") | |
| # http://msdn.microsoft.com/en-us/library/ms997538.aspx | |
| w = input.read(1) | |
| h = input.read(1) | |
| width = ord(w) | |
| height = ord(h) | |
| else: | |
| raise UnknownImageFormat(FILE_UNKNOWN) | |
| if orientation in _EXIF_SWAP_ORIENTATIONS: | |
| width, height = height, width | |
| return Image(path=file_path, | |
| type=imgtype, | |
| file_size=size, | |
| width=width, | |
| height=height) | |
| import unittest | |
| class Test_get_image_size(unittest.TestCase): | |
| data = [{ | |
| 'path': 'lookmanodeps.png', | |
| 'width': 251, | |
| 'height': 208, | |
| 'file_size': 22228, | |
| 'type': 'PNG'}] | |
| def setUp(self): | |
| pass | |
| def test_get_image_size_from_bytesio(self): | |
| img = self.data[0] | |
| p = img['path'] | |
| with io.open(p, 'rb') as fp: | |
| b = fp.read() | |
| fp = io.BytesIO(b) | |
| sz = len(b) | |
| output = get_image_size_from_bytesio(fp, sz) | |
| self.assertTrue(output) | |
| self.assertEqual(output, | |
| (img['width'], | |
| img['height'])) | |
| def test_get_image_metadata_from_bytesio(self): | |
| img = self.data[0] | |
| p = img['path'] | |
| with io.open(p, 'rb') as fp: | |
| b = fp.read() | |
| fp = io.BytesIO(b) | |
| sz = len(b) | |
| output = get_image_metadata_from_bytesio(fp, sz) | |
| self.assertTrue(output) | |
| for field in image_fields: | |
| self.assertEqual(getattr(output, field), None if field == 'path' else img[field]) | |
| def test_get_image_metadata(self): | |
| img = self.data[0] | |
| output = get_image_metadata(img['path']) | |
| self.assertTrue(output) | |
| for field in image_fields: | |
| self.assertEqual(getattr(output, field), img[field]) | |
| def test_get_image_metadata__ENOENT_OSError(self): | |
| with self.assertRaises(OSError): | |
| get_image_metadata('THIS_DOES_NOT_EXIST') | |
| def test_get_image_metadata__not_an_image_UnknownImageFormat(self): | |
| with self.assertRaises(UnknownImageFormat): | |
| get_image_metadata('README.rst') | |
| def test_get_image_size(self): | |
| img = self.data[0] | |
| output = get_image_size(img['path']) | |
| self.assertTrue(output) | |
| self.assertEqual(output, | |
| (img['width'], | |
| img['height'])) | |
| def tearDown(self): | |
| pass | |
| def main(argv=None): | |
| """ | |
| Print image metadata fields for the given file path. | |
| Keyword Arguments: | |
| argv (list): commandline arguments (e.g. sys.argv[1:]) | |
| Returns: | |
| int: zero for OK | |
| """ | |
| import logging | |
| import optparse | |
| import sys | |
| prs = optparse.OptionParser( | |
| usage="%prog [-v|--verbose] [--json|--json-indent] <path0> [<pathN>]", | |
| description="Print metadata for the given image paths " | |
| "(without image library bindings).") | |
| prs.add_option('--json', | |
| dest='json', | |
| action='store_true') | |
| prs.add_option('--json-indent', | |
| dest='json_indent', | |
| action='store_true') | |
| prs.add_option('-v', '--verbose', | |
| dest='verbose', | |
| action='store_true', ) | |
| prs.add_option('-q', '--quiet', | |
| dest='quiet', | |
| action='store_true', ) | |
| prs.add_option('-t', '--test', | |
| dest='run_tests', | |
| action='store_true', ) | |
| argv = list(argv) if argv is not None else sys.argv[1:] | |
| (opts, args) = prs.parse_args(args=argv) | |
| loglevel = logging.INFO | |
| if opts.verbose: | |
| loglevel = logging.DEBUG | |
| elif opts.quiet: | |
| loglevel = logging.ERROR | |
| logging.basicConfig(level=loglevel) | |
| log = logging.getLogger() | |
| log.debug('argv: %r', argv) | |
| log.debug('opts: %r', opts) | |
| log.debug('args: %r', args) | |
| if opts.run_tests: | |
| import sys | |
| sys.argv = [sys.argv[0]] + args | |
| import unittest | |
| return unittest.main() | |
| output_func = Image.to_str_row | |
| if opts.json_indent: | |
| import functools | |
| output_func = functools.partial(Image.to_str_json, indent=2) | |
| elif opts.json: | |
| output_func = Image.to_str_json | |
| elif opts.verbose: | |
| output_func = Image.to_str_row_verbose | |
| EX_OK = 0 | |
| EX_NOT_OK = 2 | |
| if len(args) < 1: | |
| prs.print_help() | |
| print('') | |
| prs.error("You must specify one or more paths to image files") | |
| errors = [] | |
| for path_arg in args: | |
| try: | |
| img = get_image_metadata(path_arg) | |
| print(output_func(img)) | |
| except KeyboardInterrupt: | |
| raise | |
| except OSError as e: | |
| log.error((path_arg, e)) | |
| errors.append((path_arg, e)) | |
| except Exception as e: | |
| log.exception(e) | |
| errors.append((path_arg, e)) | |
| pass | |
| if len(errors): | |
| import pprint | |
| print("ERRORS", file=sys.stderr) | |
| print("======", file=sys.stderr) | |
| print(pprint.pformat(errors, indent=2), file=sys.stderr) | |
| return EX_NOT_OK | |
| return EX_OK | |
| is_window_shown = False | |
| display_lock = threading.Lock() | |
| current_img = None | |
| update_event = threading.Event() | |
| def update_image(img, name): | |
| global current_img | |
| with display_lock: | |
| current_img = (img, name) | |
| update_event.set() | |
| def display_image_in_thread(): | |
| global is_window_shown | |
| def display_img(): | |
| global current_img | |
| while True: | |
| update_event.wait() | |
| with display_lock: | |
| if current_img: | |
| img, name = current_img | |
| cv2.imshow(name, img) | |
| current_img = None | |
| update_event.clear() | |
| if cv2.waitKey(1) & 0xFF == 27: # Esc key to stop | |
| cv2.destroyAllWindows() | |
| print('\nESC pressed, stopping') | |
| break | |
| if not is_window_shown: | |
| is_window_shown = True | |
| threading.Thread(target=display_img, daemon=True).start() | |
| def show_img(img, name='AI Toolkit'): | |
| img = np.clip(img, 0, 255).astype(np.uint8) | |
| update_image(img[:, :, ::-1], name) | |
| if not is_window_shown: | |
| display_image_in_thread() | |
| def show_tensors(imgs: torch.Tensor, name='AI Toolkit'): | |
| if len(imgs.shape) == 4: | |
| img_list = torch.chunk(imgs, imgs.shape[0], dim=0) | |
| else: | |
| img_list = [imgs] | |
| img = torch.cat(img_list, dim=3) | |
| img = img / 2 + 0.5 | |
| img_numpy = img.to(torch.float32).detach().cpu().numpy() | |
| img_numpy = np.clip(img_numpy, 0, 1) * 255 | |
| img_numpy = img_numpy.transpose(0, 2, 3, 1) | |
| img_numpy = img_numpy.astype(np.uint8) | |
| show_img(img_numpy[0], name=name) | |
| def save_tensors(imgs: torch.Tensor, path='output.png', fps=None): | |
| if len(imgs.shape) == 5 and imgs.shape[0] == 1: | |
| imgs = imgs.squeeze(0) | |
| if len(imgs.shape) == 4: | |
| img_list = torch.chunk(imgs, imgs.shape[0], dim=0) | |
| else: | |
| img_list = [imgs] | |
| num_frames = len(img_list) | |
| print(f"Saving {num_frames} frames to {path} at {fps} fps") | |
| if fps is not None and num_frames > 1: | |
| img = torch.cat(img_list, dim=0) | |
| else: | |
| img = torch.cat(img_list, dim=3) | |
| img = img / 2 + 0.5 | |
| img_numpy = img.to(torch.float32).detach().cpu().numpy() | |
| img_numpy = np.clip(img_numpy, 0, 1) * 255 | |
| img_numpy = img_numpy.transpose(0, 2, 3, 1) | |
| img_numpy = img_numpy.astype(np.uint8) | |
| if fps is not None and num_frames > 1: | |
| img_list = [PILImage.fromarray(img_numpy[i]) for i in range(num_frames)] | |
| duration = int(1000 / fps) | |
| img_list[0].save(path, save_all=True, append_images=img_list[1:], duration=duration, loop=0, quality=95) | |
| else: | |
| # concat images to one | |
| img_numpy = np.concatenate(img_numpy, axis=1) | |
| # conver to pil | |
| img_pil = PILImage.fromarray(img_numpy) | |
| img_pil.save(path) | |
| def show_latents(latents: torch.Tensor, vae: 'AutoencoderTiny', name='AI Toolkit'): | |
| if vae.device == 'cpu': | |
| vae.to(latents.device) | |
| latents = latents / vae.config['scaling_factor'] | |
| imgs = vae.decode(latents).sample | |
| show_tensors(imgs, name=name) | |
| def on_exit(): | |
| if is_window_shown: | |
| cv2.destroyAllWindows() | |
| def reduce_contrast(tensor, factor): | |
| # Ensure factor is between 0 and 1 | |
| factor = max(0, min(factor, 1)) | |
| # Calculate the mean of the tensor | |
| mean = torch.mean(tensor) | |
| # Reduce contrast | |
| adjusted_tensor = (tensor - mean) * factor + mean | |
| # Clip values to ensure they stay within -1 to 1 range | |
| return torch.clamp(adjusted_tensor, -1.0, 1.0) | |
| atexit.register(on_exit) | |
| if __name__ == "__main__": | |
| import sys | |
| sys.exit(main(argv=sys.argv[1:])) | |