Framing widget done
This commit is contained in:
@ -1,7 +1,7 @@
|
||||
import dearpygui.dearpygui as dpg
|
||||
import numpy as np
|
||||
import time
|
||||
from scipy.ndimage import rotate
|
||||
import scipy.ndimage as snd
|
||||
|
||||
from .stage_viewer_widget import PipelineStageViewer
|
||||
|
||||
@ -15,74 +15,71 @@ class FramingWidget(PipelineStageViewer):
|
||||
def __init__(self, manager, logger):
|
||||
super().__init__(manager, logger)
|
||||
|
||||
# Rotation line endpoints (canvas coords)
|
||||
self.rot_start = None # (x, y)
|
||||
self.rot_end = None # (x, y)
|
||||
self.angle = 0.0 # computed deskew angle
|
||||
self.needs_publishing = False
|
||||
|
||||
# Throttle publishing to a few Hz
|
||||
self._last_pub_time = 0.0
|
||||
self._publish_interval = 0.5 # seconds
|
||||
# Rotation line endpoints (canvas coords)
|
||||
self.rot_start = None # (x, y)
|
||||
self.rot_end = None # (x, y)
|
||||
self.angle = 0.0 # computed deskew angle
|
||||
|
||||
# Crop rect endpoints (canvas coords)
|
||||
self.crop_start = None # (x, y)
|
||||
self.crop_end = None # (x, y)
|
||||
|
||||
self.manager.bus.subscribe("img_clicked", self.on_click)
|
||||
self.manager.bus.subscribe("img_dragged", self.on_drag)
|
||||
self.manager.bus.subscribe("img_scrolled", self.on_scroll)
|
||||
self.manager.bus.subscribe("img_released", self.on_release)
|
||||
|
||||
def create_pipeline_stage_content(self):
|
||||
super().create_pipeline_stage_content()
|
||||
|
||||
def on_full_res_pipeline_data(self, img):
|
||||
if img is None:
|
||||
return
|
||||
img = self.rotate(img)
|
||||
self.publish_stage(self.crop(img))
|
||||
|
||||
def on_pipeline_data(self, img):
|
||||
if img is None:
|
||||
return
|
||||
self.img = img.copy()
|
||||
self._publish_rotated_and_cropped()
|
||||
self.original_img = img.copy()
|
||||
self.needs_update = True
|
||||
|
||||
# Apply transformation and publish
|
||||
self.img = self.rotate(img)
|
||||
self.publish_stage(self.crop(self.img))
|
||||
|
||||
def update_texture(self, img):
|
||||
super().update_texture(img)
|
||||
# Draw rotation guide if active
|
||||
|
||||
# Draw rotation guide
|
||||
if self.rot_start and self.rot_end:
|
||||
p0 = self._pos_to_canvas(self.rot_start)
|
||||
p1 = self._pos_to_canvas(self.rot_end)
|
||||
dpg.draw_line(p1=p0, p2=p1, color=(
|
||||
255, 255, 0, 255), thickness=2, parent=self.drawlist)
|
||||
dpg.draw_line(
|
||||
p1=p0,
|
||||
p2=p1,
|
||||
color=(255, 0, 0, 255),
|
||||
thickness=2,
|
||||
parent=self.drawlist,
|
||||
)
|
||||
|
||||
def on_click(self, data):
|
||||
if data.get("obj") is not self:
|
||||
return
|
||||
x, y = data.get("pos")
|
||||
button = data.get("button")
|
||||
if button == "right":
|
||||
self.rot_start = (x, y)
|
||||
self.rot_end = (x, y)
|
||||
self.needs_update = True
|
||||
# Draw crop rect
|
||||
if self.crop_start and self.crop_end:
|
||||
p0 = self._pos_to_canvas(self.crop_start)
|
||||
p1 = self._pos_to_canvas(self.crop_end)
|
||||
dpg.draw_rectangle(
|
||||
pmin=p0,
|
||||
pmax=p1,
|
||||
color=(255, 0, 255, 255),
|
||||
thickness=2,
|
||||
parent=self.drawlist,
|
||||
)
|
||||
|
||||
def on_drag(self, data):
|
||||
if data.get("obj") is not self:
|
||||
return
|
||||
x, y = data.get("pos")
|
||||
button = data.get("button")
|
||||
if button == "right":
|
||||
self.rot_end = (x, y)
|
||||
# Update angle on rotation drag
|
||||
if self.rot_start and self.rot_end:
|
||||
dx = self.rot_end[0] - self.rot_start[0]
|
||||
dy = self.rot_end[1] - self.rot_start[1]
|
||||
self.angle = np.degrees(np.arctan2(dy, dx))
|
||||
# Throttle publishes
|
||||
now = time.time()
|
||||
if now - self._last_pub_time >= self._publish_interval:
|
||||
self._publish_rotated_and_cropped()
|
||||
self._last_pub_time = now
|
||||
self.needs_update = True
|
||||
|
||||
def on_scroll(self, data):
|
||||
print(data)
|
||||
|
||||
def _publish_rotated_and_cropped(self):
|
||||
w, h, _ = self.img.shape
|
||||
out = self.rotate_and_crop(self.img, self.angle, (0, 0, w, h))
|
||||
self.publish_stage(out)
|
||||
if self.needs_publishing:
|
||||
img = self.crop(self.img)
|
||||
self.publish_stage(img)
|
||||
|
||||
def _pos_to_canvas(self, img_pos):
|
||||
x, y = img_pos
|
||||
@ -91,28 +88,91 @@ class FramingWidget(PipelineStageViewer):
|
||||
sw, sh = self.scaled_size
|
||||
return (ix + x / iw * sw, iy + y / ih * sh)
|
||||
|
||||
def rotate_and_crop(
|
||||
self,
|
||||
img: np.ndarray,
|
||||
angle: float,
|
||||
rect: tuple[int, int, int, int],
|
||||
cval: float = 0.0
|
||||
) -> np.ndarray:
|
||||
def on_click(self, data):
|
||||
if data.get("obj") is not self:
|
||||
return
|
||||
x, y = data.get("pos")
|
||||
if data["button"] == "left":
|
||||
self.rot_start = (x, y)
|
||||
self.rot_end = None
|
||||
elif data["button"] == "right":
|
||||
self.crop_start = (x,y)
|
||||
self.crop_end = None
|
||||
self.needs_update = True
|
||||
|
||||
def on_drag(self, data):
|
||||
if data.get("obj") is not self:
|
||||
return
|
||||
x, y = data.get("pos")
|
||||
if data["button"] == "left":
|
||||
self.rot_end = (x, y)
|
||||
elif data["button"] == "right":
|
||||
self.crop_end = (x,y)
|
||||
self.needs_update = True
|
||||
|
||||
def on_scroll(self, data):
|
||||
if data.get("obj") is not self:
|
||||
return
|
||||
|
||||
def on_release(self, data):
|
||||
if data.get("obj") is not self:
|
||||
return
|
||||
if data["button"] == "left":
|
||||
# End of rotation line dragging
|
||||
dx = self.rot_end[0] - self.rot_start[0]
|
||||
dy = self.rot_end[1] - self.rot_start[1]
|
||||
self.angle += np.degrees(np.arctan2(dy, dx))
|
||||
# Do the rotation
|
||||
self.img = self.rotate(self.original_img)
|
||||
# Delete lines
|
||||
self.rot_start = self.rot_end = None
|
||||
self.needs_publishing = True
|
||||
elif data["button"] == "right":
|
||||
self.needs_publishing = True
|
||||
|
||||
self.needs_update = True
|
||||
|
||||
def rotate(self, img):
|
||||
img = snd.rotate(img, self.angle, reshape=False)
|
||||
return img
|
||||
|
||||
def crop(self, img):
|
||||
"""
|
||||
Crop `img` according to a rectangle drawn in the DISPLAY,
|
||||
by converting display‐coords to image‐coords via percentages.
|
||||
"""
|
||||
# only crop if user has dragged out a box
|
||||
if not (self.crop_start and self.crop_end):
|
||||
return img
|
||||
|
||||
# 1) get actual image dims
|
||||
h, w = img.shape[:2]
|
||||
x, y, cw, ch = rect
|
||||
rotated = np.empty_like(img)
|
||||
for c in range(img.shape[2]):
|
||||
rotated[..., c] = rotate(
|
||||
img[..., c],
|
||||
angle,
|
||||
reshape=False,
|
||||
order=1, # bilinear interpolation
|
||||
mode='constant',
|
||||
cval=cval,
|
||||
prefilter=False
|
||||
)
|
||||
x0 = max(0, min(int(x), w - 1))
|
||||
y0 = max(0, min(int(y), h - 1))
|
||||
x1 = max(0, min(int(x + cw), w))
|
||||
y1 = max(0, min(int(y + ch), h))
|
||||
return rotated[y0:y1, x0:x1]
|
||||
|
||||
# 2) get display info (set in update_texture)
|
||||
disp_h, disp_w, _ = self.img.shape
|
||||
|
||||
# 3) unpack the two drag points (in screen coords)
|
||||
sx, sy = self.crop_start
|
||||
ex, ey = self.crop_end
|
||||
|
||||
# 4) normalize into [0..1]
|
||||
nx0 = np.clip(sx / disp_w, 0.0, 1.0)
|
||||
ny0 = np.clip(sy / disp_h, 0.0, 1.0)
|
||||
nx1 = np.clip(ex / disp_w, 0.0, 1.0)
|
||||
ny1 = np.clip(ey / disp_h, 0.0, 1.0)
|
||||
|
||||
# 5) map to image‐pixel coords
|
||||
x0 = int(nx0 * w)
|
||||
y0 = int(ny0 * h)
|
||||
x1 = int(nx1 * w)
|
||||
y1 = int(ny1 * h)
|
||||
|
||||
# 6) sort & clamp again just in case
|
||||
x0, x1 = sorted((max(0, x0), min(w, x1)))
|
||||
y0, y1 = sorted((max(0, y0), min(h, y1)))
|
||||
|
||||
# 7) avoid zero‐area
|
||||
if x0 == x1 or y0 == y1:
|
||||
return img
|
||||
|
||||
return img[y0:y1, x0:x1]
|
||||
Reference in New Issue
Block a user