fix: update utils import

This commit is contained in:
arkohut 2024-08-21 23:22:58 +08:00
parent 15972b9c0f
commit 9a38b0c17d
3 changed files with 15 additions and 13 deletions

View File

@ -2,31 +2,33 @@ from PIL import Image
import piexif
from PIL.PngImagePlugin import PngInfo
import json
from pathlib import Path
def write_image_metadata(image_path, metadata):
img = Image.open(image_path)
image_path_str = str(image_path)
if image_path.lower().endswith((".jpg", ".jpeg", ".tiff")):
if image_path_str.lower().endswith((".jpg", ".jpeg", ".tiff")):
exif_dict = piexif.load(image_path)
updated_description = json.dumps(metadata).encode("utf-8")
exif_dict["0th"][piexif.ImageIFD.ImageDescription] = updated_description
exif_bytes = piexif.dump(exif_dict)
img.save(image_path, exif=exif_bytes)
elif image_path.lower().endswith(".png"):
elif image_path_str.lower().endswith(".png"):
metadata_info = PngInfo()
metadata_info.add_text("Description", json.dumps(metadata))
img.save(image_path, "PNG", pnginfo=metadata_info)
elif image_path.lower().endswith(".webp"):
elif image_path_str.lower().endswith(".webp"):
img.save(image_path, "WebP", quality=85, metadata=json.dumps(metadata))
else:
print(f"Skipping unsupported file format: {image_path}")
print(f"Skipping unsupported file format: {image_path_str}")
def get_image_metadata(image_path):
img = Image.open(image_path)
image_path_str = str(image_path)
if image_path.lower().endswith((".jpg", ".jpeg", ".tiff")):
if image_path_str.lower().endswith((".jpg", ".jpeg", ".tiff")):
exif_dict = piexif.load(image_path)
existing_description = exif_dict["0th"].get(
piexif.ImageIFD.ImageDescription, b"{}"
@ -35,18 +37,18 @@ def get_image_metadata(image_path):
return json.loads(existing_description.decode("utf-8"))
except json.JSONDecodeError:
return {}
elif image_path.lower().endswith(".png"):
elif image_path_str.lower().endswith(".png"):
existing_description = img.info.get("Description", "{}")
try:
return json.loads(existing_description)
except json.JSONDecodeError:
return {}
elif image_path.lower().endswith(".webp"):
elif image_path_str.lower().endswith(".webp"):
existing_metadata = img.info.get("metadata", "{}")
try:
return json.loads(existing_metadata)
except json.JSONDecodeError:
return {}
else:
print(f"Unsupported file format: {image_path}")
return None
print(f"Unsupported file format: {image_path_str}")
return None

View File

@ -8,7 +8,7 @@ import win32gui
import win32process
import psutil
from mss import mss
from .utils import write_image_metadata
from memos.utils import write_image_metadata
import win32api
import pywintypes
import ctypes

View File

@ -12,7 +12,7 @@ from Quartz import (
import json
import imagehash
import argparse
from .utils import write_image_metadata
from memos.utils import write_image_metadata
def get_active_window_info():
@ -186,4 +186,4 @@ def main():
if __name__ == "__main__":
main()
main()