fix(sort-images): Detect formats from file signatures
This commit is contained in:
+31
-10
@@ -14,6 +14,30 @@ from PIL import Image
|
||||
# --- Utility Functions ---
|
||||
|
||||
|
||||
class ImageFormat(enum.Enum):
|
||||
JPEG = "jpeg"
|
||||
PNG = "png"
|
||||
WEBP = "webp"
|
||||
|
||||
@property
|
||||
def suffix(self) -> str:
|
||||
return ".jpg" if self is self.JPEG else f".{self.value}"
|
||||
|
||||
|
||||
def sniff_image_format(image_path: Path) -> ImageFormat:
|
||||
"""Identify a supported image format from its file signature."""
|
||||
with image_path.open("rb") as file:
|
||||
header = file.read(12)
|
||||
|
||||
if header.startswith(b"\xff\xd8\xff"):
|
||||
return ImageFormat.JPEG
|
||||
if header.startswith(b"\x89PNG\r\n\x1a\n"):
|
||||
return ImageFormat.PNG
|
||||
if len(header) >= 12 and header[:4] == b"RIFF" and header[8:12] == b"WEBP":
|
||||
return ImageFormat.WEBP
|
||||
raise ValueError(f"Unsupported or unrecognized image format: {image_path}")
|
||||
|
||||
|
||||
def read_dims_jpeg(image_path: Path) -> tuple[int, int]:
|
||||
with image_path.open("rb") as file:
|
||||
if file.read(2) != b"\xff\xd8":
|
||||
@@ -129,14 +153,14 @@ def read_dims_webp(image_path: Path) -> tuple[int, int]:
|
||||
|
||||
|
||||
def read_dims(image_path: Path) -> tuple[int, int]:
|
||||
if image_path.suffix.lower() in [".jpg", ".jpeg"]:
|
||||
image_format = sniff_image_format(image_path)
|
||||
if image_format is ImageFormat.JPEG:
|
||||
return read_dims_jpeg(image_path)
|
||||
elif image_path.suffix.lower() == ".png":
|
||||
elif image_format is ImageFormat.PNG:
|
||||
return read_dims_png(image_path)
|
||||
elif image_path.suffix.lower() == ".webp":
|
||||
elif image_format is ImageFormat.WEBP:
|
||||
return read_dims_webp(image_path)
|
||||
else:
|
||||
raise ValueError(f"Unsupported image format: {image_path.suffix}")
|
||||
raise AssertionError(f"Unhandled image format: {image_format}")
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
@@ -188,11 +212,8 @@ def main():
|
||||
logging.warning(f"{image_path} is not a file. Skipping.")
|
||||
continue
|
||||
|
||||
if image_path.suffix.lower() not in [".jpg", ".jpeg", ".png", ".webp"]:
|
||||
logging.warning(f"{image_path} is not a JPEG, PNG, or WebP file. Skipping.")
|
||||
continue
|
||||
|
||||
try:
|
||||
image_format = sniff_image_format(image_path)
|
||||
width, height = read_dims(image_path)
|
||||
target_dir = None
|
||||
|
||||
@@ -208,7 +229,7 @@ def main():
|
||||
make_dir(target_dir)
|
||||
|
||||
target_path = target_dir / image_path.name
|
||||
target_path = target_path.with_suffix(".jpg")
|
||||
target_path = target_path.with_suffix(image_format.suffix)
|
||||
|
||||
image_path.rename(target_path)
|
||||
logging.info(f"Moved {image_path} to {target_path}")
|
||||
|
||||
Reference in New Issue
Block a user