new pyqtgraph viewer

Adrian Roth · Adrian Roth · commit 935fdc059e6e · 2021-04-08T15:32:36.000+02:00
diff --git a/pycine/cli/pfs_raw.py b/pycine/cli/pfs_raw.py
@@ -1,18 +1,31 @@
 #!/usr/bin/env python3
 import os
+import functools
 
 import click
 import cv2
 import numpy as np
 
-from pycine.color import color_pipeline, resize
+
+from pycine.color import color_pipeline
 from pycine.raw import read_frames
+from pycine.viewer import view_cine
+
+
+def image_post_processing(image, setup, bpp):
+    if setup.CFA in [3, 4]:
+        image = color_pipeline(image, setup=setup, bpp=bpp)
+    elif setup.CFA == 0:
+        pass
+    else:
+        raise ValueError("Sensor not supported")
 
+    if setup.EnableCrop:
+        image = image[setup.CropRect.top : setup.CropRect.bottom + 1, setup.CropRect.left : setup.CropRect.right + 1]
 
-def display(image_8bit):
-    cv2.imshow("image", image_8bit)
-    cv2.waitKey(0)
-    cv2.destroyAllWindows()
+    if setup.EnableResample:
+        image = cv2.resize(image, (setup.ResampleWidth, setup.ResampleHeight))
+    return image
 
 
 @click.command()
@@ -29,40 +42,24 @@ def cli(
     out_path: str,
     cine_file: str,
 ):
-    raw_images, setup, bpp = read_frames(cine_file, start_frame=start_frame, count=count)
-
-    if setup.CFA in [3, 4]:
-        # FIXME: the color pipeline is not at all ready for production!
-        images = (color_pipeline(raw_image, setup=setup, bpp=bpp) for raw_image in raw_images)
-
-    elif setup.CFA == 0:
-        images = raw_images
+    images, setup, bpp = read_frames(cine_file, start_frame=start_frame, count=count)
+    images.post_processing = functools.partial(image_post_processing, setup=setup, bpp=bpp)
 
-    else:
-        raise ValueError("Sensor not supported")
+    if not out_path:
+        # cine count is not used here
+        view_cine(images)
+        return
 
     for i, rgb_image in enumerate(images):
         frame_number = start_frame + i
 
-        if setup.EnableCrop:
-            rgb_image = rgb_image[
-                setup.CropRect.top : setup.CropRect.bottom + 1, setup.CropRect.left : setup.CropRect.right + 1
-            ]
-
-        if setup.EnableResample:
-            rgb_image = cv2.resize(rgb_image, (setup.ResampleWidth, setup.ResampleHeight))
-
-        if out_path:
-            ending = file_format.strip(".")
-            name = os.path.splitext(os.path.basename(cine_file))[0]
-            out_name = f"{name}-{frame_number:06d}.{ending}"
-            out_file = os.path.join(out_path, out_name)
-            print(f"Writing File {out_file}")
-            interpolated = np.interp(rgb_image, [0, 2 ** bpp - 1], [0, 2 ** 16 - 1]).astype(np.uint16)
-            cv2.imwrite(out_file, interpolated)
-
-        else:
-            display(resize(rgb_image, 720))
+        ending = file_format.strip(".")
+        name = os.path.splitext(os.path.basename(cine_file))[0]
+        out_name = f"{name}-{frame_number:06d}.{ending}"
+        out_file = os.path.join(out_path, out_name)
+        print(f"Writing File {out_file}")
+        interpolated = np.interp(rgb_image, [0, 2 ** bpp - 1], [0, 2 ** 16 - 1]).astype(np.uint16)
+        cv2.imwrite(out_file, interpolated)
 
 
 if __name__ == "__main__":
diff --git a/pycine/raw.py b/pycine/raw.py
@@ -12,36 +12,104 @@
 logger = logging.getLogger()
 
 
