st-ten-1/src/components/galaxy_camera.py

import pathlib
import sys
from itertools import cycle

import cv2
import imutils
import numpy as np
from PyQt5.QtCore import QMutex, Qt, QThread, pyqtSignal
from PyQt5.QtGui import QImage, QPixmap
from PyQt5.QtWidgets import (QDialog, QFormLayout, QLabel,
                             QPushButton, QSizePolicy, QSlider)
try:
    from . import gxipy as gx
except Exception:
    from dummies import gxpy as gx

if "--sim-camera" in sys.argv:
    from components.dummies.gxpy import DummyCamera

from datetime import datetime

from .component import Component


class GalaxyCamera(Component):
    _edits_new_frame = pyqtSignal(object)

    def __init__(self, config=None, name=None, period=1, lazy=True, paused=False, threaded=True, registers=None):
        super().__init__(config=config, name=name, period=period, lazy=lazy, paused=paused, threaded=threaded)
        self.lock = QMutex()
        self.simulate = "--sim-camera" in sys.argv
        self._last_frame = None

    def config_changed(self):
        if self.simulate:
            self.device_manager = None
            self.camera = DummyCamera()
            self.sim_imgs = cycle(sorted(pathlib.Path("data/simulation_images/").glob("*.png")))
        else:
            self.device_manager = gx.DeviceManager()
            # create a device manager
            dev_num, dev_info_list = self.device_manager.update_device_list()
            if dev_num == 0:
                raise AssertionError(f"camera not detected: dev_num = {dev_num!r}")
            self.camera = self.device_manager.open_device_by_index(1)
        self.resolution = {
            "w": int(self.round_4(self.config[self.name]["horizontal_resolution"])),
            "h": int(self.round_4(self.config[self.name]["vertical_resolution"])),
        }
        self._period = int(self.config[self.name].get("frame_time_ms", 300)) / 1000
        self.exposure_time = int(self.config[self.name]["exposure_time"])
        self.roi = {
            "x": int(self.round_4(self.config[self.name].get("horizontal_crop_offset", 0))),
            "w": int(self.round_4(self.config[self.name].get("horizontal_crop_resolution", self.resolution["w"]))),
            "y": int(self.round_4(self.config[self.name].get("vertical_crop_offset", 0))),
            "h": int(self.round_4(self.config[self.name].get("vertical_crop_resolution", self.resolution["h"]))),
            "r": int(self.config[self.name].get("rotate_90_clockwise_times", 0)),
        }
        self.auto_white_balance = self.config[self.name].get("auto_white_balance", "").lower() in {"on", "1", "y", "yes", "true", "enable", "enabled", }
        self.balance = {
            "r": float(self.config[self.name].get("balance_red", 1)),
            "g": float(self.config[self.name].get("balance_green", 1)),
            "b": float(self.config[self.name].get("balance_blue", 1)),
        }
        self.lock.lock()
        self.camera.stream_off()
        self.camera.TriggerMode.set(gx.GxSwitchEntry.OFF)
        self.camera.ExposureTime.set(self.exposure_time)
        self.camera.Gain.set(1.0)
        self.camera.OffsetX.set(self.roi["x"])
        self.camera.Width.set(self.roi["w"])
        self.camera.OffsetY.set(self.roi["y"])
        self.camera.Height.set(self.roi["h"])
        if self.auto_white_balance:
            self.camera.BalanceWhiteAuto.set(gx.GxAutoEntry.ONCE)
            QThread.msleep(3000)
            self.camera.BalanceRatioSelector.set(gx.GxBalanceRatioSelectorEntry.RED)
            self.balance["r"] = self.camera.BalanceRatio.get()
            self.camera.BalanceRatioSelector.set(gx.GxBalanceRatioSelectorEntry.GREEN)
            self.balance["g"] = self.camera.BalanceRatio.get()
            self.camera.BalanceRatioSelector.set(gx.GxBalanceRatioSelectorEntry.BLUE)
            self.balance["b"] = self.camera.BalanceRatio.get()
        else:
            self.camera.BalanceRatioSelector.set(gx.GxBalanceRatioSelectorEntry.RED)
            self.camera.BalanceRatio.set(self.balance["r"])
            self.camera.BalanceRatioSelector.set(gx.GxBalanceRatioSelectorEntry.GREEN)
            self.camera.BalanceRatio.set(self.balance["g"])
            self.camera.BalanceRatioSelector.set(gx.GxBalanceRatioSelectorEntry.BLUE)
            self.camera.BalanceRatio.set(self.balance["b"])
        if self.camera.GammaParam.is_readable():
            self.gamma_lut = gx.Utility.get_gamma_lut(self.camera.GammaParam.get())
        else:
            self.gamma_lut = None
        if self.camera.ContrastParam.is_readable():
            self.contrast_lut = gx.Utility.get_contrast_lut(self.camera.ContrastParam.get())
        else:
            self.contrast_lut = None
        if self.camera.ColorCorrectionParam.is_readable():
            self.color_correction = self.camera.ColorCorrectionParam.get()
        else:
            self.color_correction = 0
        self.camera.TriggerMode.set(gx.GxSwitchEntry.ON)
        self.camera.TriggerSource.set(gx.GxTriggerSourceEntry.SOFTWARE)
        self.camera.stream_on()
        self.edits = None
        self.lock.unlock()
        self.edits_enabled = "--camera-edits" in sys.argv
        if self.edits_enabled:
            self.init_edits()

    @staticmethod
    def round_4(x):
        return 4 * round(int(x) / 4)

    @Component.reconfig_on_error
    def _get(self):
        # print("GALAXY CAMERA", str(int(QThread.currentThreadId())), flush=True)
        frame = None
        self.lock.lock()
        if self.simulate:
            img_path = next(self.sim_imgs)
            self.log.debug(f"loading image {img_path}")
            frame = cv2.cvtColor(cv2.imread(str(img_path)), cv2.COLOR_BGR2RGB)
            QThread.msleep(1000)
        else:
            self.camera.TriggerSoftware.send_command()
            frame = self.camera.data_stream[0].get_image()
            if frame is not None:
                if frame.get_status() == gx.GxFrameStatusList.INCOMPLETE:
                    self.log.error("incomplete frame")
                    frame = None
                else:
                    frame = frame.convert("RGB")
                    frame.image_improvement(self.color_correction, self.contrast_lut, self.gamma_lut)
                    frame = frame.get_numpy_array()
                    frame = np.rot90(frame, self.roi["r"])
        self.lock.unlock()
        if frame is None:
            self.log.error("failed to get frame")
        elif self.edits_enabled:
            frame = self.edit(frame, self.edits)
            self._edits_new_frame.emit([frame])
        super()._get([frame])

    def __del__(self, event=None):
        if self.camera is not None:
            self.camera.stream_off()
            self.camera.close_device()

    def init_edits(self):
        self.edits_enabled = True
        self.edits_dialog = EditsDialog(self.roi)
        self.edits_dialog.edits_changed.connect(self.set_edits)
        self._edits_new_frame.connect(self.edits_dialog.save_and_show_edits_new_frame)
        self.edits_dialog.edits_pause.connect(self.toggle_paused)

    def toggle_paused(self):
        if self.running:
            self.pause()
        else:
            self.resume()

    def set_edits(self, edits=None):
        self.edits = edits

    def edit(self, img, edits=None):
        if edits is None:
            return img
        # BRIGHTNESS AND CONTRAST
        contrast = float(edits.get("contrast", 0))
        brightness = float(edits.get("brightness", 0))
        if not (contrast == 0 and brightness == 0):
            img = imutils.adjust_brightness_contrast(img, brightness=brightness, contrast=contrast / 255 * 500)
        # ROTATE AND SCALE
        rotation = -float(edits.get("rotation", 0))
        scale = float(edits.get("scale", 1))
        if not (rotation == 0 and scale == 1):
            img = imutils.rotate(img, rotation, scale=scale)
        # TRANSLATE
        translation_x = float(edits.get("translation_x", 0))
        translation_y = float(edits.get("translation_y", 0))
        if not (translation_x == 0 and translation_y == 0):
            img = imutils.translate(img, translation_x, translation_y)
        return img


