Files
easy_prompts/main.py
Julian Freeman d15e2e59ba mark bg color
2025-08-11 00:50:35 -04:00

375 lines
14 KiB
Python
Raw Permalink Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
#!/usr/bin/env python3
# coding: utf-8
import os
import json
import sys
from functools import partial
from PySide6 import QtCore, QtGui, QtWidgets
from PIL import Image, PngImagePlugin, ImageQt
import piexif
# ============ 配置区 ============
PROJECT_COLUMNS = ["Runway", "Hailuo", "Kling", "Dreamina", "Vidu", "Pixverse"] # 自定义列名
THUMBNAIL_SIZE = (320, 240) # 缩略图大小(宽, 高)
# ==============================
def load_notes_from_image(path):
"""从图片中读取 JSON 格式的备注(如果存在),返回 dict否则返回 {}"""
try:
ext = path.lower().split(".")[-1]
img = Image.open(path)
if ext in ("jpg", "jpeg"):
exif_bytes = img.info.get("exif")
if exif_bytes:
try:
exif = piexif.load(exif_bytes)
raw = exif.get("Exif", {}).get(piexif.ExifIFD.UserComment)
if raw:
# piexif UserComment 通常有两个前缀字节,照处理
if isinstance(raw, bytes) and raw.startswith(b'\x00\x00'):
raw = raw[2:]
if isinstance(raw, bytes):
text = raw.decode("utf-8", errors="ignore")
else:
text = str(raw)
return json.loads(text) if text else {}
except Exception:
return {}
return {}
elif ext == "png":
# Pillow 将 text chunks 放在 img.info
note = img.info.get("notes") or img.info.get("Notes") or img.info.get("text") or img.info.get("TEXT")
if note:
try:
return json.loads(note)
except Exception:
return {}
return {}
else:
return {}
except Exception as e:
print(f"[load_notes_from_image] 读取失败 {path}: {e}")
return {}
def save_notes_to_image(path, note_dict):
"""把 note_dict (python dict) 写入图片文件:
JPEG -> Exif.UserComment
PNG -> tEXt 'notes'
"""
try:
ext = path.lower().split(".")[-1]
note_json = json.dumps(note_dict, ensure_ascii=False)
if ext in ("jpg", "jpeg"):
img = Image.open(path)
# 获取已有 exif如果没有就创建结构
try:
exif_dict = piexif.load(img.info.get("exif", b""))
except Exception:
exif_dict = {"0th": {}, "Exif": {}, "GPS": {}, "1st": {}, "Interop": {}, "thumbnail": None}
# 写入 UserComment前两个 0 字节用来代表编码,保持之前的写法)
exif_dict.setdefault("Exif", {})
exif_dict["Exif"][piexif.ExifIFD.UserComment] = b'\x00\x00' + note_json.encode("utf-8")
exif_bytes = piexif.dump(exif_dict)
# 覆盖保存 jpeg注意会重新压缩图片
img_rgb = img.convert("RGB") # 确保jpeg格式
img_rgb.save(path, "jpeg", exif=exif_bytes)
elif ext == "png":
img = Image.open(path)
meta = PngImagePlugin.PngInfo()
# 添加文本 chunk
meta.add_text("notes", note_json)
# 保存时保留原模式(注意:这也会重新写文件)
img.save(path, "png", pnginfo=meta)
else:
raise ValueError("不支持的图片格式")
return True, None
except Exception as e:
return False, str(e)
class NoteEditor(QtWidgets.QPlainTextEdit):
"""
自定义编辑器:
- Shift+Enter 插入换行
- Enter (无 Shift) 触发保存emit save_requested 信号)
- 右键菜单设置背景色(绿色、红色、黄色、移除)
"""
save_requested = QtCore.Signal()
def __init__(self, parent=None):
super().__init__(parent)
# 可选:设置合适的样式/行高
self.setTabChangesFocus(False)
self.setPlaceholderText("按 Enter 保存Shift+Enter 换行)")
# 右键菜单
self.setContextMenuPolicy(QtCore.Qt.CustomContextMenu)
self.customContextMenuRequested.connect(self._show_bg_menu)
self.bg_color = "" # 可选值: "", "green", "red", "yellow"
def keyPressEvent(self, event: QtGui.QKeyEvent):
if event.key() in (QtCore.Qt.Key_Return, QtCore.Qt.Key_Enter):
# 如果按下 Shift则插入换行
if event.modifiers() & QtCore.Qt.ShiftModifier:
super().keyPressEvent(event)
return
# 否则触发保存信号(并阻止默认换行)
# 但先让文档内容同步(光标位置等)
event.accept()
# 发出保存信号(外部会连接并执行保存操作)
self.save_requested.emit()
return
super().keyPressEvent(event)
def _show_bg_menu(self, pos):
menu = QtWidgets.QMenu(self)
green_action = menu.addAction("设置绿色背景")
red_action = menu.addAction("设置红色背景")
yellow_action = menu.addAction("设置黄色背景")
menu.addSeparator()
clear_action = menu.addAction("移除背景色")
action = menu.exec(self.mapToGlobal(pos))
if action == green_action:
self.setStyleSheet("background-color: #234d2a; color: #e0e0e0;")
self.bg_color = "green"
elif action == red_action:
self.setStyleSheet("background-color: #4d2323; color: #e0e0e0;")
self.bg_color = "red"
elif action == yellow_action:
self.setStyleSheet("background-color: #4d4d23; color: #e0e0e0;")
self.bg_color = "yellow"
elif action == clear_action:
self.setStyleSheet("")
self.bg_color = ""
def set_bg_color(self, color):
"""用于初始化时恢复背景色"""
self.bg_color = color
if color == "green":
self.setStyleSheet("background-color: #234d2a; color: #e0e0e0;")
elif color == "red":
self.setStyleSheet("background-color: #4d2323; color: #e0e0e0;")
elif color == "yellow":
self.setStyleSheet("background-color: #4d4d23; color: #e0e0e0;")
else:
self.setStyleSheet("")
class ImageRowWidget(QtWidgets.QWidget):
"""每一行的组合组件:缩略图 + 多个 NoteEditor"""
def __init__(self, folder_path, filename, projects, thumbnail_size):
super().__init__()
self.folder_path = folder_path
self.filename = filename
self.path = os.path.join(folder_path, filename)
self.projects = projects
h = QtWidgets.QHBoxLayout()
h.setContentsMargins(4, 4, 4, 4)
h.setSpacing(8)
# 缩略图
self.thumb_label = QtWidgets.QLabel()
self.thumb_label.setFixedSize(thumbnail_size[0], thumbnail_size[1])
self.thumb_label.setAlignment(QtCore.Qt.AlignCenter)
pix = self._load_thumbnail(self.path, thumbnail_size)
if pix:
self.thumb_label.setPixmap(pix)
else:
self.thumb_label.setText("无法加载")
h.addWidget(self.thumb_label)
# 读取已有元数据
self.note_dict = load_notes_from_image(self.path)
# editors 按项目顺序
self.editors = {}
for proj in projects:
editor = NoteEditor()
# 初始化内容(如果没有该项目,则为空)
editor.setPlainText(self.note_dict.get(proj, ""))
editor.setMinimumHeight(80)
# 恢复背景色
bg = self.note_dict.get(f"{proj}_bg", "")
editor.set_bg_color(bg)
# 当触发保存信号时,调用保存回调(只保存当前图片)
editor.save_requested.connect(partial(self._on_save_requested))
self.editors[proj] = editor
h.addWidget(editor, 1)
self.setLayout(h)
def _load_thumbnail(self, image_path, size):
"""用 PIL 创建缩略并转为 QPixmap"""
try:
img = Image.open(image_path)
img.thumbnail(size, Image.LANCZOS)
# 转为Qt pixmap
qim = ImageQt.ImageQt(img.convert("RGBA"))
pix = QtGui.QPixmap.fromImage(qim)
return pix
except Exception as e:
print(f"[thumbnail] {image_path} load fail: {e}")
return None
def collect_notes(self):
"""收集当前编辑器里所有项目字段和背景色,返回 dict"""
d = {}
for proj, editor in self.editors.items():
text = editor.toPlainText().strip()
d[proj] = text
d[f"{proj}_bg"] = editor.bg_color
return d
@QtCore.Slot()
def _on_save_requested(self):
"""当某个编辑器触发保存时,写入该图片的元数据"""
notes = self.collect_notes()
success, err = save_notes_to_image(self.path, notes)
if success:
# 简单的视觉反馈:临时背景闪烁或状态提示
for editor in self.editors.values():
# 设置短暂背景色
editor.setStyleSheet("background-color: #e6ffe6;")
QtCore.QTimer.singleShot(350, self._clear_editor_styles)
else:
# 弹出错误提示框
QtWidgets.QMessageBox.warning(self, "保存失败", f"图片 {self.filename} 写入失败:\n{err}")
def _clear_editor_styles(self):
for proj, editor in self.editors.items():
editor.set_bg_color(editor.bg_color)
class MainWindow(QtWidgets.QMainWindow):
def __init__(self, folder, projects):
super().__init__()
# ...existing code...
self.folder = folder
self.projects = projects
self.setWindowTitle("简易记录提示词")
self.resize(1500, 800)
central = QtWidgets.QWidget()
vbox = QtWidgets.QVBoxLayout()
central.setLayout(vbox)
self.setCentralWidget(central)
# ======= 修改:图片目录选择栏 =======
dir_widget = QtWidgets.QWidget()
dir_layout = QtWidgets.QHBoxLayout()
dir_layout.setContentsMargins(4, 4, 4, 4)
dir_widget.setLayout(dir_layout)
self.dir_edit = QtWidgets.QLineEdit(self.folder)
self.dir_edit.setPlaceholderText("输入或选择图片目录路径")
dir_layout.addWidget(QtWidgets.QLabel("图片目录:"))
dir_layout.addWidget(self.dir_edit, 1)
select_btn = QtWidgets.QPushButton("选择目录")
select_btn.clicked.connect(self._on_select_folder)
dir_layout.addWidget(select_btn)
vbox.addWidget(dir_widget)
# ======= 修改结束 =======
# 顶部说明栏
instr = QtWidgets.QLabel(
"说明:在任意单元格中编辑,按 Enter不按 Shift保存该图片的所有字段"
"Shift+Enter 插入换行。JPEG 写入 Exif.UserCommentPNG 写入 tEXt 'notes'"
)
instr.setWordWrap(True)
vbox.addWidget(instr)
# 标题行(缩略 + 每个项目列名)
header = QtWidgets.QWidget()
header_h = QtWidgets.QHBoxLayout()
header_h.setContentsMargins(4, 4, 4, 4)
header.setLayout(header_h)
thumb_title = QtWidgets.QLabel("缩略")
thumb_title.setFixedSize(THUMBNAIL_SIZE[0], 30)
thumb_title.setAlignment(QtCore.Qt.AlignCenter)
header_h.addWidget(thumb_title)
for proj in projects:
lbl = QtWidgets.QLabel(proj)
lbl.setMinimumHeight(30)
lbl.setAlignment(QtCore.Qt.AlignCenter)
header_h.addWidget(lbl)
vbox.addWidget(header)
# Scroll area for rows
scroll = QtWidgets.QScrollArea()
scroll.setWidgetResizable(True)
container = QtWidgets.QWidget()
self.rows_layout = QtWidgets.QVBoxLayout()
container.setLayout(self.rows_layout)
scroll.setWidget(container)
vbox.addWidget(scroll)
# Load images and populate
if self.folder and os.path.isdir(self.folder):
self._load_images()
# 状态栏
self.status = QtWidgets.QStatusBar()
self.setStatusBar(self.status)
def _load_images(self):
# 清空
for i in reversed(range(self.rows_layout.count())):
w = self.rows_layout.itemAt(i).widget()
if w:
w.setParent(None)
# 修复:仅在目录有效时加载图片,否则直接返回
if not self.folder or not os.path.isdir(self.folder):
return
# 获取图片列表(按文件名排序)
try:
files = [f for f in os.listdir(self.folder)
if f.lower().endswith((".jpg", ".jpeg", ".png"))]
files.sort()
except Exception as e:
QtWidgets.QMessageBox.critical(self, "错误", f"无法读取文件夹 {self.folder}\n{e}")
return
if not files:
lbl = QtWidgets.QLabel("文件夹中没有 JPEG/PNG 图片")
self.rows_layout.addWidget(lbl)
return
for fname in files:
row = ImageRowWidget(self.folder, fname, self.projects, THUMBNAIL_SIZE)
self.rows_layout.addWidget(row)
# 填充底部伸缩项
self.rows_layout.addStretch(1)
def _on_select_folder(self):
"""弹出目录选择对话框,选择后填充并加载图片"""
dlg = QtWidgets.QFileDialog(self, "选择图片目录")
dlg.setFileMode(QtWidgets.QFileDialog.Directory)
dlg.setOption(QtWidgets.QFileDialog.ShowDirsOnly, True)
if dlg.exec():
selected = dlg.selectedFiles()
if selected:
folder = selected[0]
self.dir_edit.setText(folder)
self.folder = folder
self._load_images()
def main():
app = QtWidgets.QApplication(sys.argv)
# 默认目录可设为空字符串或当前目录
win = MainWindow("", PROJECT_COLUMNS)
win.show()
sys.exit(app.exec())
if __name__ == "__main__":
main()