-def frame_reader(
-    cine_file: Union[str, bytes, PathLike],
-    header: Header,
-    start_frame: int = 1,
-    count: int = None,
-) -> Generator[np.ndarray, Any, None]:
-    frame = start_frame
-    if not count:
-        count = header["cinefileheader"].ImageCount
+class Frame_reader:
+    def __init__(self, cine_file, header=None, start_index=0, count=None, post_processing=None):
+        """
+        Create an object for reading frames from a cine file. It can either be used to get specific frames with __getitem__ or as an iterator.
+
+        Parameters
+        ----------
+        cine_file : str or file-like object
+            A string containing a path to a cine file
+        header : dict (optional)
+            A dictionary contains header information of the cine file
+        start_index : int
+            First image in a pile of images in cine file. Only used with the iterator in the object.
+        count : int
+            Maximum number of frames to get if this object is used as an iterator.
+        post_processing : function
+            Function that takes one image parameter and returns a new processed image
+            If provided this function will be applied to the raw image before returning.
+
+        Returns
+        -------
+        the created object
+        """
+        self.cine_file = cine_file
+        self.cine_file_stream = open(cine_file, "rb")
+        self.post_processing = post_processing
+
+        if header is None:
+            header = read_header(cine_file)
+        self.header = header
+
+        self.start_index = start_index
+        self.i = self.start_index
+
+        self.full_size = self.header["cinefileheader"].ImageCount
+        if not count:
+            count = self.full_size - self.start_index
+        self.end_index = self.start_index + count
+        if self.end_index > self.full_size:
+            raise ValueError("end_index {} is larger than the maximum {}".format(self.end_index, self.full_size))
+        self.size = count
+
+    def __getitem__(self, frame_index):
+        """
+        Dunder method to be able to use [] for retrieving images from the object.
+
+        Parameters
+        ----------
+        frame_index : int
+            the index for the image in cine file to retrieve.
+        """
+        f = self.cine_file_stream
+        try:
+            f.seek(self.header["pImage"][frame_index])
+        except IndexError:
+            raise IndexError(
+                "Index {} is out of bounds for cine_file with size {}.".format(frame_index, self.full_size)
+            )
 