class EditsDialog(QDialog):
    edits_changed = pyqtSignal(object)
    edits_pause = pyqtSignal()

    def __init__(self, roi):
        super().__init__()
        self.frame = None
        self.edits = {
            "brightness": 0,
            "contrast": 0,
            "rotation": 0,
            "scale": 1,
            "translation_x": 0,
            "translation_y": 0,
        }
        self.edits_specs = {
            "brightness": [[-255, 255], 1, QSlider(Qt.Horizontal), QLabel()],
            "contrast": [[-255, 255], 1, QSlider(Qt.Horizontal), QLabel()],
            "rotation": [[-180, 180], 1, QSlider(Qt.Horizontal), QLabel()],
            "scale": [[0, 5], 100, QSlider(Qt.Horizontal), QLabel()],
            "translation_x": [[-roi["w"], roi["w"]], 1, QSlider(Qt.Horizontal), QLabel()],
            "translation_y": [[-roi["h"], roi["h"]], 1, QSlider(Qt.Horizontal), QLabel()],
        }
        layout = QFormLayout()
        self.edits_frame_l = QLabel()
        self.edits_frame_l.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Expanding)
        layout.addRow(self.edits_frame_l)
        self.edits_pause_b = QPushButton("toggle paused")
        layout.addRow(self.edits_pause_b)
        self.edits_pause_b.clicked.connect(self.edits_pause)
        for edit, limits in self.edits_specs.items():
            limit, multiplier, slider, label = limits
            slider.setRange(limit[0] * multiplier, limit[1] * multiplier)
            slider.setSingleStep(1)
            slider.setValue(self.edits[edit] * multiplier)
            label.setText(f"{edit}: {self.edits[edit]}")
            layout.addRow(label, slider)
            slider.valueChanged.connect(self.update_edits)
        self.edits_save_frame_b = QPushButton("save frame")
        layout.addRow(self.edits_save_frame_b)
        self.edits_save_frame_b.clicked.connect(self.edits_save_frame)
        self.setLayout(layout)
        self.update_edits()
        self.show()

    def update_edits(self):
        for edit, limits in self.edits_specs.items():
            limit, multiplier, slider, label = limits
            self.edits[edit] = slider.value() / multiplier
            label.setText(f"{edit}: {self.edits[edit]}")
        self.edits_changed.emit(self.edits)

    def save_and_show_edits_new_frame(self, frame):
        self.frame = frame[0]
        if self.frame is not None:
            self.edits_frame_l.setPixmap(
                QPixmap.fromImage(
                    QImage(
                        self.frame.data,
                        self.frame.shape[1],  # width
                        self.frame.shape[0],  # height
                        self.frame.shape[2] * self.frame.shape[1],  # width * channels
                        QImage.Format_RGB888
                    )
                ).scaled(
                    max(self.edits_frame_l.width(), 640),
                    max(self.edits_frame_l.height(), 480),
                    Qt.KeepAspectRatio,
                    Qt.SmoothTransformation
                )
            )

    def edits_save_frame(self):
        if self.frame is None:
            return
        out_path = f"./{datetime.now().isoformat()}.png"
        self.log.info(f"saving frame: {out_path!r}")
        img = cv2.cvtColor(self.frame, cv2.COLOR_RGB2BGR)
        cv2.imwrite(out_path, img)
        return out_path