ArgbImage from arbitrary (but square) image data + guess rgb-mask types
This commit is contained in:
@@ -1,5 +1,6 @@
|
||||
#!/usr/bin/env python3
|
||||
from typing import Union, Iterator, Optional
|
||||
from math import sqrt
|
||||
from . import IcnsType, PackBytes
|
||||
try:
|
||||
from PIL import Image
|
||||
@@ -72,12 +73,23 @@ class ArgbImage:
|
||||
if is_argb or data[:4] == b'\x00\x00\x00\x00':
|
||||
data = data[4:] # remove ARGB and it32 header
|
||||
|
||||
idat = PackBytes.unpack(data)
|
||||
iType = IcnsType.match_maxsize(len(idat), 'argb' if is_argb else 'rgb')
|
||||
uncompressed_data = PackBytes.unpack(data)
|
||||
|
||||
self.size = iType.size
|
||||
self.channels = iType.channels or 0
|
||||
self.a, self.r, self.g, self.b = iType.split_channels(idat)
|
||||
self.channels = 4 if is_argb else 3
|
||||
per_channel = len(uncompressed_data) // self.channels
|
||||
w = sqrt(per_channel)
|
||||
if w != int(w):
|
||||
raise NotImplementedError(
|
||||
'Could not determine square image size. Or unknown type.')
|
||||
self.size = (int(w), int(w))
|
||||
if self.channels == 3:
|
||||
self.a = [255] * per_channel # opaque alpha channel for rgb
|
||||
else:
|
||||
self.a = uncompressed_data[:per_channel]
|
||||
i = 1 if is_argb else 0
|
||||
self.r = uncompressed_data[(i + 0) * per_channel:(i + 1) * per_channel]
|
||||
self.g = uncompressed_data[(i + 1) * per_channel:(i + 2) * per_channel]
|
||||
self.b = uncompressed_data[(i + 2) * per_channel:(i + 3) * per_channel]
|
||||
|
||||
def load_mask(self, *, file: Optional[str] = None,
|
||||
data: Optional[bytes] = None) -> None:
|
||||
@@ -99,15 +111,12 @@ class ArgbImage:
|
||||
return bytes(PackBytes.msb_stream(self.a, bits=bits))
|
||||
|
||||
def rgb_data(self, *, compress: bool = True) -> bytes:
|
||||
return b''.join(self._raw_rgb_channels(compress=compress))
|
||||
return b''.join(PackBytes.pack(x) if compress else bytes(x)
|
||||
for x in (self.r, self.g, self.b))
|
||||
|
||||
def argb_data(self, *, compress: bool = True) -> bytes:
|
||||
return b'ARGB' + self.mask_data(compress=compress) + \
|
||||
b''.join(self._raw_rgb_channels(compress=compress))
|
||||
|
||||
def _raw_rgb_channels(self, *, compress: bool = True) -> Iterator[bytes]:
|
||||
for x in (self.r, self.g, self.b):
|
||||
yield (PackBytes.pack(x) if compress else bytes(x))
|
||||
return b'ARGB' + self.mask_data(compress=compress) \
|
||||
+ self.rgb_data(compress=compress)
|
||||
|
||||
def _load_png(self, fname: str) -> None:
|
||||
if not PIL_ENABLED:
|
||||
|
||||
@@ -54,20 +54,6 @@ class Media:
|
||||
return self.desc # guaranteed to be icon, mask, or iconmask
|
||||
return self.types[-1]
|
||||
|
||||
def split_channels(self, uncompressed_data: List[int]) -> Iterator[
|
||||
List[int]]:
|
||||
if self.channels not in [3, 4]:
|
||||
raise NotImplementedError('Only RGB and ARGB data supported.')
|
||||
if len(uncompressed_data) != self.maxsize:
|
||||
raise ValueError(
|
||||
'Data does not match expected uncompressed length. '
|
||||
'{} != {}'.format(len(uncompressed_data), self.maxsize))
|
||||
per_channel = self.maxsize // self.channels
|
||||
if self.channels == 3:
|
||||
yield [255] * per_channel # opaque alpha channel for rgb
|
||||
for i in range(self.channels):
|
||||
yield uncompressed_data[per_channel * i:per_channel * (i + 1)]
|
||||
|
||||
def decompress(self, data: bytes, ext: Optional[str] = '-?-') -> Optional[
|
||||
List[int]]:
|
||||
''' Returns None if media is not decompressable. '''
|
||||
@@ -230,7 +216,7 @@ def get(key: Media.KeyT) -> Media:
|
||||
def match_maxsize(total: int, typ: str) -> Media:
|
||||
assert(typ == 'argb' or typ == 'rgb')
|
||||
ret = [x for x in _TYPES.values() if x.is_type(typ) and x.maxsize == total]
|
||||
return ret[0] # TODO: handle cases with multiple options? eg: is32 icp4
|
||||
return _best_option(ret, typ)
|
||||
|
||||
|
||||
def guess(data: bytes, filename: Optional[str] = None) -> Media:
|
||||
@@ -247,38 +233,64 @@ def guess(data: bytes, filename: Optional[str] = None) -> Media:
|
||||
if bname in _TYPES:
|
||||
return _TYPES[bname]
|
||||
|
||||
ext = RawData.determine_file_ext(data)
|
||||
if not ext and filename and filename.endswith('.rgb'):
|
||||
ext = 'rgb'
|
||||
|
||||
# Guess by image size and retina flag
|
||||
size = RawData.determine_image_size(data, ext) if ext else None
|
||||
retina = bname.lower().endswith('@2x') if filename else False
|
||||
if size == (1024, 1024):
|
||||
retina = True # stupid double usage of ic10
|
||||
|
||||
# Icns specific names
|
||||
# Filter attributes
|
||||
desc = None
|
||||
if ext == 'icns' and filename:
|
||||
size = None
|
||||
maxsize = None
|
||||
retina = False
|
||||
|
||||
# Guess extension
|
||||
ext = RawData.determine_file_ext(data)
|
||||
if not ext and filename:
|
||||
if filename.endswith('.rgb'):
|
||||
ext = 'rgb'
|
||||
elif filename.endswith('.mask'):
|
||||
maxsize = len(data)
|
||||
desc = 'mask'
|
||||
|
||||
# Guess image size
|
||||
if ext:
|
||||
size = RawData.determine_image_size(data, ext)
|
||||
|
||||
# if filename is set, then bname is also set (see above)
|
||||
if filename:
|
||||
# Guess retina flag
|
||||
retina = bname.lower().endswith('@2x')
|
||||
# Guess icns-specific type
|
||||
if ext == 'icns':
|
||||
for candidate in ['template', 'selected', 'dark']:
|
||||
if filename.endswith(candidate + '.icns'):
|
||||
if bname.endswith(candidate):
|
||||
desc = candidate
|
||||
break
|
||||
|
||||
# stupid double usage of ic10, enforce retina flag
|
||||
if size == (1024, 1024):
|
||||
retina = True
|
||||
|
||||
choices = []
|
||||
for x in _TYPES.values():
|
||||
if retina != x.retina: # png + jp2
|
||||
continue
|
||||
if desc and desc != x.desc: # icns or rgb-mask
|
||||
continue
|
||||
if ext:
|
||||
if size != x.size or not x.is_type(ext):
|
||||
continue
|
||||
if desc and desc != x.desc: # icns only
|
||||
continue
|
||||
else: # not ext
|
||||
if x.ext_certain:
|
||||
continue
|
||||
if maxsize and x.maxsize and maxsize != x.maxsize: # mask only
|
||||
continue
|
||||
choices.append(x)
|
||||
|
||||
return _best_option(choices, ext)
|
||||
|
||||
|
||||
def _best_option(choices: List[Media], ext: Optional[str] = None) -> Media:
|
||||
'''
|
||||
Get most favorable media type.
|
||||
If more than one option exists, choose based on order index of ext.
|
||||
'''
|
||||
if len(choices) == 1:
|
||||
return choices[0]
|
||||
# Try get most favorable type (sort order of types)
|
||||
@@ -294,7 +306,7 @@ def guess(data: bytes, filename: Optional[str] = None) -> Media:
|
||||
best_choice.append(x)
|
||||
if len(best_choice) == 1:
|
||||
return best_choice[0]
|
||||
# Else
|
||||
raise CanNotDetermine(
|
||||
'Could not determine type for file: "{}" – one of {}.'.format(
|
||||
filename, [x.key for x in choices]))
|
||||
choices = best_choice
|
||||
|
||||
raise CanNotDetermine('Could not determine type – one of {}.'.format(
|
||||
[x.key for x in choices]))
|
||||
|
||||
Reference in New Issue
Block a user