Spaces:
Runtime error
Runtime error
File size: 3,322 Bytes
4f8ad24 |
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 89 90 91 92 93 94 95 96 97 98 99 |
import json
import os.path
import pickle
from dataclasses import dataclass
from typing import Optional
from PIL import Image
from hbutils.encoding import base64_decode, base64_encode
from hbutils.reflection import quick_import_object
NoneType = type(None)
_TYPE_META = '__type'
_BASE64_META = 'base64'
def load_meta(data, path=()):
if isinstance(data, (int, float, str, NoneType)):
return data
elif isinstance(data, list):
return [load_meta(item, (*path, i)) for i, item in enumerate(data)]
elif isinstance(data, dict):
if _TYPE_META not in data:
return {key: load_meta(value, (*path, key)) for key, value in data.items()}
else:
cls, _, _ = quick_import_object(data[_TYPE_META])
binary = base64_decode(data[_BASE64_META])
obj = pickle.loads(binary)
if isinstance(obj, cls):
return obj
else:
raise TypeError(f'{cls!r} expected but {obj!r} found at {path!r}.')
else:
raise TypeError(f'Unknown type {data!r} at {path!r}.')
def dump_meta(data, path=()):
if isinstance(data, (int, float, str, NoneType)):
return data
elif isinstance(data, list):
return [dump_meta(item, (*path, i)) for i, item in enumerate(data)]
elif isinstance(data, dict):
return {key: dump_meta(value, (*path, key)) for key, value in data.items()}
else:
cls = type(data)
type_str = f'{cls.__module__}.{cls.__name__}' if hasattr(cls, '__module__') else cls.__name__
base64_str = base64_encode(pickle.dumps(data))
return {
_TYPE_META: type_str,
_BASE64_META: base64_str
}
@dataclass
class ImageItem:
image: Image.Image
meta: dict
def __init__(self, image: Image.Image, meta: Optional[dict] = None):
self.image = image
self.meta = meta or {}
@classmethod
def _image_file_to_meta_file(cls, image_file):
directory, filename = os.path.split(image_file)
filebody, _ = os.path.splitext(filename)
meta_file = os.path.join(directory, f'.{filebody}_meta.json')
return meta_file
@classmethod
def load_from_image(cls, image_file):
image = Image.open(image_file)
meta_file = cls._image_file_to_meta_file(image_file)
if os.path.exists(meta_file):
with open(meta_file, 'r', encoding='utf-8') as f:
meta = load_meta(json.load(f))
else:
meta = {}
return cls(image, meta)
def save(self, image_file, no_meta: bool = False, skip_when_image_exist: bool = False):
if not skip_when_image_exist or not os.path.exists(image_file):
self.image.save(image_file)
if not no_meta and self.meta:
meta_file = self._image_file_to_meta_file(image_file)
with open(meta_file, 'w', encoding='utf-8') as f:
json.dump(dump_meta(self.meta), f)
def __repr__(self):
values = {'size': self.image.size}
for key, value in self.meta.items():
if isinstance(value, (int, float, str)):
values[key] = value
content = ', '.join(f'{key}: {values[key]!r}' for key in sorted(values.keys()))
return f'<{self.__class__.__name__} {content}>'
|