This commit is contained in:
arkohut 2024-08-21 23:24:13 +08:00
commit a76adde798
3 changed files with 15 additions and 13 deletions

View File

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

View File

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

View File

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