diff --git a/pycine/cli/pfs_raw.py b/pycine/cli/pfs_raw.py index 63b2d28..a115c1a 100755 --- 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 index feb95a9..25d9ac8 100644 --- 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 new file mode 100644 index 0000000..ba66e33 --- /dev/null +++ b/pycine/viewer.py @@ -0,0 +1,104 @@ +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_readers): + super(Cine_viewer, self).__init__() + + if not isinstance(frame_readers, list): + frame_readers = [frame_readers] + self.frame_readers = frame_readers + self.n_readers = len(frame_readers) + + self.setWindowTitle("Cine viewer") + size = (640 * self.n_readers, 480) + self.resize(*size) + + layout = QtWidgets.QGridLayout() + + # check that all fullsizes are the same + full_size = frame_readers[0].full_size + for i in range(1, self.n_readers): + assert full_size == frame_readers[i].full_size + # use only the start_index of the first one + start_frame = frame_readers[0].start_index + 1 + + self.image_views = [] + for i in range(self.n_readers): + image_view = pg.ImageView() + self.image_views.append(image_view) + layout.addWidget(image_view, 0, i) + + if full_size == 1: + for i in range(self.n_readers): + image = self.frame_readers[i][0] + self.image_views[i].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(full_size) + frame_slider.setTracking(True) # call update on mouse drag + frame_slider.setSingleStep(1) + frame_slider.setValue(start_frame) + self.frame_slider = frame_slider + layout.addWidget(frame_slider, 1, 0, 1, self.n_readers) + + slider_label_left = QtWidgets.QLabel() + slider_label_left.setText("{}".format(frame_slider.minimum())) + layout.addWidget(slider_label_left, 2, 0, 1, self.n_readers) + + slider_label = QtWidgets.QLabel() + slider_label.setAlignment(QtCore.Qt.AlignCenter) + self.slider_label = slider_label + layout.addWidget(slider_label, 2, 0, 1, self.n_readers) + + 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, 1, self.n_readers) + + for i in range(self.n_readers): + self.update_frame(start_frame, auto=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, auto=False): + for i in range(self.n_readers): + # keep the transform when changing images + transform = self.image_views[i].imageItem.transform() if not auto else None + image = self.frame_readers[i][frame_index - 1] + self.image_views[i].setImage(image, autoRange=auto, autoLevels=auto, transform=transform) + self.slider_label.setText("Frame: {}".format(frame_index)) + + +def view_cine(frame_readers): + """Start an interactive viewer for a cine file + + Parameters + ---------- + frame_readers : pycine.raw.Frame_reader or list of Frame_readers + Object(s) for the cine_file to be viewed + + Returns + ------- + None + """ + app = QtWidgets.QApplication([]) + window = Cine_viewer(frame_readers) + window.show() + + # Start the event loop. + app.exec_() diff --git a/setup.py b/setup.py index 1faab22..bfb1d0f 100644 --- a/setup.py +++ b/setup.py @@ -7,9 +7,9 @@ author="Ben Hagen", author_email="ben@ottomatic.io", 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"]}, + entry_points={"console_scripts": ["cine_meta = pycine.cli.pfs_meta:cli", "cine_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",