diff --git a/.gitignore b/.gitignore index fa98d919..b95acb05 100644 --- a/.gitignore +++ b/.gitignore @@ -6,3 +6,9 @@ parallax/*.log *.npy *.pkl *.csv +*.json +.idea +*.egg-info +*-workspace +*.log +console_history diff --git a/parallax/__init__.py b/parallax/__init__.py index b1b8bd02..4570f55c 100644 --- a/parallax/__init__.py +++ b/parallax/__init__.py @@ -5,5 +5,5 @@ # allow multiple OpenMP instances os.environ['KMP_DUPLICATE_LIB_OK'] = 'True' -# change workdir to package root -os.chdir(os.path.dirname(os.path.realpath(__file__))) +# # change workdir to package root +# os.chdir(os.path.dirname(os.path.realpath(__file__))) diff --git a/parallax/accuracy_test.py b/parallax/accuracy_test.py index 13d50cd5..0834e77f 100644 --- a/parallax/accuracy_test.py +++ b/parallax/accuracy_test.py @@ -93,11 +93,11 @@ def __init__(self, model, parent=None): self.setMinimumWidth(300) def start_accuracy_test(self): - print('TODO: start_accuracy_test') + self.model.start_accuracy_test(self.get_params()) def get_params(self): params = {} - params['stage'] = self.stage_dropdown.get_current_stage() + params['stage'] = self.stage_dropdown.current_stage() params['cal'] = self.model.calibrations[self.cal_dropdown.currentText()] params['npoints'] = int(self.npoints_edit.text()) params['extent_um'] = float(self.extent_edit.text()) @@ -230,9 +230,9 @@ def __init__(self, params): self.ready_to_go = False - def register_corr_points(self, lcorr, rcorr): - xyz_recon = self.cal.triangulate(lcorr, rcorr) - self.results.append(self.last_stage_point + xyz_recon.tolist()) + def register_corr_points(self, corr_pt): + xyz_recon = self.cal.triangulate(corr_pt) + self.results.append(self.last_stage_point + xyz_recon.coordinates.tolist()) def carry_on(self): self.ready_to_go = True diff --git a/parallax/calibration.py b/parallax/calibration.py index 98eb9b9a..73e41153 100755 --- a/parallax/calibration.py +++ b/parallax/calibration.py @@ -1,99 +1,158 @@ -#!/usr/bin/python3 - +import time +import pickle import numpy as np import cv2 as cv +import coorx from . import lib -from .helper import WF, HF - - -imtx1 = [[1.81982227e+04, 0.00000000e+00, 2.59310865e+03], - [0.00000000e+00, 1.89774632e+04, 1.48105977e+03], - [0.00000000e+00, 0.00000000e+00, 1.00000000e+00]] - -imtx2 = [[1.55104298e+04, 0.00000000e+00, 1.95422363e+03], - [0.00000000e+00, 1.54250418e+04, 1.64814750e+03], - [0.00000000e+00, 0.00000000e+00, 1.00000000e+00]] - -idist1 = [[ 1.70600649e+00, -9.85797706e+01, 4.53808433e-03, -2.13200143e-02, 1.79088477e+03]] - -idist2 = [[-4.94883798e-01, 1.65465770e+02, -1.61013572e-03, 5.22601960e-03, -8.73875986e+03]] class Calibration: - def __init__(self, name): - self.name = name - self.set_initial_intrinsics_default() - - def set_name(self, name): - self.name = name - - def set_origin(self, origin): - self.origin = origin + date_format = r"%Y-%m-%d-%H-%M-%S" + file_regex = r'([^-]+)-([^-]+)-((\d\d\d\d-\d\d-\d\d)-(\d+)-(\d+)-(\d+)).pkl' - def get_origin(self): - return self.origin + @staticmethod + def load(filename): + with open(filename, 'rb') as f: + return pickle.load(f) - def set_initial_intrinsics(self, mtx1, mtx2, dist1, dist2): + def __init__(self, img_size): + self.img_size = img_size + self.img_points1 = [] + self.img_points2 = [] + self.obj_points = [] + self.template_images = {} - self.imtx1 = mtx1 - self.imtx2 = mtx2 - self.idist1 = dist1 - self.idist2 = dist2 + def save(self, filename): + with open(filename, 'wb') as f: + pickle.dump(self, f) - def set_initial_intrinsics_default(self): + @property + def name(self): + date = time.strftime(self.date_format, self.timestamp) + return f"{self.from_cs}-{self.to_cs}-{date}" - self.imtx1 = np.array(imtx1, dtype=np.float32) - self.imtx2 = np.array(imtx2, dtype=np.float32) - self.idist1 = np.array(idist1, dtype=np.float32) - self.idist2 = np.array(idist2, dtype=np.float32) + def add_points(self, img_pt1, img_pt2, obj_pt): + self.img_points1.append(img_pt1) + self.img_points2.append(img_pt2) + self.obj_points.append(obj_pt) - def triangulate(self, lcorr, rcorr): + def triangulate(self, img_point): """ l/rcorr = [xc, yc] """ - - img_points1_cv = np.array([[lcorr]], dtype=np.float32) - img_points2_cv = np.array([[rcorr]], dtype=np.float32) - - # undistort - img_points1_cv = lib.undistort_image_points(img_points1_cv, self.mtx1, self.dist1) - img_points2_cv = lib.undistort_image_points(img_points2_cv, self.mtx2, self.dist2) - - img_point1 = img_points1_cv[0,0] - img_point2 = img_points2_cv[0,0] - obj_point_reconstructed = lib.triangulate_from_image_points(img_point1, img_point2, self.proj1, self.proj2) - - return obj_point_reconstructed # [x,y,z] - - def calibrate(self, img_points1, img_points2, obj_points, origin): - - self.set_origin(origin) - + return self.transform.map(img_point) + + def calibrate(self): + cam1 = self.img_points1[0].system.name + cam2 = self.img_points2[0].system.name + stage = self.obj_points[0].system.name + + self.from_cs = f"{cam1}+{cam2}" + self.to_cs = stage + self.camera_names = (cam1, cam2) + self.stage_name = stage + + self.timestamp = time.localtime() + + self.transform = StereoCameraTransform(from_cs=self.from_cs, to_cs=self.to_cs) + self.transform.set_mapping( + np.array(self.img_points1), + np.array(self.img_points2), + np.array(self.obj_points), + self.img_size + ) + + +class CameraTransform(coorx.BaseTransform): + """Maps from camera sensor pixels to undistorted UV. + """ + # initial intrinsic / distortion coefficients + imtx = np.array([ + [4.44851950e+04, 0.00000000e+00, 2.59310865e+03], + [0.00000000e+00, 4.44738925e+04, 1.48105977e+03], + [0.00000000e+00, 0.00000000e+00, 1.00000000e+00]], + dtype='float32' + ) + idist = np.array([[0., 0., 0., 0., 0.]], dtype='float32') + + def __init__(self, mtx=None, dist=None, **kwds): + super().__init__(dims=(2, 2), **kwds) + self.set_coeff(mtx, dist) + + def set_coeff(self, mtx, dist): + self.mtx = mtx + self.dist = dist + self._inverse_transform = None + + def _map(self, pts): + return lib.undistort_image_points(pts, self.mtx, self.dist) + + def _imap(self, pts): + if self._inverse_transform is None: + atr1 = coorx.AffineTransform(matrix=self.mtx[:2, :2], offset=self.mtx[:2, 2]) + ltr1 = coorx.nonlinear.LensDistortionTransform(self.dist[0]) + self._inverse_transform = coorx.CompositeTransform([atr.inverse, ltr, atr]) + return self._inverse_transform.map(pts) + + def set_mapping(self, img_pts, obj_pts, img_size): # undistort calibration points - img_points1 = lib.undistort_image_points(img_points1, self.imtx1, self.idist1) - img_points2 = lib.undistort_image_points(img_points2, self.imtx2, self.idist2) - - # calibrate each camera against these points - my_flags = cv.CALIB_USE_INTRINSIC_GUESS + cv.CALIB_FIX_PRINCIPAL_POINT - rmse1, mtx1, dist1, rvecs1, tvecs1 = cv.calibrateCamera(obj_points, img_points1, - (WF, HF), self.imtx1, self.idist1, - flags=my_flags) - rmse2, mtx2, dist2, rvecs2, tvecs2 = cv.calibrateCamera(obj_points, img_points2, - (WF, HF), self.imtx2, self.idist2, - flags=my_flags) - - # calculate projection matrices - proj1 = lib.get_projection_matrix(mtx1, rvecs1[0], tvecs1[0]) - proj2 = lib.get_projection_matrix(mtx2, rvecs2[0], tvecs2[0]) - - self.mtx1 = mtx1 - self.mtx2 = mtx2 - self.dist1 = dist1 - self.dist2 = dist2 - self.proj1 = proj1 - self.proj2 = proj2 - self.rmse1 = rmse1 - self.rmse2 = rmse2 - - + img_pts_undist = lib.undistort_image_points(img_pts, self.imtx, self.idist) + + # calibrate against correspondence points + rmse, mtx, dist, rvecs, tvecs = cv.calibrateCamera( + obj_pts.astype('float32')[np.newaxis, ...], + img_pts_undist[np.newaxis, ...], + img_size, + self.imtx, self.idist, + flags=cv.CALIB_USE_INTRINSIC_GUESS + cv.CALIB_FIX_PRINCIPAL_POINT, + ) + + # calculate projection matrix + self.proj_matrix = lib.get_projection_matrix(mtx, rvecs[0], tvecs[0]) + + # record results + self.set_coeff(mtx, dist) + self.calibration_result = {'rmse': rmse, 'mtx': mtx, 'dist': dist, 'rvecs': rvecs, 'tvecs': tvecs} + + +class StereoCameraTransform(coorx.BaseTransform): + """Maps from dual camera sensor pixels to 3D object space. + """ + def __init__(self, **kwds): + super().__init__(dims=(4, 3), **kwds) + self.camera_tr1 = CameraTransform() + self.camera_tr2 = CameraTransform() + self.proj1 = None + self.proj2 = None + + def set_mapping(self, img_points1, img_points2, obj_points, img_size): + self.camera_tr1.set_mapping(img_points1, obj_points, img_size) + self.camera_tr2.set_mapping(img_points2, obj_points, img_size) + + self.proj1 = self.camera_tr1.proj_matrix + self.proj2 = self.camera_tr2.proj_matrix + + self.rmse1 = self.camera_tr1.calibration_result['rmse'] + self.rmse2 = self.camera_tr2.calibration_result['rmse'] + + def triangulate(self, img_point1, img_point2): + x,y,z = lib.DLT(self.proj1, self.proj2, img_point1, img_point2) + return np.array([x,y,z]) + + def _map(self, arr2d): + # undistort + img_pts1 = self.camera_tr1.map(arr2d[:, 0:2]) + img_pts2 = self.camera_tr2.map(arr2d[:, 2:4]) + + # triangulate + n_pts = arr2d.shape[0] + obj_points = [self.triangulate(*img_pts) for img_pts in zip(img_pts1, img_pts2)] + return np.vstack(obj_points) + + def _imap(self, arr2d): + itr1, itr2 = self._inverse_transforms + ret = np.empty((len(arr2d), 4)) + ret[:, 0:2] = itr1.map(arr2d) + ret[:, 2:4] = itr2.map(arr2d) + return ret diff --git a/parallax/calibration_worker.py b/parallax/calibration_worker.py index f0fd6ab2..16eb0787 100644 --- a/parallax/calibration_worker.py +++ b/parallax/calibration_worker.py @@ -1,41 +1,45 @@ +import time from PyQt5.QtCore import QObject, pyqtSignal import numpy as np -import time +import cv2 as cv +import queue +from .calibration import Calibration +from .helper import WF, HF +from .model import Model +from .config import config class CalibrationWorker(QObject): finished = pyqtSignal() calibration_point_reached = pyqtSignal(int, int, float, float, float) + suggested_corr_points = pyqtSignal(object) RESOLUTION_DEFAULT = 3 EXTENT_UM_DEFAULT = 2000 - def __init__(self, name, stage, resolution=RESOLUTION_DEFAULT, extent_um=EXTENT_UM_DEFAULT, + template_match_method = cv.TM_CCORR_NORMED + template_radius = 200 + + def __init__(self, stage, cameras, resolution=RESOLUTION_DEFAULT, extent_um=EXTENT_UM_DEFAULT, parent=None): # resolution is number of steps per dimension, for 3 dimensions # (so default value of 3 will yield 3^3 = 27 calibration points) # extent_um is the extent in microns for each dimension, centered on zero QObject.__init__(self) - self.name = name self.stage = stage + self.cameras = cameras self.resolution = resolution self.extent_um = extent_um - self.object_points = [] # units are mm self.num_cal = self.resolution**3 - - self.img_points1 = [] - self.img_points2 = [] + self.calibration = Calibration(img_size=(WF, HF)) + self.corr_point_queue = queue.Queue() self.complete = False self.alive = True def register_corr_points(self, lcorr, rcorr): - self.img_points1.append(lcorr) - self.img_points2.append(rcorr) - - def carry_on(self): - self.ready_to_go = True + self.corr_point_queue.put((lcorr, rcorr)) def stop(self): self.alive = False @@ -47,14 +51,37 @@ def run(self): for x in np.linspace(mn, mx, self.resolution): for y in np.linspace(mn, mx, self.resolution): for z in np.linspace(mn, mx, self.resolution): + + # move stage to next point and let the user know we are ready for clicks self.stage.move_to_target_3d(x,y,z, relative=True, safe=False) - self.calibration_point_reached.emit(n,self.num_cal, x,y,z) - self.ready_to_go = False - while self.alive and not self.ready_to_go: - time.sleep(0.1) - if not self.alive: - break - self.object_points.append([x,y,z]) + self.calibration_point_reached.emit(n, self.num_cal, x, y, z) + + # If we are using a mock stage, automatically + # select known correspondence points + if config['mock_sim']['auto_select_corr_points'] and hasattr(self.stage, 'get_tip_position'): + tip_pos = self.stage.get_tip_position() + pts = {} + for screen in Model.instance.main_window.screens(): + camera = screen.camera + pos = camera.camera_tr.map(tip_pos.coordinates) + pts[camera.name()] = pos[:2] + self.suggested_corr_points.emit(pts) + else: + # Attempt to select correspondence points based on template match + if n > 0: + time.sleep(0.25) # let camera catch up + self.match_templates() + + # wait for user to confirm correspondence points + lcorr, rcorr = self.corr_point_queue.get() + + if n == 0: + # Add template images to calibration (maybe this could be a + # running average instead of one-time?) + self.collect_templates([lcorr, rcorr]) + + # add points to calibration + self.calibration.add_points(lcorr, rcorr, self.stage.get_position()) n += 1 else: continue @@ -66,10 +93,63 @@ def run(self): self.complete = True self.finished.emit() - def get_image_points(self): - return np.array([self.img_points1], dtype=np.float32), \ - np.array([self.img_points2], dtype=np.float32) + def get_calibration(self): + self.calibration.calibrate() + return self.calibration + + def collect_templates(self, pts): + r = self.template_radius + for cam, pt in zip(self.cameras, pts): + img = cam.get_last_image_data() + row, col = int(pt[1]), int(pt[0]) + template = img[row-r:row+r, col-r:col+r] + self.calibration.template_images[cam.name()] = template + + def match_templates(self): + method = self.template_match_method + result = {} + for cam in self.cameras: + img = cam.get_last_image_data() + template = self.calibration.template_images[cam.name()] + res, mx = fast_template_match(img, template, method) + result[cam.name()] = mx[::-1] + self.template_radius + + self.suggested_corr_points.emit(result) + + +def template_match(img, template, method): + # should we look for min or max in match results + ext_method = 'argmax' + if method in (cv.TM_SQDIFF, cv.TM_SQDIFF_NORMED): + ext_method = 'argmin' + res = cv.matchTemplate(img, template, method) + ext = getattr(res, ext_method)() + mx = np.array(np.unravel_index(ext, res.shape)) + return res, mx + + +def fast_template_match(img, template, method, downsample=10): + """2-stage template match for better performance + """ + + # first convert to greyscale + img = img.mean(axis=2).astype('ubyte') + template = template.mean(axis=2).astype('ubyte') + + # do a quick template match on downsampled images + img2 = img[::downsample, ::downsample] + template2 = template[::downsample, ::downsample] + + res, mx = template_match(img2, template2, method) + mx = mx * downsample + + crop = [ + (max(0, mx[0] - downsample*2), mx[0] + template.shape[0] + downsample*2), + (max(0, mx[1] - downsample*2), mx[1] + template.shape[1] + downsample*2), + ] - def get_object_points(self): - return np.array([self.object_points], dtype=np.float32) + img3 = img[crop[0][0]:crop[0][1], crop[1][0]:crop[1][1]] + res, mx = template_match(img3, template, method) + mx = mx + [crop[0][0], crop[1][0]] + return res, mx diff --git a/parallax/camera.py b/parallax/camera.py index 7c2a828a..0367b407 100755 --- a/parallax/camera.py +++ b/parallax/camera.py @@ -3,6 +3,8 @@ import threading import numpy as np import logging +import pyqtgraph as pg +from .mock_sim import MockSim logger = logging.getLogger(__name__) @@ -15,12 +17,13 @@ def list_cameras(): - global pyspin_cameras, pyspin_instance cameras = [] + cameras.extend([ + MockCamera(camera_params={'pitch': 70, 'yaw': 120}), + MockCamera(camera_params={'pitch': 70, 'yaw': 150}), + ]) if PySpin is not None: cameras.extend(PySpinCamera.list_cameras()) - while len(cameras) < 2: - cameras.append(MockCamera()) return cameras @@ -160,12 +163,24 @@ def capture_loop(self): class MockCamera: n_cameras = 0 - def __init__(self): - self._name = f"MockCamera{MockCamera.n_cameras}" + def __init__(self, camera_params): + self._name = f"mock_camera_{MockCamera.n_cameras}" MockCamera.n_cameras += 1 - self.data = np.random.randint(0, 255, size=(5, 3000, 4000), dtype='ubyte') + self.sensor_size = (4000, 3000) + self.noise = np.random.randint(0, 15, size=(5, self.sensor_size[1], self.sensor_size[0], 3), dtype='ubyte') self._next_frame = 0 + self.camera_params = dict( + pitch=30, + yaw=0, + fov=5, + distance=200e3, + distortion=(-0.1, 0, 0, 0, 0), + ) + self.camera_params.update(camera_params) + self.sim = MockSim.instance() + self.sim.add_camera(self) + def name(self): return self._name @@ -173,6 +188,26 @@ def get_last_image_data(self): """ Return last image as numpy array with shape (height, width, 3) for RGB or (height, width) for mono. """ - frame = self.data[self._next_frame] - self._next_frame = (self._next_frame + 1) % self.data.shape[0] - return frame + + image = self.sim.get_camera_frame(self) + + # # add noise + # noise = self.noise[self._next_frame] + # self._next_frame = (self._next_frame + 1) % self.noise.shape[0] + # image = image.copy() + # # squeeze to leave room for noise (and prevent overflows) + # np.multiply(image, 224/255, out=image, casting='unsafe') + # image += noise + + return image + + def save_last_image(self, filename): + arr = self.get_last_image_data()[..., [2,1,0]] + img = pg.makeQImage(arr, alpha=False, transpose=False) + img.save(filename) + + @property + def camera_tr(self): + """Transform mapping from simulated global coordinate system to camera image pixels. + """ + return self.sim.cameras[self]['view'].camera_tr diff --git a/parallax/config.py b/parallax/config.py new file mode 100644 index 00000000..25c99a06 --- /dev/null +++ b/parallax/config.py @@ -0,0 +1,42 @@ +import os, json, argparse + +package_path = os.path.split(os.path.dirname(__file__))[0] + +# global configuration +config = { + "views": [{}, {}], + "cameras": [], + "stages": [], + "mock_sim": { + "show_checkers": True, + "show_axes": True, + "auto_select_corr_points": True, + }, + "calibration_path": "./", + "console_history_file": os.path.join(package_path, 'console.log'), + "console_edit_command": "code -g {fileName}:{lineNum}", +} + + +def parse_cli_args(): + parser = argparse.ArgumentParser(prog='parallax') + parser.add_argument('--config', type=str, default=None, help='configuration file to load at startup') + args = parser.parse_args() + return args + + +def init_config(args): + global config + if args.config is not None: + loaded_config = json.load(open(args.config, 'r')) + for k,v in loaded_config.items(): + if k not in config: + raise KeyError(f"Invalid config key {k}") + config[k] = v + +def post_init_config(model, main_window): + for i,view in enumerate(config['views']): + screen_widget = main_window.widget.add_screen() + if 'default_camera' in view: + camera = model.get_camera(view['default_camera']) + screen_widget.set_camera(camera) diff --git a/parallax/control_panel.py b/parallax/control_panel.py index 7b3dbf22..96fb0884 100644 --- a/parallax/control_panel.py +++ b/parallax/control_panel.py @@ -1,11 +1,16 @@ +import time from PyQt5.QtWidgets import QWidget, QLabel, QPushButton, QFrame from PyQt5.QtWidgets import QVBoxLayout, QGridLayout from PyQt5.QtCore import pyqtSignal, Qt -from PyQt5.QtGui import QIcon +from PyQt5.QtGui import QIcon, QGuiApplication +import pyqtgraph as pg +import coorx from .helper import FONT_BOLD from .dialogs import StageSettingsDialog, TargetDialog from .stage_dropdown import StageDropdown +from .calibration import Calibration +from .config import config JOG_UM_DEFAULT = 250 CJOG_UM_DEFAULT = 50 @@ -53,7 +58,7 @@ class ControlPanel(QFrame): def __init__(self, model): QFrame.__init__(self) self.model = model - + self.calibration = None # widgets self.main_label = QLabel('Stage Control') @@ -64,9 +69,13 @@ def __init__(self, model): self.dropdown.activated.connect(self.handle_stage_selection) self.settings_button = QPushButton() - self.settings_button.setIcon(QIcon('../img/gear.png')) + self.settings_button.setIcon(QIcon('./img/gear.png')) self.settings_button.clicked.connect(self.handle_settings) + self.calibration_label = QLabel("") + self.cal_pt_btn = QPushButton('Copy Calibration Point') + self.cal_pt_btn.clicked.connect(self.copy_cal_pt) + self.xcontrol = AxisControl('x') self.xcontrol.jog_requested.connect(self.jog) self.xcontrol.center_requested.connect(self.center) @@ -80,19 +89,39 @@ def __init__(self, model): self.zero_button = QPushButton('Set Relative Origin') self.zero_button.clicked.connect(self.zero) - self.move_target_button = QPushButton('Move to Target') + self.move_target_button = QPushButton('Move to ...') self.move_target_button.clicked.connect(self.move_to_target) + self.move_selected_button = QPushButton('Move to Selected') + self.move_selected_button.clicked.connect(self.move_to_selected) + + self.approach_selected_button = QPushButton('Approach Selected') + self.approach_selected_button.clicked.connect(self.approach_selected) + self.approach_distance_spin = pg.SpinBox(value=3e-3, suffix='m', siPrefix=True, bounds=[1e-6, 20e-3], dec=True, step=0.5, minStep=1e-6, compactHeight=False) + + self.move_to_depth_button = QPushButton('Advance Depth') + self.move_to_depth_button.clicked.connect(self.move_to_depth) + self.depth_spin = pg.SpinBox(value=1e-3, suffix='m', siPrefix=True, bounds=[1e-6, 20e-3], dec=True, step=0.5, minStep=1e-6, compactHeight=False) + self.depth_speed_spin = pg.SpinBox(value=10e-6, suffix='m/s', siPrefix=True, bounds=[1e-6, 1e-3], dec=True, step=0.5, minStep=1e-6, compactHeight=False) + # layout main_layout = QGridLayout() main_layout.addWidget(self.main_label, 0,0, 1,3) main_layout.addWidget(self.dropdown, 1,0, 1,2) main_layout.addWidget(self.settings_button, 1,2, 1,1) - main_layout.addWidget(self.xcontrol, 2,0, 1,1) - main_layout.addWidget(self.ycontrol, 2,1, 1,1) - main_layout.addWidget(self.zcontrol, 2,2, 1,1) - main_layout.addWidget(self.zero_button, 3,0, 1,3) - main_layout.addWidget(self.move_target_button, 4,0, 1,3) + main_layout.addWidget(self.calibration_label, 2,0, 1,2) + main_layout.addWidget(self.cal_pt_btn, 2,2, 1,1) + main_layout.addWidget(self.xcontrol, 3,0, 1,1) + main_layout.addWidget(self.ycontrol, 3,1, 1,1) + main_layout.addWidget(self.zcontrol, 3,2, 1,1) + main_layout.addWidget(self.zero_button, 4,0, 1,3) + main_layout.addWidget(self.move_target_button, 5,0, 1,1) + main_layout.addWidget(self.move_selected_button, 5,1, 1,2) + main_layout.addWidget(self.approach_selected_button, 6,0, 1,1) + main_layout.addWidget(self.approach_distance_spin, 6,1, 1,1) + main_layout.addWidget(self.move_to_depth_button, 7,0, 1,1) + main_layout.addWidget(self.depth_spin, 7,1, 1,1) + main_layout.addWidget(self.depth_speed_spin, 7,2, 1,1) self.setLayout(main_layout) # frame border @@ -103,6 +132,9 @@ def __init__(self, model): self.jog_um = JOG_UM_DEFAULT self.cjog_um = CJOG_UM_DEFAULT + self.model.calibrations_changed.connect(self.update_calibration) + self.model.corr_pts_changed.connect(self.update_cal_pt) + def update_coordinates(self, *args): xa, ya, za = self.stage.get_position() xo, yo, zo = self.stage.get_origin() @@ -115,34 +147,42 @@ def update_relative_origin(self): self.zero_button.setText('Set Relative Origin: (%d %d %d)' % (x, y, z)) def handle_stage_selection(self, index): - stage_name = self.dropdown.currentText() - self.set_stage(self.model.stages[stage_name]) + stage = self.dropdown.current_stage() + self.set_stage(stage) self.update_coordinates() + self.update_calibration() def set_stage(self, stage): + if self.stage is not None: + self.stage.position_changed.disconnect(self.stage_position_changed) self.stage = stage + self.stage.position_changed.connect(self.stage_position_changed) self.update_relative_origin() def move_to_target(self, *args): dlg = TargetDialog(self.model) if dlg.exec_(): params = dlg.get_params() - x = params['x'] - y = params['y'] - z = params['z'] + pt = params['point'] if self.stage: - self.stage.move_to_target_3d(x, y, z, relative=params['relative'], safe=True) - if params['relative']: - self.msg_posted.emit('Moved to relative position: ' - '[{0:.2f}, {1:.2f}, {2:.2f}]'.format(x, y, z)) - else: - self.msg_posted.emit('Moved to absolute position: ' - '[{0:.2f}, {1:.2f}, {2:.2f}]'.format(x, y, z)) + self.move_to_point(pt, params['relative']) self.update_coordinates() self.target_reached.emit() else: self.msg_posted.emit('Move to target: no stage selected.') + def move_to_point(self, pt, relative, **kwds): + if isinstance(pt, coorx.Point): + if pt.system.name != self.stage.get_name(): + raise Exception(f"Not moving stage {self.stage.get_name()} to coordinate in system {pt.system.name}") + if relative: + raise Exception(f"Not moving to relative point in system {pt.system.name}") + + x, y, z = pt + self.stage.move_to_target_3d(x, y, z, relative=relative, safe=True, **kwds) + absrel = "relative" if relative else "absolute" + self.msg_posted.emit(f'Moved to {absrel} position: [{x:.2f}, {y:.2f}, {z:.2f}]') + def handle_settings(self, *args): if self.stage: dlg = StageSettingsDialog(self.stage, self.jog_um, self.cjog_um) @@ -184,4 +224,58 @@ def halt(self): # doesn't actually work now because we need threading self.stage.halt() + def stage_position_changed(self, stage, pos): + self.update_coordinates() + + def update_calibration(self): + if self.stage is None: + return + + # search for and load calibration appropriate for the selected stage + cal = self.model.get_calibration(self.stage) + self.calibration = cal + if cal is None: + self.calibration_label.setText('(no calibration)') + else: + ts_str = time.strftime(r"%Y-%m-%d %H:%M:%S", cal.timestamp) + self.calibration_label.setText(f'calibrated {ts_str} for {cal.to_cs}') + + def get_stage_point(self): + if self.stage is None: + raise Exception("No stage selected") + if self.calibration is None: + raise Exception(f"No calibration loaded for {self.stage.get_name()}") + img_pt = self.model.get_image_point() + if img_pt is None: + raise Exception("No correspondence points selected") + return self.calibration.triangulate(img_pt) + + def move_to_selected(self): + stage_pt = self.get_stage_point() + self.move_to_point(stage_pt, relative=False, block=False) + + def approach_selected(self): + stage_pt = self.get_stage_point() + depth = self.approach_distance_spin.value() * 1e6 + stage_pt.coordinates[2] += depth + self.move_to_point(stage_pt, relative=False, block=False) + + def move_to_depth(self): + pos = self.stage.get_position() + depth = self.depth_spin.value() * 1e6 + speed = self.depth_speed_spin.value() * 1e6 + pos.coordinates[2] -= depth + self.move_to_point(pos, relative=False, speed=speed, block=False) + + def update_cal_pt(self): + try: + stage_pt = self.get_stage_point() + x = stage_pt.coordinates + self.cal_pt_btn.setText(f'{x[0]:0.0f}, {x[1]:0.0f}, {x[2]:0.0f}') + except Exception: + self.cal_pt_btn.setText('') + + def copy_cal_pt(self): + cb = QGuiApplication.clipboard() + cb.setText(self.cal_pt_btn.text()) diff --git a/parallax/dialogs.py b/parallax/dialogs.py index 06ff4fac..3cd6b1e3 100755 --- a/parallax/dialogs.py +++ b/parallax/dialogs.py @@ -1,12 +1,10 @@ -from PyQt5.QtWidgets import QPushButton, QLabel, QRadioButton, QSpinBox +from PyQt5.QtWidgets import QPushButton, QLabel, QSpinBox from PyQt5.QtWidgets import QGridLayout from PyQt5.QtWidgets import QDialog, QLineEdit, QDialogButtonBox from PyQt5.QtCore import Qt from PyQt5.QtGui import QDoubleValidator - +import pyqtgraph as pg import numpy as np -import time -import datetime from .toggle_switch import ToggleSwitch from .helper import FONT_BOLD @@ -117,13 +115,6 @@ def __init__(self, model, parent=None): self.extent_label.setAlignment(Qt.AlignCenter) self.extent_edit = QLineEdit(str(cw.EXTENT_UM_DEFAULT)) - self.name_label = QLabel('Name') - ts = time.time() - dt = datetime.datetime.fromtimestamp(ts) - cal_default_name = 'cal_%04d%02d%02d-%02d%02d%02d' % (dt.year, - dt.month, dt.day, dt.hour, dt.minute, dt.second) - self.name_edit = QLineEdit(cal_default_name) - self.go_button = QPushButton('Start Calibration Routine') self.go_button.setEnabled(False) self.go_button.clicked.connect(self.go) @@ -135,8 +126,6 @@ def __init__(self, model, parent=None): layout.addWidget(self.resolution_box, 1,1, 1,1) layout.addWidget(self.extent_label, 2,0, 1,1) layout.addWidget(self.extent_edit, 2,1, 1,1) - layout.addWidget(self.name_label, 3,0, 1,1) - layout.addWidget(self.name_edit, 3,1, 1,1) layout.addWidget(self.go_button, 4,0, 1,2) self.setLayout(layout) @@ -144,9 +133,7 @@ def __init__(self, model, parent=None): self.setMinimumWidth(300) def get_stage(self): - ip = self.stage_dropdown.currentText() - stage = self.model.stages[ip] - return stage + return self.stage_dropdown.current_stage() def get_resolution(self): return self.resolution_box.value() @@ -154,9 +141,6 @@ def get_resolution(self): def get_extent(self): return float(self.extent_edit.text()) - def get_name(self): - return self.name_edit.text() - def go(self): self.accept() @@ -174,6 +158,7 @@ def __init__(self, model): QDialog.__init__(self) self.model = model + self.obj_point = None self.last_button = QPushButton('Last Reconstructed Point') self.last_button.clicked.connect(self.populate_last) if self.model.obj_point_last is None: @@ -207,10 +192,17 @@ def __init__(self, model): self.yedit.textChanged.connect(self.validate) self.zedit.textChanged.connect(self.validate) - self.info_label = QLabel('(units are microns)') + self.xedit.textEdited.connect(self.input_changed) + self.yedit.textEdited.connect(self.input_changed) + self.zedit.textEdited.connect(self.input_changed) + self.abs_rel_toggle.toggled.connect(self.input_changed) + + self.info_label = QLabel('') self.info_label.setAlignment(Qt.AlignCenter) self.info_label.setFont(FONT_BOLD) + self.update_info() + self.cancel_button = QPushButton('Cancel') self.cancel_button.clicked.connect(self.reject) @@ -249,23 +241,41 @@ def validate(self): self.ok_button.setEnabled(valid) def populate_last(self): - self.xedit.setText('{0:.2f}'.format(self.model.obj_point_last[0])) - self.yedit.setText('{0:.2f}'.format(self.model.obj_point_last[1])) - self.zedit.setText('{0:.2f}'.format(self.model.obj_point_last[2])) + op = self.model.obj_point_last + self.xedit.setText('{0:.2f}'.format(op[0])) + self.yedit.setText('{0:.2f}'.format(op[1])) + self.zedit.setText('{0:.2f}'.format(op[2])) + self.abs_rel_toggle.setChecked(False) + self.obj_point = op + self.update_info() def populate_random(self): + self.obj_point = None self.xedit.setText('{0:.2f}'.format(np.random.uniform(-2000, 2000))) self.yedit.setText('{0:.2f}'.format(np.random.uniform(-2000, 2000))) self.zedit.setText('{0:.2f}'.format(np.random.uniform(-2000, 2000))) + self.update_info() def get_params(self): params = {} - params['x'] = float(self.xedit.text()) - params['y'] = float(self.yedit.text()) - params['z'] = float(self.zedit.text()) - params['relative'] = self.abs_rel_toggle.isChecked() + if self.obj_point is None: + params['point'] = np.array([float(self.xedit.text()), float(self.yedit.text()), float(self.zedit.text())]) + params['relative'] = self.abs_rel_toggle.isChecked() + else: + params['point'] = self.obj_point + params['relative'] = False return params + def update_info(self): + info = "(units are μm)" + if self.obj_point is not None: + info = f'coord sys: {self.obj_point.system.name}\n{info}' + self.info_label.setText(info) + + def input_changed(self): + self.obj_point = None + self.update_info() + class CsvDialog(QDialog): @@ -362,3 +372,61 @@ def get_params(self): y = float(self.yedit.text()) z = float(self.zedit.text()) return x,y,z + + +class TrainingDataDialog(QDialog): + + def __init__(self, model): + QDialog.__init__(self) + self.model = model + + self.setWindowTitle('Training Data Generator') + + self.stage_label = QLabel('Select a Stage:') + self.stage_label.setAlignment(Qt.AlignCenter) + self.stage_label.setFont(FONT_BOLD) + + self.stage_dropdown = StageDropdown(self.model) + self.stage_dropdown.activated.connect(self.update_status) + + self.img_count_label = QLabel('Image Count:') + self.img_count_label.setAlignment(Qt.AlignCenter) + self.img_count_box = QSpinBox() + self.img_count_box.setMinimum(1) + self.img_count_box.setValue(100) + + self.extent_label = QLabel('Extent:') + self.extent_label.setAlignment(Qt.AlignCenter) + self.extent_spin = pg.SpinBox(value=4e-3, suffix='m', siPrefix=True, bounds=[0.1e-3, 20e-3], dec=True, step=0.5, minStep=1e-6, compactHeight=False) + + self.go_button = QPushButton('Start Data Collection') + self.go_button.setEnabled(False) + self.go_button.clicked.connect(self.go) + + layout = QGridLayout() + layout.addWidget(self.stage_label, 0,0, 1,1) + layout.addWidget(self.stage_dropdown, 0,1, 1,1) + layout.addWidget(self.img_count_label, 1,0, 1,1) + layout.addWidget(self.img_count_box, 1,1, 1,1) + layout.addWidget(self.extent_label, 2,0, 1,1) + layout.addWidget(self.extent_spin, 2,1, 1,1) + layout.addWidget(self.go_button, 4,0, 1,2) + self.setLayout(layout) + + self.setMinimumWidth(300) + + def get_stage(self): + return self.stage_dropdown.current_stage() + + def get_img_count(self): + return self.img_count_box.value() + + def get_extent(self): + return self.extent_spin.value() * 1e6 + + def go(self): + self.accept() + + def update_status(self): + if self.stage_dropdown.is_selected(): + self.go_button.setEnabled(True) diff --git a/parallax/geometry_panel.py b/parallax/geometry_panel.py index 9d1d118c..225e30aa 100644 --- a/parallax/geometry_panel.py +++ b/parallax/geometry_panel.py @@ -1,5 +1,5 @@ -from PyQt5.QtWidgets import QGridLayout, QVBoxLayout, QHBoxLayout -from PyQt5.QtWidgets import QPushButton, QFrame, QWidget, QComboBox, QLabel +from PyQt5.QtWidgets import QGridLayout, QVBoxLayout +from PyQt5.QtWidgets import QPushButton, QFrame, QComboBox, QLabel from PyQt5.QtWidgets import QFileDialog from PyQt5.QtCore import pyqtSignal, Qt, QThread @@ -9,9 +9,9 @@ from .helper import FONT_BOLD from .dialogs import CalibrationDialog from .rigid_body_transform_tool import RigidBodyTransformTool, PointTransformWidget -from .calibration import Calibration from .calibration_worker import CalibrationWorker - +from .calibration import Calibration +from .config import config class GeometryPanel(QFrame): msg_posted = pyqtSignal(str) @@ -74,25 +74,29 @@ def __init__(self, model): self.setLineWidth(2) def triangulate(self): - - if (self.cal_combo.currentIndex() < 0): + cal_selected = self.selected_calibration() + if cal_selected is None: self.msg_posted.emit('No calibration selected.') return - else: - cal_selected = self.model.calibrations[self.cal_combo.currentText()] - if not (self.model.lcorr and self.model.rcorr): + corr_pt = self.model.get_image_point() + if corr_pt is None: self.msg_posted.emit('No correspondence points selected.') return - else: - lcorr, rcorr = self.model.lcorr, self.model.rcorr - obj_point = cal_selected.triangulate(lcorr, rcorr) + obj_point = cal_selected.triangulate(corr_pt) self.model.set_last_object_point(obj_point) x,y,z = obj_point - self.msg_posted.emit('Reconstructed object point: ' - '[{0:.2f}, {1:.2f}, {2:.2f}]'.format(x, y, z)) + self.msg_posted.emit( + f'Reconstructed object point: [{x:.2f}, {y:.2f}, {z:.2f}] in {obj_point.system.name}' + ) + + def selected_calibration(self): + if (self.cal_combo.currentIndex() < 0): + return None + else: + return self.model.calibrations[self.cal_combo.currentText()] def cal_start_stop(self): if self.cal_start_stop_button.text() == 'Start': @@ -101,18 +105,18 @@ def cal_start_stop(self): stage = dlg.get_stage() res = dlg.get_resolution() extent = dlg.get_extent() - name = dlg.get_name() - self.start_cal_thread(stage, res, extent, name) + self.start_cal_thread(stage, res, extent) elif self.cal_start_stop_button.text() == 'Stop': self.stop_cal_thread() - def start_cal_thread(self, stage, res, extent, name): + def start_cal_thread(self, stage, res, extent): self.model.cal_in_progress = True self.cal_thread = QThread() - self.cal_worker = CalibrationWorker(name, stage, res, extent) + self.cal_worker = CalibrationWorker(stage, self.model.cameras, res, extent) self.cal_worker.moveToThread(self.cal_thread) self.cal_thread.started.connect(self.cal_worker.run) self.cal_worker.calibration_point_reached.connect(self.handle_cal_point_reached) + self.cal_worker.suggested_corr_points.connect(self.show_suggested_corr_points) self.cal_thread.finished.connect(self.handle_cal_finished) self.cal_worker.finished.connect(self.cal_thread.quit) self.cal_thread.finished.connect(self.cal_thread.deleteLater) @@ -131,23 +135,18 @@ def handle_cal_point_reached(self, n, num_cal, x,y,z): def register_corr_points_cal(self): lcorr, rcorr = self.model.lcorr, self.model.rcorr - if (lcorr and rcorr): + if None not in (lcorr, rcorr): self.cal_worker.register_corr_points(lcorr, rcorr) self.msg_posted.emit('Correspondence points registered: (%d,%d) and (%d,%d)' % \ (lcorr[0],lcorr[1], rcorr[0],rcorr[1])) - self.cal_worker.carry_on() else: self.msg_posted.emit('Highlight correspondence points and press C to continue') def handle_cal_finished(self): if self.cal_worker.complete: - cal = Calibration(self.cal_worker.name) - img_points1, img_points2 = self.cal_worker.get_image_points() - obj_points = self.cal_worker.get_object_points() - origin = self.cal_worker.stage.get_origin() - cal.calibrate(img_points1, img_points2, obj_points, origin) + cal = self.cal_worker.get_calibration() self.msg_posted.emit('Calibration finished. RMSE1 = %f, RMSE2 = %f' % \ - (cal.rmse1, cal.rmse2)) + (cal.transform.rmse1, cal.transform.rmse2)) self.model.add_calibration(cal) self.update_cals() else: @@ -156,29 +155,26 @@ def handle_cal_finished(self): self.cal_start_stop_button.setText('Start') def load_cal(self): - filename = QFileDialog.getOpenFileName(self, 'Load calibration file', '.', + filename = QFileDialog.getOpenFileName(self, 'Load calibration file', config['calibration_path'], 'Pickle files (*.pkl)')[0] if filename: - with open(filename, 'rb') as f: - cal = pickle.load(f) - self.model.add_calibration(cal) + cal = Calibration.load(filename) + self.model.add_calibration(cal) self.update_cals() def save_cal(self): - if (self.cal_combo.currentIndex() < 0): + cal_selected = self.selected_calibration() + if cal_selected is None: self.msg_posted.emit('No calibration selected.') return - else: - cal_selected = self.model.calibrations[self.cal_combo.currentText()] - suggested_filename = os.path.join(os.getcwd(), cal_selected.name + '.pkl') + suggested_filename = os.path.join(config['calibration_path'], cal_selected.name + '.pkl') filename = QFileDialog.getSaveFileName(self, 'Save calibration file', suggested_filename, 'Pickle files (*.pkl)')[0] if filename: - with open(filename, 'wb') as f: - pickle.dump(cal_selected, f) + cal_selected.save(filename) self.msg_posted.emit('Saved calibration %s to: %s' % (cal_selected.name, filename)) def update_cals(self): @@ -231,3 +227,8 @@ def show_rbt_tool(self): self.rbt_tool.generated.connect(self.update_transforms) self.rbt_tool.show() + def show_suggested_corr_points(self, pts): + screens = {screen.camera.name():screen for screen in self.model.main_window.screens()} + for cam_name, pt in pts.items(): + screens[cam_name].set_selected(pt) + diff --git a/parallax/lib.py b/parallax/lib.py index 9e1d150a..41bb80b8 100644 --- a/parallax/lib.py +++ b/parallax/lib.py @@ -25,7 +25,7 @@ def undistort_image_points(img_points, mtx, dist): - img_points_corrected_normalized = cv.undistortPoints(img_points, mtx, dist) + img_points_corrected_normalized = cv.undistortPoints(img_points.astype('float32'), mtx, dist) fx = mtx[0,0] fy = mtx[1,1] cx = mtx[0,2] @@ -35,8 +35,8 @@ def undistort_image_points(img_points, mtx, dist): x,y = img_point[0] x = x * fx + cx y = y * fy + cy - img_points_corrected.append(np.array([x,y])) - return np.array([img_points_corrected], dtype=np.float32) + img_points_corrected.append([x, y]) + return np.array(img_points_corrected, dtype=np.float32) def get_projection_matrix(mtx, r, t): R, jacobian = cv.Rodrigues(r) @@ -59,6 +59,31 @@ def DLT(P1, P2, point1, point2): U, s, vh = linalg.svd(B, full_matrices = False) return vh[3,0:3]/vh[3,3] -def triangulate_from_image_points(img_point1, img_point2, proj1, proj2): - x,y,z = DLT(proj1, proj2, img_point1, img_point2) - return np.array([x,y,z], dtype=np.float32) + +def find_checker_corners(img, board_shape, show=False): + """https://docs.opencv.org/4.x/dc/dbb/tutorial_py_calibration.html + """ + if show: + view = pg.image(img) + + criteria = (cv.TERM_CRITERIA_EPS + cv.TERM_CRITERIA_MAX_ITER, 30, 0.001) + + # prepare object points, like (0,0,0), (1,0,0), (2,0,0) ....,(6,5,0) + objp = np.zeros((board_shape[0] * board_shape[1], 3), dtype='float32') + objp[:, :2] = np.mgrid[0:board_shape[0], 0:board_shape[1]].T.reshape(-1, 2) + + gray = cv.cvtColor(img, cv.COLOR_BGR2GRAY) + + # Find the chess board corners + ret, corners = cv.findChessboardCorners(gray, board_shape, None) + if not ret: + return None, None + + # If found, add object points, image points (after refining them) + imgp = cv.cornerSubPix(gray, corners, (11,11), (-1,-1), criteria) + + if show: + plt = pg.ScatterPlotItem(x=imgp[:, 0, 1], y=imgp[:, 0, 0]) + view.view.addItem(plt) + + return objp, imgp diff --git a/parallax/main_window.py b/parallax/main_window.py index 431ec616..0cd9be8b 100644 --- a/parallax/main_window.py +++ b/parallax/main_window.py @@ -1,4 +1,4 @@ -from PyQt5.QtWidgets import QApplication, QWidget, QVBoxLayout, QHBoxLayout, QMainWindow, QAction +from PyQt5.QtWidgets import QApplication, QWidget, QHBoxLayout, QGridLayout, QMainWindow, QAction, QSplitter from PyQt5.QtCore import Qt, QTimer from PyQt5.QtGui import QIcon import pyqtgraph.console @@ -10,6 +10,8 @@ from .geometry_panel import GeometryPanel from .dialogs import AboutDialog from .stage_manager import StageManager +from .config import config +from .training_data import TrainingDataCollector from .rigid_body_transform_tool import RigidBodyTransformTool from .template_tool import TemplateTool from .accuracy_test import AccuracyTestTool @@ -21,8 +23,14 @@ def __init__(self, model): QMainWindow.__init__(self) self.model = model + # allow main window to be accessed globally + model.main_window = self + + self.data_collector = None + self.widget = MainWidget(model) self.setCentralWidget(self.widget) + self.resize(1280, 900) # menubar actions self.save_frames_action = QAction("Save Camera Frames") @@ -44,6 +52,8 @@ def __init__(self, model): self.accutest_action.triggered.connect(self.launch_accutest) self.console_action = QAction("Python Console") self.console_action.triggered.connect(self.show_console) + self.training_data_action = QAction("Collect Training Data") + self.training_data_action.triggered.connect(self.collect_training_data) self.about_action = QAction("About") self.about_action.triggered.connect(self.launch_about) @@ -65,6 +75,7 @@ def __init__(self, model): self.tools_menu.addAction(self.tt_action) self.tools_menu.addAction(self.accutest_action) self.tools_menu.addAction(self.console_action) + self.tools_menu.addAction(self.training_data_action) self.help_menu = self.menuBar().addMenu("Help") self.help_menu.addAction(self.about_action) @@ -75,7 +86,7 @@ def __init__(self, model): self.console = None self.refresh_cameras() - self.model.scan_for_usb_stages() + self.model.scan_for_stages() self.refresh_focus_controllers() def launch_stage_manager(self): @@ -102,7 +113,7 @@ def new_transform(self, name, tr): self.model.add_transform(name, tr) def screens(self): - return self.widget.lscreen, self.widget.rscreen + return self.widget.screens[:] def refresh_cameras(self): self.model.scan_for_cameras() @@ -111,7 +122,12 @@ def refresh_cameras(self): def show_console(self): if self.console is None: - self.console = pyqtgraph.console.ConsoleWidget() + self.console = pyqtgraph.console.ConsoleWidget( + historyFile=config['console_history_file'], + editor=config['console_edit_command'], + namespace={'model': self.model, 'win': self} + ) + self.console.catchNextException() self.console.show() def closeEvent(self, ev): @@ -123,37 +139,42 @@ def refresh_focus_controllers(self): for screen in self.screens(): screen.update_focus_control_menu() + def collect_training_data(self): + self.data_collector = TrainingDataCollector(self.model) + self.data_collector.start() -class MainWidget(QWidget): +class MainWidget(QSplitter): def __init__(self, model): - QWidget.__init__(self) + QSplitter.__init__(self, Qt.Vertical) self.model = model - self.screens = QWidget() - hlayout = QHBoxLayout() - self.lscreen = ScreenWidget(model=self.model) - self.rscreen = ScreenWidget(model=self.model) - hlayout.addWidget(self.lscreen) - hlayout.addWidget(self.rscreen) - self.screens.setLayout(hlayout) + self.screens_widget = QWidget() + self.screens_layout = QHBoxLayout() + self.screens_widget.setLayout(self.screens_layout) + self.screens = [] # screens are added by init config + self.addWidget(self.screens_widget) self.controls = QWidget() self.control_panel1 = ControlPanel(self.model) self.control_panel2 = ControlPanel(self.model) self.geo_panel = GeometryPanel(self.model) - hlayout = QHBoxLayout() - hlayout.addWidget(self.control_panel1) - hlayout.addWidget(self.geo_panel) - hlayout.addWidget(self.control_panel2) - self.controls.setLayout(hlayout) + self.msg_log = MessageLog() + self.controls_layout = QGridLayout() + self.controls_layout.addWidget(self.control_panel1, 0, 0) + self.controls_layout.addWidget(self.geo_panel, 0, 1) + self.controls_layout.addWidget(self.control_panel2, 0, 2) + self.controls_layout.addWidget(self.msg_log, 1, 0, 1, 3) + self.controls.setLayout(self.controls_layout) + self.addWidget(self.controls) + + self.setSizes([550, 350]) self.refresh_timer = QTimer() self.refresh_timer.timeout.connect(self.refresh) self.refresh_timer.start(125) # connections - self.msg_log = MessageLog() self.control_panel1.msg_posted.connect(self.msg_log.post) self.control_panel2.msg_posted.connect(self.msg_log.post) self.control_panel1.target_reached.connect(self.zoom_out) @@ -162,16 +183,14 @@ def __init__(self, model): self.geo_panel.cal_point_reached.connect(self.clear_selected) self.geo_panel.cal_point_reached.connect(self.zoom_out) self.model.msg_posted.connect(self.msg_log.post) - self.lscreen.selected.connect(self.model.set_lcorr) - self.lscreen.cleared.connect(self.model.clear_lcorr) - self.rscreen.selected.connect(self.model.set_rcorr) - self.rscreen.cleared.connect(self.model.clear_rcorr) - main_layout = QVBoxLayout() - main_layout.addWidget(self.screens) - main_layout.addWidget(self.controls) - main_layout.addWidget(self.msg_log) - self.setLayout(main_layout) + def add_screen(self): + screen = ScreenWidget(model=self.model) + self.screens_layout.addWidget(screen) + self.screens.append(screen) + screen.selected.connect(self.update_corr) + screen.cleared.connect(self.update_corr) + return screen def keyPressEvent(self, e): if e.key() == Qt.Key_R: @@ -190,16 +209,17 @@ def keyPressEvent(self, e): self.geo_panel.triangulate() def refresh(self): - self.lscreen.refresh() - self.rscreen.refresh() + for screen in self.screens: + screen.refresh() def clear_selected(self): - self.lscreen.clear_selected() - self.rscreen.clear_selected() + for screen in self.screens: + screen.clear_selected() + self.model.set_correspondence_points((None, None)) def zoom_out(self): - self.lscreen.zoom_out() - self.rscreen.zoom_out() + for screen in self.screens: + screen.zoom_out() def save_camera_frames(self): for i,camera in enumerate(self.model.cameras): @@ -208,4 +228,8 @@ def save_camera_frames(self): camera.save_last_image(filename) self.msg_log.post('Saved camera frame: %s' % filename) + def update_corr(self): + # send correspondence points to model + pts = [ctrl.get_selected() for ctrl in self.screens] + self.model.set_correspondence_points(pts) diff --git a/parallax/message_log.py b/parallax/message_log.py index 1e1b9e32..19ace843 100644 --- a/parallax/message_log.py +++ b/parallax/message_log.py @@ -15,6 +15,7 @@ def __init__(self, parent=None): main_layout = QVBoxLayout() main_layout.addWidget(self.message_log) + main_layout.setContentsMargins(0, 0, 0, 0) self.setLayout(main_layout) def post(self, message, **kwargs): diff --git a/parallax/mock_sim.py b/parallax/mock_sim.py new file mode 100644 index 00000000..6c4fbdac --- /dev/null +++ b/parallax/mock_sim.py @@ -0,0 +1,376 @@ +import pyqtgraph as pg +import cv2 as cv +import numpy as np +import coorx +from .calibration import CameraTransform, StereoCameraTransform +from .lib import find_checker_corners +from .config import config +from .threadrun import runInGuiThread + + +class CameraTransform(coorx.CompositeTransform): + def __init__(self): + self.view = coorx.SRT3DTransform() + self.proj = coorx.linear.PerspectiveTransform() + self.dist = coorx.nonlinear.LensDistortionTransform() + self.dist_embed = coorx.util.AxisSelectionEmbeddedTransform(axes=[0, 1], transform=self.dist, dims=(3, 3)) + self.screen = coorx.STTransform(dims=(3, 3)) + super().__init__([self.view, self.proj, self.dist_embed, self.screen]) + + def set_camera(self, center, look, fov, screen_size, up=(0, 0, 1), distortion=(0, 0, 0, 0, 0)): + center = np.asarray(center) + look = np.asarray(look) + up = np.asarray(up) + + aspect_ratio = screen_size[0] / screen_size[1] + look_dist = np.linalg.norm(look - center) + forward = look - center + forward /= np.linalg.norm(forward) + right = np.cross(forward, up) + right /= np.linalg.norm(right) + up = np.cross(right, forward) + up /= np.linalg.norm(up) + + pts1 = np.array([center, center + forward, center + right, center + up]) + pts2 = np.array([[0, 0, 0], [0, 0, -1], [1, 0, 0], [0, 1, 0]]) + self._view_mapping_err = self.view.set_mapping(pts1, pts2) + + + near = look_dist / 100 + far = look_dist * 100 + + # set perspective with aspect=1 so that distortion is performed on + # isotropic coordinates + self.proj.set_perspective(fov, 1.0, near, far) + + # correct for aspect in the screen transform instead + self.screen.set_mapping( + [[-1, -1 / aspect_ratio, -1], [1, 1 / aspect_ratio, 1]], + [[0, screen_size[1], 0], [screen_size[0], 0, 1]] + ) + + self.dist.set_coeff(distortion) + + +class GraphicsItem: + def __init__(self, views): + self._transform = coorx.SRT3DTransform() + self.items = [] + self.item_views = [] + for view in views: + self.add_view(view) + + def add_view(self, view): + self.item_views.append(GraphicsItemView(self, view)) + self.render() + + @property + def transform(self): + return self._transform + + @transform.setter + def transform(self, tr): + self._transform = tr + for view in self.item_views: + view.update_transform() + + def add_items(self, items): + self.items.extend(items) + self.render() + + def render(self): + for view in self.item_views: + view.render() + + +class GraphicsItemView: + def __init__(self, item, view): + self.item = item + self.view = view + self.rendered = False + self.full_transform = coorx.CompositeTransform([]) + self.update_transform() + self.scene = view.scene() + view.prepare_for_paint.connect(self.render_if_needed) + self.full_transform.add_change_callback(self.transform_changed) + + def update_transform(self): + self.full_transform.transforms = [self.item.transform, self.view.camera_tr] + + def transform_changed(self, event): + self.rendered = False + runInGuiThread(self.view.update) + + def render(self): + self.clear_graphics_items() + for item in self.item.items: + pts = None + if 'points' in item: + pts = self.full_transform.map(item['points']) + pts_xy = pts[:, :2] + pts_z = pts[:, 2] + + if item['type'] == 'poly': + polygon = pg.QtGui.QPolygonF([pg.QtCore.QPointF(*pt) for pt in pts_xy]) + gfx_item = pg.QtWidgets.QGraphicsPolygonItem(polygon) + elif item['type'] == 'line': + gfx_item = pg.QtWidgets.QGraphicsLineItem(*pts_xy.flatten()) + elif item['type'] == 'plot': + gfx_item = pg.PlotCurveItem(x=pts_xy[:,0], y=pts_xy[:,1]) + else: + raise TypeError(item['type']) + + if 'pen' in item: + pen = pg.mkPen(item['pen']) + gfx_item.setPen(pen) + + if 'brush' in item: + brush = pg.mkBrush(item['brush']) + gfx_item.setBrush(brush) + + gfx_item.setZValue(-pts_z.mean()) + item.setdefault('graphicsItems', {}) + item['graphicsItems'][self] = gfx_item + self.scene.addItem(gfx_item) + self.rendered = True + + def render_if_needed(self): + if not self.rendered: + self.render() + + def clear_graphics_items(self): + for item in self.item.items: + gfxitems = item.get('graphicsItems', {}) + gfxitem = gfxitems.pop(self, None) + if gfxitem is not None: + self.scene.removeItem(gfxitem) + + +class CheckerBoard(GraphicsItem): + def __init__(self, views, size, colors): + super().__init__(views) + + for i in range(size): + for j in range(size): + self.items.append({ + 'type': 'poly', + 'points': [[i, j, 0], [i+1, j, 0], [i+1, j+1, 0], [i, j+1, 0], [i, j, 0]], + 'pen': None, + 'brush': colors[(i + j) % 2], + }) + self.render() + + +class Axis(GraphicsItem): + def __init__(self, views): + super().__init__(views) + + self.items = [ + {'type': 'line', 'points': [[0, 0, 0], [1, 0, 0]], 'pen': {'color': 'r', 'width': 5}}, + {'type': 'line', 'points': [[0, 0, 0], [0, 1, 0]], 'pen': {'color': 'g', 'width': 5}}, + {'type': 'line', 'points': [[0, 0, 0], [0, 0, 1]], 'pen': {'color': 'b', 'width': 5}}, + ] + self.render() + + +class Electrode(GraphicsItem): + def __init__(self, views): + super().__init__(views) + w = 70 + l = 10e3 + self.items = [ + {'type': 'poly', 'pen': None, 'brush': 0.2, 'points': [ + [0, 0, 0], [w, 0, 2*w], [w, 0, l], [-w, 0, l], [-w, 0, 2*w], [0, 0, 0] + ]}, + ] + self.render() + + +class GraphicsView3D(pg.GraphicsView): + + prepare_for_paint = pg.QtCore.Signal() + + def __init__(self, **kwds): + self.cached_frame = None + self.camera_tr = CameraTransform() + self.press_event = None + self.camera_params = {'look': [0, 0, 0], 'pitch': 30, 'yaw': 0, 'distance': 10, 'fov': 45, 'distortion': (0, 0, 0, 0, 0)} + super().__init__(**kwds) + self.setRenderHint(pg.QtGui.QPainter.Antialiasing) + self.set_camera(look=[0, 0, 0], pitch=30, yaw=0, distance=10, fov=45) + + def set_camera(self, **kwds): + for k,v in kwds.items(): + assert k in self.camera_params + self.camera_params[k] = v + self.update_camera() + + def update_camera(self): + p = self.camera_params + look = np.asarray(p['look']) + pitch = p['pitch'] * np.pi/180 + hdist = p['distance'] * np.cos(pitch) + yaw = p['yaw'] * np.pi/180 + cam_pos = look + np.array([ + hdist * np.cos(yaw), + hdist * np.sin(yaw), + p['distance'] * np.sin(pitch) + ]) + self.camera_tr.set_camera(center=cam_pos, look=look, fov=p['fov'], screen_size=[self.width(), self.height()], distortion=p['distortion']) + + def mousePressEvent(self, ev): + self.press_event = ev + self.last_mouse_pos = ev.pos() + ev.accept() + + def mouseMoveEvent(self, ev): + if self.press_event is None: + return + dif = ev.pos() - self.last_mouse_pos + self.last_mouse_pos = ev.pos() + + self.camera_params['pitch'] = np.clip(self.camera_params['pitch'] + dif.y(), -90, 90) + self.camera_params['yaw'] -= dif.x() + self.update_camera() + + def mouseReleaseEvent(self, ev): + self.press_event = None + + def wheelEvent(self, event): + self.camera_params['distance'] *= 1.01**event.angleDelta().y() + self.update_camera() + + def resizeEvent(self, ev): + super().resizeEvent(ev) + self.update_camera() + + def paintEvent(self, ev): + self.prepare_for_paint.emit() + return super().paintEvent(ev) + + def item_changed(self, item): + self.clear_frames() + self.update() + + def get_array(self): + if self.cached_frame is None: + self.prepare_for_paint.emit() + img_arr = runInGuiThread(self.grab) + self.cached_frame = pg.imageToArray(pg.QtGui.QImage(img_arr), copy=True, transpose=False)[..., :3] + return self.cached_frame + + def update(self): + self.cached_frame = None + self.scene().update() + super().update() + + +def generate_calibration_data(view, n_images, cb_size): + app = pg.QtWidgets.QApplication.instance() + imgp = [] + objp = [] + cam = view.camera_params.copy() + for i in range(n_images*4): + pitch = np.random.uniform(45, 89) + yaw = np.random.uniform(0, 360) + distance = np.random.uniform(5, 15) + view.set_camera(pitch=pitch, yaw=yaw, distance=distance) + op,ip = find_checker_corners(view.get_array(), cb_size) + if op is None: + continue + objp.append(op) + imgp.append(ip) + if len(imgp) >= n_images: + break + app.processEvents() + view.set_camera(**cam) + return objp, imgp + + +def calibrate_camera(view, cb_size, n_images=40): + objp, imgp = generate_calibration_data(view, n_images=n_images, cb_size=cb_size) + ret, mtx, dist, rvecs, tvecs = cv.calibrateCamera(objp, imgp, (view.width(), view.height()), None, None) + return ret, mtx, dist, rvecs, tvecs + + +def undistort_image(img, mtx, dist, optimize=False, crop=False): + h, w = img.shape[:2] + if optimize: + new_camera_mtx, roi = cv.getOptimalNewCameraMatrix(mtx, dist, (w,h), 1, (w,h)) + else: + assert crop is False + new_camera_mtx = mtx + + # undistort + mapx, mapy = cv.initUndistortRectifyMap(mtx, dist, None, new_camera_mtx, (w,h), 5) + dst = cv.remap(img, mapx, mapy, cv.INTER_LINEAR) + + # crop the image + if crop: + x, y, w, h = roi + dst = dst[y:y+h, x:x+w] + return dst + + +class MockSim(pg.QtCore.QObject): + + _instance = None + + stage_moved = pg.QtCore.Signal(object) + + @classmethod + def instance(cls): + if cls._instance is None: + cls._instance = MockSim() + return cls._instance + + def __init__(self): + pg.QtCore.QObject.__init__(self) + self.cameras = {} + self.stages = {} + + self.items = [] + s = 1e3 + if config['mock_sim']['show_checkers']: + cb_size = 8 + checkers = CheckerBoard(views=[], size=cb_size, colors=[0.4, 0.6]) + checkers.transform.set_params(offset=[-s*cb_size/2, -s*cb_size/2, -4000], scale=[s, s, s]) + self.items.append(checkers) + + if config['mock_sim']['show_axes']: + axis = Axis(views=[]) + axis.transform.set_params(scale=[s, s, s]) + self.items.append(axis) + axis = Axis(views=[]) + axis.transform.set_params(scale=[s, s, s], offset=[0, 0, -4000]) + self.items.append(axis) + axis = Axis(views=[]) + axis.transform.set_params(scale=[s, s, s], offset=[2000, 0, 0]) + self.items.append(axis) + axis = Axis(views=[]) + axis.transform.set_params(scale=[s, s, s], offset=[0, 2000, 0]) + self.items.append(axis) + + def add_camera(self, cam): + view = GraphicsView3D(background=(128, 128, 128)) + view.resize(*cam.sensor_size) + view.set_camera(**cam.camera_params) + view.scene().changed.connect(self.clear_frames) + self.cameras[cam] = {'view': view} + + for item in self.items: + item.add_view(view) + + def clear_frames(self): + for v in self.cameras.values(): + v['frame'] = None + + def get_camera_frame(self, cam): + view = self.cameras[cam]['view'] + return view.get_array() + + def add_stage(self, stage): + views = [c['view'] for c in self.cameras.values()] + item = Electrode(views=views) + item.transform = stage.transform + self.stages[stage] = {'item': item} diff --git a/parallax/model.py b/parallax/model.py index a0b90b68..2ab7f26a 100755 --- a/parallax/model.py +++ b/parallax/model.py @@ -1,26 +1,32 @@ from PyQt5.QtCore import pyqtSignal from PyQt5.QtCore import QObject, pyqtSignal, QThread - +import os, re, time import numpy as np +import coorx import serial.tools.list_ports - from mis_focus_controller import FocusController -from newscale.interfaces import NewScaleSerial from .camera import list_cameras, close_cameras -from .stage import Stage +from .stage import list_stages, close_stages +from .config import config +from .calibration import Calibration from .accuracy_test import AccuracyTestWorker class Model(QObject): msg_posted = pyqtSignal(str) + calibrations_changed = pyqtSignal() + corr_pts_changed = pyqtSignal() + + instance = None def __init__(self): QObject.__init__(self) + Model.instance = self self.cameras = [] self.focos = [] - self.init_stages() + self.stages = [] self.calibration = None self.calibrations = {} @@ -28,7 +34,7 @@ def __init__(self): self.cal_in_progress = False self.accutest_in_progress = False - self.lcorr, self.rcorr = False, False + self.lcorr, self.rcorr = None, None self.obj_point_last = None self.transforms = {} @@ -37,39 +43,74 @@ def __init__(self): def ncameras(self): return len(self.cameras) + def get_camera(self, camera_name): + for cam in self.cameras: + if cam.name() == camera_name: + return cam + raise NameError(f"No camera named {camera_name}") + def set_last_object_point(self, obj_point): self.obj_point_last = obj_point + def get_image_point(self): + if None in (self.lcorr, self.rcorr): + return None + concat = np.hstack([self.lcorr, self.rcorr]) + return coorx.Point(concat, f'{self.lcorr.system.name}+{self.rcorr.system.name}') + def add_calibration(self, cal): self.calibrations[cal.name] = cal + self.calibrations_changed.emit() + + def list_calibrations(self): + """List all known calibrations, including those loaded in memory and + those stored in a standard location / filename. + """ + cal_path = config['calibration_path'] + cal_files = os.listdir(cal_path) + calibrations = [] + for cf in sorted(cal_files, reverse=True): + m = re.match(Calibration.file_regex, cf) + if m is None: + continue + mg = m.groups() + ts = time.strptime(mg[2], Calibration.date_format) + calibrations.append({'file': os.path.join(cal_path, cf), 'from_cs': mg[0], 'to_cs': mg[1], 'timestamp': ts}) + for cal in self.calibrations.values(): + calibrations.append({'calibration': cal, 'from_cs': cal.from_cs, 'to_cs': cal.to_cs, 'timestamp': cal.timestamp}) + + return calibrations + + def get_calibration(self, stage): + """Return the most recent calibration known for this stage + """ + cals = self.list_calibrations() + cals = [cal for cal in cals if cal['to_cs'] == stage.get_name()] + cals = sorted(cals, key=lambda cal: cal['timestamp']) + + if len(cals) == 0: + return None + else: + cal_spec = cals[-1] + if 'calibration' in cal_spec: + cal = cal_spec['calibration'] + else: + cal = Calibration.load(cal_spec['file']) + return cal def set_calibration(self, calibration): self.calibration = calibration - def set_lcorr(self, xc, yc): - self.lcorr = [xc, yc] - - def clear_lcorr(self): - self.lcorr = False - - def set_rcorr(self, xc, yc): - self.rcorr = [xc, yc] - - def clear_rcorr(self): - self.rcorr = False - - def init_stages(self): - self.stages = {} + def set_correspondence_points(self, pts): + self.lcorr = pts[0] + self.rcorr = pts[1] + self.corr_pts_changed.emit() def scan_for_cameras(self): self.cameras = list_cameras() - def scan_for_usb_stages(self): - instances = NewScaleSerial.get_instances() - self.init_stages() - for instance in instances: - stage = Stage(serial=instance) - self.add_stage(stage) + def scan_for_stages(self): + self.stages = list_stages() def scan_for_focus_controllers(self): self.focos = [] @@ -86,13 +127,10 @@ def add_stage(self, stage): def clean(self): close_cameras() - self.clean_stages() - - def clean_stages(self): - pass + close_stages() def halt_all_stages(self): - for stage in self.stages.values(): + for stage in self.stages: stage.halt() self.msg_posted.emit('Halting all stages.') @@ -102,6 +140,18 @@ def add_transform(self, name, transform): def get_transform(self, name): return self.transforms[name] + def set_lcorr(self, xc, yc): + self.lcorr = [xc, yc] + + def clear_lcorr(self): + self.lcorr = False + + def set_rcorr(self, xc, yc): + self.rcorr = [xc, yc] + + def clear_rcorr(self): + self.rcorr = False + def handle_accutest_point_reached(self, i, npoints): self.msg_posted.emit('Accuracy test point %d (of %d) reached.' % (i+1,npoints)) self.clear_lcorr() @@ -109,11 +159,10 @@ def handle_accutest_point_reached(self, i, npoints): self.msg_posted.emit('Highlight correspondence points and press C to continue') def register_corr_points_accutest(self): - lcorr, rcorr = self.lcorr, self.rcorr - if (lcorr and rcorr): - self.accutest_worker.register_corr_points(lcorr, rcorr) - self.msg_posted.emit('Correspondence points registered: (%d,%d) and (%d,%d)' % \ - (lcorr[0],lcorr[1], rcorr[0],rcorr[1])) + corr_pt = self.get_image_point() + if corr_pt is not None: + self.accutest_worker.register_corr_points(corr_pt) + self.msg_posted.emit('Correspondence points registered.') self.accutest_worker.carry_on() else: self.msg_posted.emit('Highlight correspondence points and press C to continue') diff --git a/parallax/screen_widget.py b/parallax/screen_widget.py index a6b7cfc9..4d53904f 100755 --- a/parallax/screen_widget.py +++ b/parallax/screen_widget.py @@ -4,9 +4,8 @@ from PyQt5.QtCore import pyqtSignal, Qt from PyQt5 import QtCore import pyqtgraph as pg +import coorx import inspect -import importlib - from . import filters from . import detectors @@ -58,6 +57,8 @@ def __init__(self, filename=None, model=None, parent=None): self.detector_menu = self.parallax_menu.addMenu("Detectors") self.view_box.menu.insertMenu(self.view_box.menu.actions()[0], self.parallax_menu) + self.update_camera_menu() + self.update_focus_control_menu() self.update_filter_menu() self.update_detector_menu() @@ -152,10 +153,15 @@ def set_detector(self, detector): def get_selected(self): if self.click_target.isVisible(): pos = self.click_target.pos() - return pos.x(), pos.y() + return coorx.Point([pos.x(), pos.y()], self.camera.name()) else: return None + def set_selected(self, pos): + self.click_target.setPos(pos) + self.click_target.setVisible(True) + self.selected.emit(*self.get_selected()) + def wheelEvent(self, e): forward = bool(e.angleDelta().y() > 0) control = bool(e.modifiers() & Qt.ControlModifier) diff --git a/parallax/stage.py b/parallax/stage.py index 1c925967..fb9d80e5 100755 --- a/parallax/stage.py +++ b/parallax/stage.py @@ -1,11 +1,57 @@ +import time, queue, threading +import numpy as np +from PyQt5 import QtCore +import coorx +from coorx import Point from newscale.multistage import USBXYZStage, PoEXYZStage -from newscale.interfaces import USBInterface +from newscale.interfaces import USBInterface, NewScaleSerial +from .mock_sim import MockSim -class Stage: +def list_stages(): + stages = [] + stages.extend(NewScaleStage.scan_for_stages()) + if len(stages) == 0: + tr1 = coorx.AffineTransform(dims=(3, 3)) + tr1.translate([0, 0, 500]) + tr1.rotate(60, (1, 0, 0)) + tr1.rotate(90, (0, 0, 1)) + tr2 = tr1.copy() + tr2.rotate(50, (0, 0, 1)) - def __init__(self, ip=None, serial=None): + stages.extend([ + MockStage(transform=tr1), + MockStage(transform=tr2), + ]) + return stages + + +def close_stages(): + NewScaleStage.close_stages() + MockStage.close_stages() + + +class NewScaleStage: + + stages = {} + + @classmethod + def scan_for_stages(cls): + instances = NewScaleSerial.get_instances() + stages = [] + for serial in instances: + if serial not in cls.stages: + cls.stages[serial] = NewScaleStage(serial=serial) + stages.append(cls.stages[serial]) + return stages + + @classmethod + def close_stages(cls): + for stage in cls.stages.values(): + stage.close() + def __init__(self, ip=None, serial=None): + super().__init__() if ip is not None: self.ip = ip self.name = ip @@ -17,6 +63,9 @@ def __init__(self, ip=None, serial=None): self.initialize() + def close(self): + pass + def calibrate_frequency(self): self.device.calibrate_all() @@ -39,7 +88,9 @@ def get_position(self, relative=False): x -= self.origin[0] y -= self.origin[1] z -= self.origin[2] - return x,y,z + return Point([x,y,z], self.get_name()+'_rel') + else: + return Point([x,y,z], self.get_name()) def move_to_target_1d(self, axis, position, relative=False): if axis == 'x': @@ -85,3 +136,171 @@ def get_accel(self): def halt(self): pass + + +class MockStage(QtCore.QObject): + + n_mock_stages = 0 + + position_changed = QtCore.pyqtSignal(object, object) # self, pos + + def __init__(self, transform): + super().__init__() + self.base_transform = transform + self.transform = coorx.AffineTransform(dims=(3, 3)) + self.speed = 8000 # um/s + self.accel = 5000 # um/s^2 + self.pos = np.array([0, 0, 0]) + self.name = f"mock_stage_{MockStage.n_mock_stages}" + self.transform.set_systems(self.name, "global") + MockStage.n_mock_stages += 1 + + self.update_pos(self.pos) + + self.move_callbacks = [] + + self.move_queue = queue.Queue() + self.move_thread = threading.Thread(target=self.thread_loop, daemon=True) + self.move_thread.start() + + self.sim = MockSim.instance() + self.sim.add_stage(self) + + def get_origin(self): + return [0, 0, 0] + + @classmethod + def close_stages(self): + pass + + def get_name(self): + return self.name + + def get_speed(self): + return self.speed + + def set_speed(self, speed): + self.speed = speed + + def get_accel(self): + return self.accel + + def get_position(self): + return Point(self.pos.copy(), self.get_name()) + + def move_to_target_3d(self, x, y, z, relative=False, safe=False, speed=None, block=True): + if relative: + xo,yo,zo = self.get_origin() + x += xo + y += yo + z += zo + + speed = speed or self.speed + move_cmd = MoveFuture(self, pos=np.array([x, y, z]), speed=speed, accel=self.accel) + self.move_queue.put(move_cmd) + if block: + move_cmd.wait() + return move_cmd + + def move_distance_1d(self, axis, distance): + ax_ind = 'xyz'.index(axis) + pos = self.get_position().coordinates + pos[ax_ind] += distance + return self.move_to_target_3d(*pos) + + def halt(self): + self.move_queue.put('halt') + + def update_pos(self, pos): + # stage has moved; update and emit + self.pos = pos + + # update transform used by mock graphics + tr = coorx.AffineTransform(dims=(3, 3)) + tr.translate(pos) + tr2 = self.base_transform * tr + self.transform.set_params(matrix=tr2.matrix, offset=tr2.offset) + + self.position_changed.emit(self, pos) + + def thread_loop(self): + current_move = None + while True: + next_move = None + while not self.move_queue.empty(): + next_move = self.move_queue.get() + + if next_move is not None: + if current_move is not None: + if next_move == 'halt': + message = 'interrupted by halt request' + else: + message = "interrupted by another move request" + current_move.finish(interrupted=True, message=message) + if next_move == 'halt': + next_move = None + current_move = next_move + last_update = time.perf_counter() + + if current_move is not None: + now = time.perf_counter() + dt = now - last_update + last_update = now + + pos = self.get_position().coordinates + target = current_move.pos + dx = target - pos + dist_to_go = np.linalg.norm(dx) + max_dist_per_step = current_move.speed * dt + if dist_to_go > max_dist_per_step: + # take a step + direction = dx / dist_to_go + dist_this_step = min(dist_to_go, max_dist_per_step) + step = direction * dist_this_step + self.update_pos(pos + step) + else: + self.update_pos(target.copy()) + current_move.finish(interrupted=False) + current_move = None + + for cb in self.move_callbacks[:]: + cb(self) + + time.sleep(100e-3) + + def add_move_callback(self, cb): + self.move_callbacks.append(cb) + + def orientation(self): + p1 = self.transform.map(Point([0, 0, 0], self.name)) + p2 = self.transform.map(Point([0, 0, 1], self.name)) + axis = p1 - p2 + r = np.linalg.norm(axis[:2]) + phi = np.arctan2(r, axis[2]) * 180 / np.pi + theta = np.arctan2(axis[1], axis[0]) * 180 / np.pi + return theta, phi + + def get_tip_position(self): + return self.transform.map(Point([0, 0, 0], self.name)) + + +class MoveFuture: + def __init__(self, stage, pos, speed, accel): + self.stage = stage + self.pos = pos + self.speed = speed + self.accel = accel + self.finished = threading.Event() + self.interrupted = False + + def finish(self, interrupted, message=None): + self.interrupted = interrupted + self.message = message + self.finished.set() + + @property + def succeeded(self): + return self.finished and not self.interrupted + + def wait(self, timeout=None): + self.finished.wait(timeout=timeout) diff --git a/parallax/stage_dropdown.py b/parallax/stage_dropdown.py index f4c95d56..5be3bba9 100644 --- a/parallax/stage_dropdown.py +++ b/parallax/stage_dropdown.py @@ -7,6 +7,7 @@ class StageDropdown(QComboBox): def __init__(self, model): QComboBox.__init__(self) self.model = model + self.stages = [] self.selected = False self.setFocusPolicy(Qt.NoFocus) @@ -18,6 +19,9 @@ def set_selected(self): def is_selected(self): return self.selected + def current_stage(self): + return self.stages[self.currentIndex()] + def showPopup(self): self.populate() QComboBox.showPopup(self) @@ -27,7 +31,6 @@ def get_current_stage(self): def populate(self): self.clear() - for sn in self.model.stages.keys(): - self.addItem(sn) - - + self.stages = self.model.stages[:] + for stage in self.stages: + self.addItem(stage.get_name()) diff --git a/parallax/stage_manager.py b/parallax/stage_manager.py index 3de771e7..23c60217 100755 --- a/parallax/stage_manager.py +++ b/parallax/stage_manager.py @@ -7,7 +7,7 @@ from serial.tools.list_ports import comports as list_comports from serial.serialutil import SerialException -from .stage import Stage +from .stage import NewScaleStage from .helper import PORT_NEWSCALE @@ -69,7 +69,7 @@ def run(self): else: print('ip = ', ip) s.close() - self.stages.append((ip, Stage(ip=ip))) + self.stages.append((ip, NewScaleStage(ip=ip))) self.progress_made.emit(i) self.finished.emit() diff --git a/parallax/threadrun.py b/parallax/threadrun.py new file mode 100644 index 00000000..4dfb0ae7 --- /dev/null +++ b/parallax/threadrun.py @@ -0,0 +1,340 @@ +"""Borrowed from ACQ4 +""" +import time, sys, threading, traceback, functools + +from PyQt5 import QtCore, QtWidgets + + +def runInThread(thread, func, *args, **kwds): + """Run a function in another thread and return the result. + + The remote thread must be running a Qt event loop. + """ + return ThreadCallFuture(thread, func, *args, **kwds)() + + +def runInGuiThread(func, *args, **kwds): + """Run a function the main GUI thread and return the result. + """ + isGuiThread = QtCore.QThread.currentThread() == QtCore.QCoreApplication.instance().thread() + if isGuiThread: + return func(*args, **kwds) + return ThreadCallFuture(None, func, *args, **kwds)() + + +class Future(QtCore.QObject): + """Used to track the progress of an asynchronous task. + + The simplest subclasses reimplement percentDone() and call _taskDone() when finished. + """ + sigFinished = QtCore.pyqtSignal(object) # self + sigStateChanged = QtCore.pyqtSignal(object, object) # self, state + + class StopRequested(Exception): + """Raised by _checkStop if stop() has been invoked. + """ + + class Timeout(Exception): + """Raised by wait() if the timeout period elapses. + """ + + def __init__(self): + QtCore.QObject.__init__(self) + + self.startTime = time.perf_counter() + + self._isDone = False + self._wasInterrupted = False + self._errorMessage = None + self._excInfo = None + self._stopRequested = False + self._state = 'starting' + self._errorMonitorThread = None + self.finishedEvent = threading.Event() + + def currentState(self): + """Return the current state of this future. + + The state can be any string used to indicate the progress this future is making in its task. + """ + return self._state + + def setState(self, state): + """Set the current state of this future. + + The state can be any string used to indicate the progress this future is making in its task. + """ + if state == self._state: + return + self._state = state + self.sigStateChanged.emit(self, state) + + def percentDone(self): + """Return the percent of the task that has completed. + + Must be reimplemented in subclasses. + """ + raise NotImplementedError("method must be reimplmented in subclass") + + def stop(self, reason="task stop requested"): + """Stop the task (nicely). + + This method may return another future if stopping the task is expected to + take time. + + Subclasses may extend this method and/or use _checkStop to determine whether + stop() has been called. + """ + if self.isDone(): + return + + if reason is not None: + self._errorMessage = reason + self._stopRequested = True + + def _taskDone(self, interrupted=False, error=None, state=None, excInfo=None): + """Called by subclasses when the task is done (regardless of the reason) + """ + if self._isDone: + raise Exception("_taskDone has already been called.") + self._isDone = True + if error is not None: + # error message may have been set earlier + self._errorMessage = error + self._excInfo = excInfo + self._wasInterrupted = interrupted + if interrupted: + self.setState(state or 'interrupted: %s' % error) + else: + self.setState(state or 'complete') + self.finishedEvent.set() + self.sigFinished.emit(self) + + def wasInterrupted(self): + """Return True if the task was interrupted before completing (due to an error or a stop request). + """ + return self._wasInterrupted + + def isDone(self): + """Return True if the task has completed successfully or was interrupted. + """ + return self._isDone + + def errorMessage(self): + """Return a string description of the reason for a task failure, + or None if there was no failure (or if the reason is unknown). + """ + return self._errorMessage + + def wait(self, timeout=None, updates=False, pollInterval=0.1): + """Block until the task has completed, has been interrupted, or the + specified timeout has elapsed. + + If *updates* is True, process Qt events while waiting. + + If a timeout is specified and the task takes too long, then raise Future.Timeout. + If the task ends incomplete for another reason, then raise RuntimeError. + """ + start = time.perf_counter() + while True: + if (timeout is not None) and (time.perf_counter() > start + timeout): + raise self.Timeout("Timeout waiting for task to complete.") + + if self.isDone(): + break + + if updates is True: + Qt.QTest.qWait(min(1, int(pollInterval * 1000))) + else: + self._wait(pollInterval) + + if self.wasInterrupted(): + err = self.errorMessage() + if err is None: + # This would be a fantastic place to "raise from self._excInfo[1]" once we move to py3 + raise RuntimeError(f"Task {self} did not complete (no error message).") + else: + raise RuntimeError(f"Task {self} did not complete: {err}") + + def _wait(self, duration): + """Default sleep implementation used by wait(); may be overridden to return early. + """ + self.finishedEvent.wait(timeout=duration) + + def _checkStop(self, delay=0): + """Raise self.StopRequested if self.stop() has been called. + + This may be used by subclasses to periodically check for stop requests. + + The optional *delay* argument causes this method to sleep while periodically + checking for a stop request. + """ + if delay == 0 and self._stopRequested: + raise self.StopRequested() + + stop = time.perf_counter() + delay + while True: + now = time.perf_counter() + if now > stop: + return + + time.sleep(max(0, min(0.1, stop-now))) + if self._stopRequested: + raise self.StopRequested() + + def sleep(self, duration, interval=0.2): + """Sleep for the specified duration (in seconds) while checking for stop requests. + """ + start = time.time() + while time.time() < start + duration: + self._checkStop() + time.sleep(interval) + + def waitFor(self, futures, timeout=20.0): + """Wait for multiple futures to complete while also checking for stop requests on self. + """ + if not isinstance(futures, (list, tuple)): + futures = [futures] + if len(futures) == 0: + return + start = time.time() + while True: + self._checkStop() + allDone = True + for fut in futures[:]: + try: + fut.wait(0.1) + futures.remove(fut) + except fut.Timeout: + allDone = False + break + if allDone: + break + if timeout is not None and time.time() - start > timeout: + raise futures[0].Timeout("Timed out waiting for %r" % futures) + + def raiseErrors(self, message, pollInterval=1.0): + """Monitor this future for errors and raise if any occur. + + This allows the caller to discard a future, but still expect errors to be delivered to the user. Note + that errors are raised from a background thread. + + Parameters + ---------- + message : str + Exception message to raise. May include "{stack}" to insert the stack trace of the caller, and "{error}" + to insert the original formatted exception. + pollInterval : float | None + Interval in seconds to poll for errors. This is only used with Futures that require a poller; + Futures that immediately report errors when they occur will not use a poller. + """ + if self._errorMonitorThread is not None: + return + originalFrame = sys._getframe().f_back + monitorFn = functools.partial(self._monitorErrors, message=message, pollInterval=pollInterval, originalFrame=originalFrame) + self._errorMonitorThread = threading.Thread(target=monitorFn, daemon=True) + self._errorMonitorThread.start() + + def _monitorErrors(self, message, pollInterval, originalFrame): + try: + self.wait(pollInterval=pollInterval) + except Exception as exc: + if '{stack}' in message: + stack = ''.join(traceback.format_stack(originalFrame)) + else: + stack = None + + try: + formattedMsg = message.format(stack=stack, error=traceback.format_exception_only(type(exc), exc)) + except Exception as exc2: + formattedMsg = message + f" [additional error formatting error message: {exc2}]" + raise RuntimeError(formattedMsg) from exc + + +class ThreadCallFuture(Future): + sigRequestCall = QtCore.pyqtSignal() + def __init__(self, thread, func, *args, **kwds): + Future.__init__(self) + self.func = func + self.args = args + self.kwds = kwds + self.exc = None + + if thread is None: + thread = QtWidgets.QApplication.instance().thread() + self.moveToThread(thread) + self.sigRequestCall.connect(self._callRequested) + self.sigRequestCall.emit() + + def _callRequested(self): + try: + self.ret = self.func(*self.args, **self.kwds) + self._taskDone() + except Exception as exc: + self.exc = exc + err = ''.join(traceback.format_exception(*sys.exc_info())) + self._taskDone(interrupted=True, error=err) + + def __call__(self): + self.wait() + if self.exc is not None: + raise self.exc + else: + return self.ret + + + + + +class _FuturePollThread(threading.Thread): + """Thread used to poll the state of a future. + + Used when a Future subclass does not automatically call _taskDone, but instead requires + a periodic check. May + """ + def __init__(self, future, pollInterval, originalFrame): + threading.Thread.__init__(self, daemon=True) + self.future = future + self.pollInterval = pollInterval + self._stop = False + + def run(self): + while not self._stop: + if self.future.isDone(): + break + if self.future._raiseErrors: + raise + time.sleep(self.pollInterval) + + def stop(self): + self._stop = True + self.join() + + +class MultiFuture(Future): + """Future tracking progress of multiple sub-futures. + """ + def __init__(self, futures): + self.futures = futures + Future.__init__(self) + + def stop(self, reason="task stop requested"): + for f in self.futures: + f.stop(reason=reason) + return Future.stop(self, reason) + + def percentDone(self): + return min([f.percentDone() for f in self.futures]) + + def wasInterrupted(self): + return any([f.wasInterrupted() for f in self.futures]) + + def isDone(self): + return all([f.isDone() for f in self.futures]) + + def errorMessage(self): + return "; ".join([f.errorMessage() or '' for f in self.futures]) + + def currentState(self): + return "; ".join([f.currentState() or '' for f in self.futures]) + diff --git a/parallax/training_data.py b/parallax/training_data.py new file mode 100644 index 00000000..6d5a1c22 --- /dev/null +++ b/parallax/training_data.py @@ -0,0 +1,78 @@ +import threading, pickle, os +import numpy as np +from PyQt5 import QtWidgets, QtCore +from .dialogs import TrainingDataDialog + + +class TrainingDataCollector(QtCore.QObject): + def __init__(self, model): + QtCore.QObject.__init__(self) + self.model = model + + def start(self): + dlg = TrainingDataDialog(self.model) + dlg.exec_() + if dlg.result() != dlg.Accepted: + return + + self.stage = dlg.get_stage() + self.img_count = dlg.get_img_count() + self.extent = dlg.get_extent() + self.path = QtWidgets.QFileDialog.getExistingDirectory(parent=None, caption="Select Storage Directory") + if self.path == '': + return + + self.start_pos = self.stage.get_position() + self.stage_cal = self.model.get_calibration(self.stage) + + self.thread = threading.Thread(target=self.thread_run, daemon=True) + self.thread.start() + + def thread_run(self): + meta_file = os.path.join(self.path, 'meta.pkl') + if os.path.exists(meta_file): + # todo: just append + raise Exception("Already data in this folder!") + trials = [] + meta = { + 'calibration': self.stage_cal, + 'stage': self.stage.get_name(), + 'trials': trials, + } + + # move electrode out of fov for background images + pos = self.start_pos.coordinates.copy() + pos[2] += 10000 + self.stage.move_to_target_3d(*pos, block=True) + imgs = self.save_images('background') + meta['background'] = imgs + + for i in range(self.img_count): + + # first image in random location + rnd = np.random.uniform(-self.extent/2, self.extent/2, size=3) + pos1 = self.start_pos.coordinates + rnd + self.stage.move_to_target_3d(*pos1, block=True) + images1 = self.save_images(f'{i:04d}-a') + + # take a second image slightly shifted + pos2 = pos1.copy() + pos2[2] += 10 + self.stage.move_to_target_3d(*pos2, block=True) + images2 = self.save_images(f'{i:04d}-b') + + trials.append([ + {'pos': pos1, 'images': images1}, + {'pos': pos2, 'images': images2}, + ]) + + with open(meta_file, 'wb') as fh: + pickle.dump(meta, fh) + + def save_images(self, name): + images = [] + for camera in self.model.cameras: + filename = f'{name}-{camera.name()}.png' + camera.save_last_image(os.path.join(self.path, filename)) + images.append({'camera': camera.name(), 'image': filename}) + return images diff --git a/run-parallax.py b/run-parallax.py index 8337b39b..3b9efeeb 100755 --- a/run-parallax.py +++ b/run-parallax.py @@ -1,8 +1,9 @@ #!/usr/bin/env python +import atexit from PyQt5.QtWidgets import QApplication from parallax.model import Model from parallax.main_window import MainWindow -import atexit +import parallax.config # set up logging to file import logging @@ -16,10 +17,15 @@ app = QApplication([]) +args = parallax.config.parse_cli_args() +parallax.config.init_config(args) + model = Model() atexit.register(model.clean) main_window = MainWindow(model) main_window.show() +parallax.config.post_init_config(model, main_window) + app.exec() diff --git a/tools/annotate_training_data.py b/tools/annotate_training_data.py new file mode 100644 index 00000000..f748a93f --- /dev/null +++ b/tools/annotate_training_data.py @@ -0,0 +1,128 @@ +import pyqtgraph as pg +import json + + +class MainWindow(pg.GraphicsView): + def __init__(self, meta_file, img_files): + super().__init__() + self.img_files = img_files + + self.view = pg.ViewBox() + self.view.invertY() + self.setCentralItem(self.view) + + self.img_item = pg.QtWidgets.QGraphicsPixmapItem() + self.view.addItem(self.img_item) + + self.line_item = pg.QtWidgets.QGraphicsLineItem() + self.line_item.setPen(pg.mkPen('r')) + self.circle_item = pg.QtWidgets.QGraphicsEllipseItem() + self.circle_item.setPen(pg.mkPen('r')) + self.view.addItem(self.line_item) + self.view.addItem(self.circle_item) + + self.next_click = 0 + self.attached_pt = None + self.loaded_file = None + + self.meta_file = meta_file + if os.path.exists(meta_file): + self.meta = json.load(open(meta_file, 'r')) + else: + self.meta = {} + + self.load_image(0) + + def keyPressEvent(self, ev): + if ev.key() == pg.QtCore.Qt.Key_Left: + self.load_image(self.current_index - 1) + elif ev.key() == pg.QtCore.Qt.Key_Right: + self.load_image(self.current_index + 1) + else: + print(ev.key()) + + def mousePressEvent(self, ev): + # print('press', ev) + if ev.button() == pg.QtCore.Qt.LeftButton: + self.attached_pt = self.next_click + self.update_pos(ev.pos()) + ev.accept() + return + # return super().mousePressEvent(ev) + + def mouseReleaseEvent(self, ev): + # print('release', ev) + self.attached_pt = None + self.next_click = (self.next_click + 1) % 2 + + def mouseMoveEvent(self, ev): + # print('move', ev) + self.update_pos(ev.pos()) + ev.accept() + + def update_pos(self, pos): + pos = self.view.mapDeviceToView(pos) + if self.attached_pt == 0: + self.set_pts(pos, None) + elif self.attached_pt == 1: + self.set_pts(None, pos) + else: + return + self.update_meta() + + def set_pts(self, pt1, pt2): + line = self.line_item.line() + if pt1 is not None: + line.setP1(pt1) + self.circle_item.setRect(pt1.x()-10, pt1.y()-10, 20, 20) + self.circle_item.setVisible(True) + if pt2 is not None: + line.setP2(pt2) + self.line_item.setVisible(True) + self.line_item.setLine(line) + + def hide_line(self): + self.line_item.setVisible(False) + self.circle_item.setVisible(False) + + def update_meta(self): + line = self.line_item.line() + self.meta[self.loaded_file] = { + 'pt1': (line.x1(), line.y1()), + 'pt2': (line.x2(), line.y2()), + } + json.dump(self.meta, open(self.meta_file, 'w')) + + def load_image(self, index): + filename = self.img_files[index] + pxm = pg.QtGui.QPixmap() + pxm.load(filename) + self.img_item.setPixmap(pxm) + self.img_item.pxm = pxm + self.view.autoRange(padding=0) + self.current_index = index + self.setWindowTitle(filename) + self.loaded_file = filename + + meta = self.meta.get(filename, {}) + pt1 = meta.get('pt1', None) + pt2 = meta.get('pt2', None) + if None in (pt1, pt2): + self.hide_line() + else: + self.set_pts(pg.QtCore.QPointF(*pt1), pg.QtCore.QPointF(*pt2)) + + +if __name__ == '__main__': + import os, sys + + app = pg.mkQApp() + + meta_file = sys.argv[1] + img_files = sys.argv[2:] + win = MainWindow(meta_file, img_files) + win.resize(1000, 800) + win.show() + + if sys.flags.interactive == 0: + app.exec_() \ No newline at end of file diff --git a/tools/mock_camera.py b/tools/mock_camera.py new file mode 100644 index 00000000..9f18bf47 --- /dev/null +++ b/tools/mock_camera.py @@ -0,0 +1,94 @@ +import pyqtgraph as pg +import coorx +from parallax.mock_sim import GraphicsView3D, CheckerBoard, Axis, Electrode, calibrate_camera, undistort_image +from parallax.lib import find_checker_corners + +class StereoView(pg.QtWidgets.QWidget): + def __init__(self, parent=None, background=(128, 128, 128)): + pg.QtWidgets.QWidget.__init__(self, parent) + self.layout = pg.QtWidgets.QHBoxLayout() + self.setLayout(self.layout) + self.views = [GraphicsView3D(parent=self), GraphicsView3D(parent=self)] + for v in self.views: + self.layout.addWidget(v) + v.setBackground(pg.mkColor(background)) + + def set_camera(self, cam, **kwds): + self.views[cam].set_camera(**kwds) + + + +if __name__ == '__main__': + # pg.dbg() + + app = pg.mkQApp() + win = StereoView() + win.resize(1600, 600) + win.show() + + camera_params = dict( + pitch=30, + distance=15, + distortion=(-0.1, 0.01, -0.001, 0, 0), + # distortion=(2.49765866e-02, -1.10638222e+01, -1.22811774e-04, 4.89346001e-03, -3.28053580e-01), + ) + win.set_camera(0, yaw=-5, **camera_params) + win.set_camera(1, yaw=5, **camera_params) + + cb_size = 8 + checkers = CheckerBoard(views=win.views, size=cb_size, colors=[0.1, 0.9]) + checkers.transform.set_params(offset=[-cb_size/2, -cb_size/2, 0]) + axis = Axis(views=win.views) + + s = 0.1 + # electrodes = [] + # for i in range(4): + # e = Electrode(win.views) + # electrodes.append(e) + + # e.transform = coorx.AffineTransform(dims=(3, 3)) + # e.transform.scale([s, s, s]) + # e.transform.rotate(60, [1, 0, 0]) + # e.transform.translate([0, 1, 1]) + # e.transform.rotate(15*i, [0, 0, 1]) + + + # tr = coorx.AffineTransform(dims=(3, 3)) + # tr.translate([-2.5, -2.5, 0]) + # tr.scale(0.5) + # tr.rotate(30, [1, 0, 0]) + # tr.rotate(45, [0, 0, 1]) + + def test(n_images=10): + view = win.views[0] + ret, mtx, dist, rvecs, tvecs = calibrate_camera(view, n_images=n_images, cb_size=(cb_size-1, cb_size-1)) + print(f"Distortion coefficients: {dist}") + print(f"Intrinsic matrix: {mtx}") + pg.image(undistort_image(view.get_array(), mtx, dist).transpose(1, 0, 2)) + return mtx, dist + + + def test2(n_images=10): + """Can we invert opencv's undistortion? + """ + ret, mtx, dist, rvecs, tvecs = calibrate_camera(win.views[0], n_images=n_images, cb_size=(cb_size-1, cb_size-1)) + print(mtx) + print(dist) + img = win.views[0].get_array() + uimg = undistort_image(img, mtx, dist) + v1 = pg.image(img.transpose(1, 0, 2)) + v2 = pg.image(uimg.transpose(1, 0, 2)) + tr = coorx.AffineTransform(matrix=mtx[:2, :2], offset=mtx[:2, 2]) + ltr = coorx.nonlinear.LensDistortionTransform(dist[0]) + ttr = coorx.CompositeTransform([tr.inverse, ltr, tr]) + + objp, imgp = find_checker_corners(uimg, board_shape=(cb_size-1, cb_size-1)) + undistorted_pts = imgp[:, 0, :] + + distorted_pts = ttr.map(undistorted_pts) + + s1 = pg.ScatterPlotItem(x=undistorted_pts[:,0], y=undistorted_pts[:,1], brush='r', pen=None) + v2.view.addItem(s1) + + s2 = pg.ScatterPlotItem(x=distorted_pts[:,0], y=distorted_pts[:,1], brush='r', pen=None) + v1.view.addItem(s2)