128 lines
3.6 KiB
Python
128 lines
3.6 KiB
Python
from collections.abc import Iterable
|
|
from io import BytesIO
|
|
from pathlib import Path
|
|
|
|
import PIL.Image as PILImage
|
|
from PIL import UnidentifiedImageError
|
|
from textual import on
|
|
from textual.app import ComposeResult
|
|
from textual.binding import Binding
|
|
from textual.containers import Container, Horizontal, Vertical
|
|
from textual.message import Message
|
|
from textual.screen import ModalScreen
|
|
from textual.widgets import DirectoryTree, Input, Label
|
|
from textual_image.widget import Image
|
|
|
|
IMAGE_EXTENSIONS = PILImage.registered_extensions()
|
|
|
|
|
|
def encode_jpeg(path: Path) -> bytes:
|
|
"""Re-encode an image file as RGB JPEG bytes."""
|
|
image = PILImage.open(path)
|
|
if image.mode != "RGB":
|
|
image = image.convert("RGB")
|
|
buffer = BytesIO()
|
|
image.save(buffer, format="JPEG")
|
|
return buffer.getvalue()
|
|
|
|
|
|
class ImageAdded(Message):
|
|
"""Emitted when the user picks an image to attach to the prompt."""
|
|
|
|
def __init__(self, path: Path, data: bytes) -> None:
|
|
self.path = path
|
|
self.data = data
|
|
super().__init__()
|
|
|
|
|
|
class ImageDirectoryTree(DirectoryTree):
|
|
def filter_paths(self, paths: Iterable[Path]) -> Iterable[Path]:
|
|
return [
|
|
path for path in paths if path.suffix in IMAGE_EXTENSIONS or path.is_dir()
|
|
]
|
|
|
|
|
|
class ImageSelect(ModalScreen[tuple[Path, bytes]]):
|
|
"""Modal for picking an image file, with a live preview."""
|
|
|
|
BINDINGS = [
|
|
Binding("escape", "cancel", "Cancel", show=False),
|
|
]
|
|
|
|
CSS = """
|
|
ImageSelect {
|
|
align: center middle;
|
|
background: rgba(0, 0, 0, 0.5);
|
|
}
|
|
|
|
#image-select-container {
|
|
width: 80%;
|
|
height: 80%;
|
|
background: $surface;
|
|
border: tall $primary;
|
|
padding: 1 2;
|
|
}
|
|
|
|
#image-directory-tree {
|
|
width: 40%;
|
|
}
|
|
|
|
#image-preview {
|
|
width: 60%;
|
|
}
|
|
|
|
#image-preview #image {
|
|
width: auto;
|
|
height: auto;
|
|
}
|
|
|
|
#image-select-container Input {
|
|
margin-bottom: 1;
|
|
}
|
|
"""
|
|
|
|
def action_cancel(self) -> None:
|
|
self.dismiss()
|
|
|
|
async def on_mount(self) -> None:
|
|
tree = self.query_one(ImageDirectoryTree)
|
|
tree.show_guides = False
|
|
tree.focus()
|
|
|
|
@on(DirectoryTree.FileSelected)
|
|
async def on_image_selected(self, event: DirectoryTree.FileSelected) -> None:
|
|
try:
|
|
self.dismiss((event.path, encode_jpeg(event.path)))
|
|
except UnidentifiedImageError:
|
|
self.dismiss()
|
|
|
|
@on(DirectoryTree.NodeHighlighted)
|
|
async def on_image_highlighted(self, event: DirectoryTree.NodeHighlighted) -> None:
|
|
if event.node.data is None:
|
|
return
|
|
path = event.node.data.path
|
|
preview = self.query_one(Image)
|
|
if path.suffix in IMAGE_EXTENSIONS:
|
|
try:
|
|
preview.image = PILImage.open(path.as_posix())
|
|
except UnidentifiedImageError:
|
|
preview.image = None
|
|
else:
|
|
preview.image = None
|
|
|
|
@on(Input.Changed)
|
|
async def on_root_changed(self, event: Input.Changed) -> None:
|
|
path = Path(event.value)
|
|
if path.exists() and path.is_dir():
|
|
self.query_one(ImageDirectoryTree).path = path
|
|
|
|
def compose(self) -> ComposeResult:
|
|
with Container(id="image-select-container"):
|
|
with Horizontal():
|
|
with Vertical(id="image-directory-tree"):
|
|
yield Label("Select an image:")
|
|
yield Label("Root:")
|
|
yield Input(Path("./").resolve().as_posix())
|
|
yield ImageDirectoryTree("./")
|
|
with Container(id="image-preview"):
|
|
yield Image(id="image")
|