241 lines
7.1 KiB
Python
241 lines
7.1 KiB
Python
import hashlib
|
|
import os
|
|
import sys
|
|
|
|
import imagehash
|
|
import numpy as np
|
|
from PIL import Image, ImageFilter
|
|
|
|
try:
|
|
from turbojpeg import TJPF_RGB, TurboJPEG
|
|
|
|
JPEG = TurboJPEG()
|
|
except Exception:
|
|
JPEG = None
|
|
pass
|
|
|
|
LaplaceX = ImageFilter.Kernel(size=(3, 3), kernel=(0, 0, 0, 1, -2, 1, 0, 0, 0), scale=1, offset=0)
|
|
LaplaceY = ImageFilter.Kernel(size=(3, 3), kernel=(0, 1, 0, 0, -2, 0, 0, 1, 0), scale=1, offset=0)
|
|
Border = np.zeros((10, 10), dtype=bool)
|
|
Border[0, :] = True
|
|
# Border[9, :] = True
|
|
Border[:, 0] = True
|
|
Border[:, 9] = True
|
|
|
|
|
|
class ImageMeasure:
|
|
def __init__(self, filename):
|
|
self.filename = filename
|
|
self.hash = None
|
|
self.time = None
|
|
self.size = None
|
|
self.description = None
|
|
self.width = None
|
|
self.height = None
|
|
self.portrait = None
|
|
self.fingerprint = None
|
|
self.sharpness = None
|
|
self.colors = {x: None for x in ("R", "G", "B", "BR", "BG", "BB")}
|
|
self.similarity = {"distance": 0, "color": 0, "aspect": 0}
|
|
self.tags = []
|
|
self.image = None
|
|
|
|
def __str__(self):
|
|
|
|
printable = []
|
|
for k, v in self.__dict__.items():
|
|
if k == "image":
|
|
if not self.image is None:
|
|
v = "Loaded..."
|
|
printable.append(f"{k}: {v}")
|
|
return "\n".join(printable)
|
|
|
|
def set_all(self):
|
|
self.set_filename_absolute()
|
|
self.get_hash()
|
|
self.get_time()
|
|
self.get_size()
|
|
self.get_shape()
|
|
self.get_fingerprint()
|
|
self.get_sharpness()
|
|
self.get_colors()
|
|
|
|
def set_filename_absolute(self):
|
|
self.filename = os.path.realpath(self.filename)
|
|
|
|
def get_hash(self):
|
|
"""Return hash of the file"""
|
|
if self.hash is None:
|
|
hasher = hashlib.sha1()
|
|
blk = 2**16
|
|
with open(self.filename, "rb") as f:
|
|
while True:
|
|
d = f.read(blk)
|
|
if not d:
|
|
break
|
|
hasher.update(d)
|
|
|
|
self.hash = hasher.hexdigest()
|
|
return self.hash
|
|
|
|
def get_time(self):
|
|
"""Return mtime of the file"""
|
|
if self.time is None:
|
|
self.time = int(os.path.getmtime(self.filename))
|
|
return self.time
|
|
|
|
def get_size(self):
|
|
if self.size is None:
|
|
self.size = os.path.getsize(self.filename)
|
|
return self.size
|
|
|
|
def get_width(self):
|
|
if self.width is None:
|
|
self.width, _, _ = self.get_shape()
|
|
return self.width
|
|
|
|
def get_height(self):
|
|
if self.height is None:
|
|
_, self.height, _ = self.get_shape()
|
|
return self.height
|
|
|
|
def get_portrait(self):
|
|
if self.portrait is None:
|
|
_, _, self.portrait = self.get_shape()
|
|
return self.portrait
|
|
|
|
def get_shape(self):
|
|
if self.width is None or self.height is None or self.portrait is None:
|
|
# self.height, self.width = self.get_image("numpy").shape[0:2]
|
|
self.width, self.height = read_image_size(self.filename)
|
|
self.portrait = self.height >= self.width
|
|
return self.width, self.height, self.portrait
|
|
|
|
def get_description(self):
|
|
if self.description is None:
|
|
self.description = read_image_comment(self.filename)
|
|
return self.description
|
|
|
|
|
|
def get_image(self, image_type="numpy"):
|
|
|
|
if self.image is None:
|
|
self.image, self.image_type = read_image(self.filename)
|
|
if self.image_type == "numpy":
|
|
if len(self.image.shape) > 2:
|
|
# BGR -> RGB
|
|
self.image = np.flip(self.image, axis=2)
|
|
|
|
if self.image_type == image_type:
|
|
return self.image
|
|
if image_type == "numpy":
|
|
return np.array(self.image)
|
|
if image_type == "PIL":
|
|
return Image.fromarray(self.image)
|
|
|
|
def get_fingerprint(self):
|
|
|
|
if self.fingerprint is None:
|
|
# self.fingerprint = str(imagehash.phash(self.get_image("PIL"), hash_size=8))
|
|
self.fingerprint = str(imagehash.dhash(self.get_image("PIL"), hash_size=8))
|
|
|
|
return self.fingerprint
|
|
|
|
def get_sharpness(self):
|
|
|
|
if self.sharpness is None:
|
|
try:
|
|
im = self.get_image("PIL").convert("L")
|
|
crop_box = (1, 1, im.width - 1, im.height - 1)
|
|
self.sharpness = round(
|
|
(
|
|
np.sum(np.abs(np.array(im.filter(LaplaceX).crop(crop_box)).astype(float)))
|
|
+ np.sum(np.abs(np.array(im.filter(LaplaceY).crop(crop_box)).astype(float)))
|
|
)
|
|
/ (2 * im.width * im.height),
|
|
4,
|
|
)
|
|
except Exception:
|
|
self.sharpness = 0
|
|
|
|
return self.sharpness
|
|
|
|
def get_colors(self):
|
|
|
|
def get_border(im):
|
|
return int(np.mean(im[Border]))
|
|
|
|
if self.colors["R"] is None:
|
|
im = self.get_image("PIL").convert("RGB")
|
|
th = im.copy()
|
|
th.thumbnail((1, 1), resample=Image.BILINEAR)
|
|
th = np.array(th)
|
|
im = np.array(im.resize((10, 10), resample=Image.BILINEAR))
|
|
self.colors["R"] = int(th[0][0][0])
|
|
self.colors["G"] = int(th[0][0][1])
|
|
self.colors["B"] = int(th[0][0][2])
|
|
self.colors["BR"] = get_border(im[:, :, 0])
|
|
self.colors["BG"] = get_border(im[:, :, 1])
|
|
self.colors["BB"] = get_border(im[:, :, 2])
|
|
return self.colors
|
|
|
|
def similarity_difference(self, other):
|
|
|
|
other_phash = imagehash.hex_to_hash(other.get_fingerprint())
|
|
this_phash = imagehash.hex_to_hash(self.get_fingerprint())
|
|
return other_phash - this_phash
|
|
|
|
def color_difference(self, other):
|
|
|
|
other_color = other.get_colors()
|
|
this_color = self.get_colors()
|
|
diff = round(
|
|
np.sqrt(
|
|
np.square(other_color["R"] - this_color["R"])
|
|
+ np.square(other_color["G"] - this_color["G"])
|
|
+ np.square(other_color["B"] - this_color["B"])
|
|
),
|
|
1,
|
|
)
|
|
return diff
|
|
|
|
def shape_difference(self, other):
|
|
|
|
return round(abs(float(other.width) / float(other.height) - float(self.width) / float(self.height)), 4)
|
|
|
|
|
|
EXTENSIONS = (".jpg", ".png", ".tif", ".gif", ".jpeg", ".tiff")
|
|
JPEG_EXTENSIONS = (".jpg", ".jpeg")
|
|
|
|
|
|
def is_image_extension(f):
|
|
return os.path.splitext(f.lower())[1] in EXTENSIONS
|
|
|
|
|
|
def is_jpeg(f):
|
|
return os.path.splitext(f.lower())[1] in JPEG_EXTENSIONS
|
|
|
|
|
|
def read_image(fname):
|
|
if is_jpeg(fname):
|
|
if JPEG:
|
|
try:
|
|
with open(fname, "rb") as fp:
|
|
return Image.fromarray(JPEG.decode(fp.read(), pixel_format=TJPF_RGB)), "PIL"
|
|
except Exception as e:
|
|
pass
|
|
# Do not return inside with:
|
|
im = Image.open(fname)
|
|
return im, "PIL"
|
|
|
|
|
|
def read_image_size(fname):
|
|
"""Just reading the size is faster with PIL"""
|
|
im = Image.open(fname)
|
|
return im.width, im.height
|
|
|
|
def read_image_comment(fname):
|
|
"""Just reading the comment with PIL"""
|
|
im = Image.open(fname)
|
|
return im.info.get('comment','')
|