143 lines
4.7 KiB
Python
143 lines
4.7 KiB
Python
# import dearpygui.dearpygui as dpg
|
||
# import numpy as np
|
||
|
||
# from .pipeline_stage_widget import PipelineStageWidget
|
||
|
||
|
||
# class ExportStage(PipelineStageWidget):
|
||
# name = "Export Image"
|
||
# register = True
|
||
# has_pipeline_in = True
|
||
# has_pipeline_out = False
|
||
|
||
# def __init__(self, manager, logger):
|
||
# super().__init__(manager, logger, default_stage_out="opened_image")
|
||
# self.manager.bus.subscribe(
|
||
# "process_full_res", self._on_process_full_res, True)
|
||
|
||
# def create_pipeline_stage_content(self):
|
||
# dpg.add_text("Some export fields")
|
||
|
||
# def _on_process_full_res(self, data):
|
||
# self.logger.info("Starting full res pipeline export")
|
||
|
||
# def on_pipeline_data(self, img):
|
||
# if img is None:
|
||
# return
|
||
# self.logger.info("low res image received, ignore")
|
||
|
||
# def on_full_res_pipeline_data(self, img):
|
||
# if img is None:
|
||
# return
|
||
# h, w, _ = img.shape
|
||
# self.logger.info(f"Full res image received: {w}x{h}")
|
||
|
||
|
||
import os
|
||
import dearpygui.dearpygui as dpg
|
||
import numpy as np
|
||
from PIL import Image
|
||
|
||
from .pipeline_stage_widget import PipelineStageWidget
|
||
|
||
|
||
class ExportStage(PipelineStageWidget):
|
||
name = "Export Image"
|
||
register = True
|
||
has_pipeline_in = True
|
||
has_pipeline_out = False
|
||
|
||
def __init__(self, manager, logger):
|
||
# we don’t register an output stage — this widget only consumes
|
||
super().__init__(manager, logger, default_stage_out="unused")
|
||
# tags for our “Save As” dialog
|
||
self._save_dialog_tag = dpg.generate_uuid()
|
||
self._save_path = f"{os.getcwd()}/out.png"
|
||
|
||
def create_pipeline_stage_content(self):
|
||
# Button to pop up the file-save dialog
|
||
dpg.add_button(label="Save As…",
|
||
callback=lambda s,a,u: dpg.configure_item(self._save_dialog_tag, show=True))
|
||
|
||
# File dialog for choosing export path & extension
|
||
with dpg.file_dialog(
|
||
directory_selector=False,
|
||
show=False,
|
||
callback=self._on_save_selected,
|
||
tag=self._save_dialog_tag,
|
||
width=400,
|
||
height=300
|
||
):
|
||
dpg.add_file_extension("PNG {.png}")
|
||
dpg.add_file_extension("JPEG {.jpg,.jpeg}")
|
||
dpg.add_file_extension("TIFF {.tif,.tiff}")
|
||
dpg.add_file_extension("All files {.*}")
|
||
|
||
with dpg.child_window(autosize_x=True, autosize_y=True, horizontal_scrollbar=True):
|
||
self.path_label = dpg.add_text("out.png")
|
||
|
||
def _on_save_selected(self, sender, app_data):
|
||
"""
|
||
Called when the user picks a filename in the Save As… dialog.
|
||
Stores the path for the next full-res pass.
|
||
"""
|
||
# app_data is a dict with 'current_path' and 'selections'
|
||
path = os.path.join(
|
||
app_data["current_path"],
|
||
app_data["file_name"]
|
||
)
|
||
self._save_path = path
|
||
self.logger.info(f"Export path set to: {path}")
|
||
dpg.set_value(self.path_label, path)
|
||
|
||
def on_pipeline_data(self, img: np.ndarray):
|
||
# ignore all previews
|
||
return
|
||
|
||
def on_full_res_pipeline_data(self, img: np.ndarray):
|
||
"""
|
||
Receives the full-resolution NumPy image when the user fires
|
||
the “Run full-res pipeline” action. Saves via Pillow.
|
||
"""
|
||
if img is None:
|
||
self.logger.error("on_full_res_pipeline_data called with None image")
|
||
return
|
||
if not self._save_path:
|
||
self.logger.warning("No export path set — click Save As… first")
|
||
return
|
||
|
||
# Decide bit depth by extension
|
||
ext = os.path.splitext(self._save_path)[-1].lower()
|
||
# Convert floats → uint; or leave ints alone
|
||
if np.issubdtype(img.dtype, np.floating):
|
||
if ext in (".tif", ".tiff"):
|
||
arr = np.clip(img * 65535.0, 0, 65535).astype(np.uint16)
|
||
else:
|
||
arr = np.clip(img * 255.0, 0, 255).astype(np.uint8)
|
||
else:
|
||
arr = img
|
||
|
||
# Determine PIL mode
|
||
mode = None
|
||
if arr.ndim == 2:
|
||
mode = "L"
|
||
elif arr.ndim == 3:
|
||
c = arr.shape[2]
|
||
if c == 3:
|
||
mode = "RGB"
|
||
elif c == 4:
|
||
mode = "RGBA"
|
||
|
||
try:
|
||
im = Image.fromarray(arr, mode) if mode else Image.fromarray(arr)
|
||
|
||
# JPEG doesn’t support alpha — drop it
|
||
if ext in (".jpg", ".jpeg") and im.mode == "RGBA":
|
||
im = im.convert("RGB")
|
||
|
||
im.save(self._save_path)
|
||
except Exception as e:
|
||
self.logger.error(f"Failed to save image to {self._save_path}: {e}")
|
||
else:
|
||
self.logger.info(f"Saved full-resolution image to {self._save_path}")
|