86 lines
2.6 KiB
Python
86 lines
2.6 KiB
Python
import dearpygui.dearpygui as dpg
|
|
import numpy as np
|
|
from .stage_viewer_widget import PipelineStageViewer
|
|
|
|
|
|
class CropWidget(PipelineStageViewer):
|
|
name = "Crop Image"
|
|
register = True
|
|
has_pipeline_in = True
|
|
has_pipeline_out = True
|
|
|
|
def __init__(self, manager, logger):
|
|
super().__init__(manager, logger)
|
|
self.crop_start = None # (x, y)
|
|
self.crop_end = None # (x, y)
|
|
self.crop_active = False
|
|
|
|
self.manager.bus.subscribe("img_clicked", self.on_click)
|
|
self.manager.bus.subscribe("img_dragged", self.on_drag)
|
|
|
|
def create_pipeline_stage_content(self):
|
|
super().create_pipeline_stage_content()
|
|
|
|
# def on_full_res_pipeline_data(self, img):
|
|
# pass
|
|
|
|
def on_pipeline_data(self, img):
|
|
if img is None:
|
|
return
|
|
self.img = img
|
|
|
|
if self.crop_start and self.crop_end:
|
|
x0, y0 = self.crop_start
|
|
x1, y1 = self.crop_end
|
|
x0, x1 = sorted((int(x0), int(x1)))
|
|
y0, y1 = sorted((int(y0), int(y1)))
|
|
|
|
x0 = max(0, min(x0, img.shape[1]-1))
|
|
x1 = max(0, min(x1, img.shape[1]-1))
|
|
y0 = max(0, min(y0, img.shape[0]-1))
|
|
y1 = max(0, min(y1, img.shape[0]-1))
|
|
|
|
cropped = img[y0:y1, x0:x1, :]
|
|
self.publish_stage(cropped)
|
|
else:
|
|
self.publish_stage(img)
|
|
|
|
self.needs_update = True
|
|
|
|
def on_click(self, data):
|
|
if data["obj"] is not self:
|
|
return
|
|
if data["button"] == "left":
|
|
self.crop_start = data["pos"]
|
|
self.crop_end = data["pos"]
|
|
self.crop_active = True
|
|
self.needs_update = True
|
|
|
|
def on_drag(self, data):
|
|
if data["obj"] is not self or not self.crop_active:
|
|
return
|
|
self.crop_end = data["pos"]
|
|
self.needs_update = True
|
|
|
|
def update_texture(self, img):
|
|
super().update_texture(img)
|
|
if self.crop_start and self.crop_end:
|
|
# map image coords back to screen coords
|
|
x0, y0 = self.crop_start
|
|
x1, y1 = self.crop_end
|
|
h, w, _ = self.img.shape
|
|
img_x, img_y = self.image_position
|
|
img_w, img_h = self.scaled_size
|
|
|
|
p0 = (
|
|
img_x + x0 / w * img_w,
|
|
img_y + y0 / h * img_h
|
|
)
|
|
p1 = (
|
|
img_x + x1 / w * img_w,
|
|
img_y + y1 / h * img_h
|
|
)
|
|
|
|
dpg.draw_rectangle(pmin=p0, pmax=p1, color=(255, 255, 0, 255),
|
|
fill=(255, 255, 0, 50), thickness=2, parent=self.drawlist)
|