e2evideo.image_preprocessing

  1"""
  2Image preprocessing module.
  3Main functionality:
  4    * Resize images according to a given input dimensions.
  5    * Convert images to grayscale.
  6"""
  7# pylint: disable=no-member
  8import os
  9from typing import Optional
 10import argparse
 11from dataclasses import dataclass
 12import glob
 13import numpy as np
 14import cv2
 15from skimage import img_as_float32
 16
 17
 18@dataclass
 19class ImagesConfig:
 20    """Class to hold the configuration of the images."""
 21
 22    dir: str
 23    img_format: str
 24    resize: bool
 25    gray_scale: bool
 26    output: str
 27    img_width: Optional[int] = None
 28    img_height: Optional[int] = None
 29
 30
 31class ImagePreprocessing:
 32    """Images preprocessing class."""
 33
 34    def __init__(self, config: ImagesConfig):
 35        self.config = config
 36
 37    def get_images_helper(self, folder_name):
 38        """Helper function for get_images."""
 39        video_file = []
 40        images_ = glob.glob(folder_name + "/" + self.config.img_format)
 41        for img_ in images_:
 42            img = cv2.imread(img_)  # pylint: disable=E1101
 43            if self.config.gray_scale is True:
 44                img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)  # pylint: disable=E1101
 45            if self.config.resize is True:
 46                img = cv2.resize(
 47                    img, (self.config.img_width, self.config.img_height)
 48                )  # pylint: disable=E1101
 49
 50            video_file.append(img_as_float32(img))
 51            continue
 52        return video_file
 53
 54    def get_images(self):
 55        """
 56        Read images from a given dir, using specified image format.
 57        Additionally it allows for resizing the images.
 58        Parameters:
 59            - dir_: directory contains the images.
 60            - img_format: specify the images format, it is string in the format "*.xxx".
 61            - resize: this is a boolean, set True if resizing of the images are needed.
 62            - resize_dim: set to the required image dimensions.
 63                        It takes input in the format (height, width).
 64        Returns:
 65            -- all_images: an array of array, with values from the images in the
 66            dir.
 67            This array is of size = (num_images, height, width, no_channels)
 68        """
 69        img_folders = [x[0] for x in os.walk(self.config.dir)]
 70        all_videos = []
 71        labels = []
 72        for folder_name in img_folders:
 73            video_file = self.get_images_helper(folder_name)
 74            if len(video_file) != 0:
 75                # convert video_file to numpy array
 76                video_file = np.array(video_file)
 77                all_videos.append(video_file)
 78                labels.append(folder_name.split("/")[-1])
 79
 80        assert (
 81            len(all_videos) != 0
 82        ), "The given images folder does not contain any frames"
 83        # find the maximum length of the videos (number of frames) in a video
 84        max_frames = max(len(x) for x in all_videos)
 85        # find the maximum shape of the arrays
 86        # max_shape = max([arr.shape for arr in all_images])
 87        # create a new array with the maximum shape
 88        # specify the desired shape of the padded arrays
 89        frame_dim = all_videos[0][0].shape
 90        frames_in_videos_dim = (len(all_videos), max_frames) + frame_dim
 91        frames_in_videos = np.zeros(frames_in_videos_dim, dtype=np.float64)
 92
 93        # pad the shorter videos with zeros at the end to make them all the same length
 94        for index_, video_ in enumerate(all_videos):
 95            frames_in_videos[index_][0 : len(video_)] = video_
 96        # save frames_in_videos to a file
 97        np.savez_compressed(self.config.output, frames_in_videos)
 98        # save labels to frames_labels.txt file
 99        with open("frames_labels.txt", "w", encoding="utf-8") as my_file:
100            for label in labels:
101                my_file.write(label + "\n")
102        return frames_in_videos, labels
103
104    def create_videos(self, folder_path, output_path, image_format):
105        """Create videos from the extracted frames."""
106        img_array = []
107        # check if output path has .avi extension
108        if output_path[-4:] != ".avi":
109            output_path = output_path + "new_video.avi"
110        for filename in sorted(glob.glob(folder_path + "/*." + image_format)):
111            img = cv2.imread(filename)  # pylint: disable=E1101
112            height, width, _ = img.shape
113            size = (width, height)
114            img_array.append(img)
115        # pylint: disable=E1101
116        out = cv2.VideoWriter(output_path, cv2.VideoWriter_fourcc(*"DIVX"), 10, size)
117        for _, img in enumerate(img_array):
118            out.write(img)
119        out.release()
120
121
122def main():
123    parser_ = argparse.ArgumentParser()
124    parser_.add_argument("--dir", default="./images/")
125    parser_.add_argument("--img_format", default="*.jpg")
126    parser_.add_argument("--resize", default=False)
127    parser_.add_argument("--img_width", default=224, type=int)
128    parser_.add_argument("--img_height", default=224, type=int)
129    parser_.add_argument("--gray_scale", default=False)
130    parser_.add_argument("--output", default="./results/all_images.npz")
131
132    args_ = parser_.parse_args()
133
134    images_config = ImagesConfig(
135        args_.dir,
136        args_.img_format,
137        args_.resize,
138        args_.gray_scale,
139        args_.output,
140        args_.img_width,
141        args_.img_height,
142    )
143    image_preprocessing = ImagePreprocessing(images_config)
144    _images, _ = image_preprocessing.get_images()
145
146    print("Images saved in array of array of size", str(_images.shape))
147
148
149if __name__ == "__main__":
150    main()