-    with open(cine_file, "rb") as f:
-        while count:
-            frame_index = frame - 1
-            logger.debug(f"Reading frame {frame}")
+        annotation_size = struct.unpack("I", f.read(4))[0]
+        annotation = struct.unpack("{}B".format(annotation_size - 8), f.read((annotation_size - 8) // 8))
+        self.header["Annotation"] = annotation
 
-            f.seek(header["pImage"][frame_index])
+        image_size = struct.unpack("I", f.read(4))[0]
+        data = f.read(image_size)
+        image = create_raw_array(data, self.header)
+        if not self.post_processing is None:
+            image = self.post_processing(image)
+        return image
 
-            annotation_size = struct.unpack("I", f.read(4))[0]
-            annotation = struct.unpack(f"{annotation_size - 8}B", f.read((annotation_size - 8) // 8))
-            # TODO: Save annotations
+    def __iter__(self):
+        """ Object of this class is an iterator """
+        return self
 
-            image_size = struct.unpack("I", f.read(4))[0]
+    def __next__(self):
+        """ When iterating get the next image if more are left """
+        if self.i >= self.end_index:
+            raise StopIteration
+        logger.debug("Reading frame {}".format(self.i))
+        raw_image = self.__getitem__(self.i)
+        self.i += 1
+        return raw_image
 
-            data = f.read(image_size)
+    def __len__(self):
+        return self.size
 
-            raw_image = create_raw_array(data, header)
+    def __del__(self):
+        self.cine_file_stream.close()
 
-            yield raw_image
-            frame += 1
-            count -= 1
+
+# def frame_reader(cine_file, header=None, start_frame=1, count=None):
+def frame_reader(
+    cine_file: Union[str, bytes, PathLike],
+    header: Header,
+    start_frame: int = 1,
+    count: int = None,
+) -> Generator[np.ndarray, Any, None]:
+    return Frame_reader(cine_file, header, start_frame - 1, count)
 
 
 def read_bpp(header):
diff --git a/pycine/viewer.py b/pycine/viewer.py
@@ -0,0 +1,85 @@
+from PyQt5 import QtCore, QtWidgets
+import pyqtgraph as pg
+# backward compatibility default is col-major in pyqtgraph so change
+pg.setConfigOptions(imageAxisOrder="row-major")
+
+
+class Cine_viewer(QtWidgets.QMainWindow):
+    def __init__(self, frame_reader):
+        super(Cine_viewer, self).__init__()
+        self.frame_reader = frame_reader
+
+        self.setWindowTitle("Cine viewer")
+        size = (640, 480)
+        self.resize(*size)
+
+        layout = QtWidgets.QGridLayout()
+
+        image_view = pg.ImageView()
+        self.image_view = image_view
+        layout.addWidget(image_view, 0, 0)
+
+        if frame_reader.full_size == 1:
+            image = self.frame_reader[0]
+            self.image_view.setImage(image)
+        else:
+            frame_slider = QtWidgets.QSlider()
+            frame_slider.setOrientation(QtCore.Qt.Horizontal)
+            frame_slider.setMinimum(1) # maybe change to FirstImageNo in cine file
+            frame_slider.setMaximum(frame_reader.full_size)
+            frame_slider.setTracking(True)  # call update on mouse drag
+            frame_slider.setSingleStep(1)
+            frame_slider.setValue(frame_reader.start_index + 1)
+            self.frame_slider = frame_slider
+            layout.addWidget(frame_slider, 1, 0)
+
+            slider_label_left = QtWidgets.QLabel()
+            slider_label_left.setText("{}".format(frame_slider.minimum()))
+            layout.addWidget(slider_label_left, 2, 0)
+
+            slider_label = QtWidgets.QLabel()
+            slider_label.setAlignment(QtCore.Qt.AlignCenter)
+            self.slider_label = slider_label
+            layout.addWidget(slider_label, 2, 0)
+
+            slider_label_right = QtWidgets.QLabel()
+            slider_label_right.setText("{}".format(self.frame_slider.maximum()))
+            slider_label_right.setAlignment(QtCore.Qt.AlignRight)
+            layout.addWidget(slider_label_right, 2, 0)
+
+            self.update_frame(frame_reader.start_index + 1, autoLevels=True)
+            # do this last to avoid double update with slider setValue
+            frame_slider.valueChanged.connect(self.update_frame)
+
+        widget = QtWidgets.QWidget()
+        widget.setLayout(layout)
+        self.setCentralWidget(widget)
+
+    def set_frame(self, frame_index):
+        # indirectly call update frame
+        self.frame_slider.setValue(frame_index)
+
+    def update_frame(self, frame_index, autoLevels=False):
+        image = self.frame_reader[frame_index - 1]
+        self.image_view.setImage(image, autoLevels=autoLevels)
+        self.slider_label.setText("Frame: {}".format(frame_index))
+
+
+def view_cine(frame_reader):
+    """Start an interactive viewer for a cine file
+
+    Parameters
+    ----------
+    frame_reader : pycine.raw.Frame_reader
+        Object for the cine_file to be viewed
+
+    Returns
+    -------
+    None
+    """
+    app = QtWidgets.QApplication([])
+    window = Cine_viewer(frame_reader)
+    window.show()
+
+    # Start the event loop.
+    app.exec_()
diff --git a/setup.py b/setup.py
@@ -9,7 +9,7 @@
     description="This package allows handling of .cine files created by Vision Research Phantom® cameras.",
     entry_points={"console_scripts": ["pfs_meta = pycine.cli.pfs_meta:cli", "pfs_raw = pycine.cli.pfs_raw:cli"]},
     include_package_data=True,
-    install_requires=["click", "docopt", "opencv-python", "colorama", "timecode"],
+    install_requires=["click", "docopt", "opencv-python", "colorama", "timecode", "pyqt5", "pyqtgraph"],
     long_description=long_description,
     long_description_content_type="text/markdown",
     name="pycine",
