movie-barcodes 0.0.2__py3-none-any.whl

src/main.py

@@ -0,0 +1,181 @@
+import argparse
+import logging
+import time
+
+from typing import Callable
+from os import cpu_count, path
+
+from .barcode_generation import generate_circular_barcode, generate_barcode
+
+from .utility import (
+    save_barcode_image,
+    get_dominant_color_function,
+    format_time,
+    get_video_properties,
+    validate_args,
+)
+from .video_processing import load_video, extract_colors, parallel_extract_colors
+
+MAX_PROCESSES = cpu_count() or 1
+MIN_FRAME_COUNT = 2
+
+
+def generate_and_save_barcode(args: argparse.Namespace, dominant_color_function: Callable, method: str) -> None:
+    """
+    Generate and save the barcode image based on the specified method.
+
+    :param args: argparse.Namespace object containing the command-line arguments
+    :param dominant_color_function: The function to extract the dominant color from a frame
+    :param method: The method used to extract the dominant color
+    :return: None
+    """
+    start_time = time.time()
+
+    # Get Video Properties
+    video, frame_count, frame_width, frame_height = load_video(args.input_video_path)
+    _, _, video_duration, video_size = get_video_properties(video, args)
+
+    # If the user specifies the 'workers' argument
+    if args.workers is not None:
+        if args.workers == 1:
+            # If the user explicitly sets 'workers' to 1, use sequential processing
+            colors = extract_colors(
+                args.input_video_path,
+                0,
+                frame_count - 1,
+                dominant_color_function,
+                args.width,
+            )
+        else:
+            # Perform parallel processing with the user-specified number of workers
+            colors = parallel_extract_colors(
+                args.input_video_path,
+                frame_count,
+                dominant_color_function,
+                args.workers,
+                args.width,
+            )
+    else:
+        # If 'workers' is not specified, use the maximum number of available CPU cores
+        colors = parallel_extract_colors(
+            args.input_video_path,
+            frame_count,
+            dominant_color_function,
+            MAX_PROCESSES,
+            args.width,
+        )
+
+    # Generate the appropriate type of barcode
+    if args.barcode_type == "circular":
+        # Assuming image width = video frame width for circular barcodes
+        barcode = generate_circular_barcode(colors, frame_width)
+    else:
+        barcode = generate_barcode(colors, frame_height, frame_count, args.width)
+
+    base_name = path.basename(args.input_video_path)
+    file_name_without_extension = path.splitext(base_name)[0]
+    save_barcode_image(barcode, file_name_without_extension, args, method)
+
+    # Calculate processing time
+    end_time = time.time()
+    processing_time = end_time - start_time
+
+    # Log the information
+    logging.info("Processed File: %s", file_name_without_extension)
+    logging.info("Number of Frames: %d", frame_count)
+    logging.info("Video Duration: %s", format_time(video_duration))
+    logging.info("Video Size: %.2f MB", video_size / (1024 * 1024))
+    logging.info("Processing Time: %s", format_time(processing_time))
+
+    video.release()
+
+
+def main(args: argparse.Namespace) -> None:
+    """
+    Main function to generate a barcode from a video file.
+
+    :param args: argparse.Namespace object containing the command-line arguments
+    :return: None
+    """
+    # Check if the input video file exists
+    _, frame_count, _, _ = load_video(args.input_video_path)
+
+    # Check if the arguments are valid
+    validate_args(args, frame_count, MAX_PROCESSES, MIN_FRAME_COUNT)
+
+    # Get a list of all available methods
+    methods = ["avg", "hsv", "bgr", "kmeans"]
+
+    # Check if all_methods flag is set
+    if args.all_methods:
+        for method in methods:
+            # Generate barcodes for each method
+            dominant_color_function = get_dominant_color_function(method)
+            generate_and_save_barcode(args, dominant_color_function, method)
+    else:
+        # Use the specified method to generate barcode
+        dominant_color_function = get_dominant_color_function(args.method)
+        generate_and_save_barcode(args, dominant_color_function, args.method)
+
+
+if __name__ == "__main__":
+    logging.basicConfig(
+        filename=path.join("..", "logs.txt"),
+        level=logging.INFO,
+        format="%(asctime)s - %(message)s",
+    )
+    header_msg = "=" * 40 + " NEW RUN " + "=" * 40
+    logging.info("\n%s\n", header_msg)
+
+    parser = argparse.ArgumentParser(description="Generate a color barcode from a video file.")
+    parser.add_argument("--input_video_path", type=str, help="Path to the video file.")
+    parser.add_argument(
+        "--destination_path",
+        type=str,
+        nargs="?",
+        help="Path to save the output image. If not provided, the image will be saved in a default location.",
+        default=None,
+    )
+    parser.add_argument(
+        "--barcode_type",
+        choices=["horizontal", "circular"],
+        default="horizontal",
+        help="Type of barcode to generate: horizontal or circular. Default is horizontal.",
+    )
+    parser.add_argument(
+        "--method",
+        choices=["avg", "kmeans", "hsv", "bgr"],
+        default="avg",
+        help="Method to extract dominant color: avg (average), kmeans (K-Means clustering), hsv (HSV "
+        "histogram), or bgr (BGR histogram). Default is avg.",
+    )
+    parser.add_argument(
+        "--workers",
+        type=int,
+        default=None,
+        help="Number of workers for parallel processing. Default behavior uses all available CPU cores."
+        "Setting this to 1 will use sequential processing. Do not specify a value greater than "
+        "the number of available CPU cores.",
+    )
+    parser.add_argument(
+        "--width",
+        type=int,
+        default=None,
+        help="Width of the output image. If not provided, the width will be the same as the video",
+    )
+    parser.add_argument(
+        "--output_name",
+        type=str,
+        nargs="?",
+        help="Custom name for the output barcode image. If not provided, a name will be automatically " "generated.",
+        default=None,
+    )
+    parser.add_argument(
+        "--all_methods",
+        type=bool,
+        default=False,
+        help="If provided, all methods to extract dominant color will be used to create barcodes. "
+        "Overrides --method argument.",
+    )
+    arguments = parser.parse_args()
+    main(arguments)

src/utility.py

@@ -0,0 +1,168 @@
+import argparse
+from os import path, access, W_OK, makedirs
+from typing import Callable
+import cv2
+import numpy as np
+from PIL import Image
+
+
+from .color_extraction import (
+    get_dominant_color_mean,
+    get_dominant_color_kmeans,
+    get_dominant_color_hsv,
+    get_dominant_color_bgr,
+)
+
+
+def validate_args(args: argparse.Namespace, frame_count: int, MAX_PROCESSES: int, MIN_FRAME_COUNT: int) -> None:
+    """
+    Validate command-line arguments for logical errors.
+
+    :param argparse.Namespace args: The command-line arguments.
+    :param int frame_count: The number of frames in the video.
+    :param int MAX_PROCESSES: The maximum number of processes to use.
+    :param int MIN_FRAME_COUNT: The minimum number of frames required in the video.
+    :return: None
+    :raises FileNotFoundError: If the input video file does not exist.
+    :raises ValueError: If the video file has an invalid extension, the destination path is not writable, the number of
+        workers is invalid, the width is invalid, the frame count is invalid, or the method is invalid.
+    :raises PermissionError: If the destination path is not writable.
+    """
+    # Check if input video file exists
+    if not path.exists(args.input_video_path):
+        raise FileNotFoundError(f"The specified input video file '{args.input_video_path}' does not exist.")
+
+    valid_extensions = [".mp4", ".webm"]
+    if path.splitext(args.input_video_path)[1].lower() not in valid_extensions:
+        raise ValueError("The specified video file must have a valid video extension (e.g., .mp4).")
+
+    # Check if the destination path is writable
+    if args.destination_path is not None:
+        destination_dir = path.dirname(args.destination_path)
+        if not access(destination_dir, W_OK):
+            raise PermissionError(f"The specified destination path '{args.destination_path}' is not writable.")
+
+    if args.workers is not None:
+        if args.workers < 1:
+            raise ValueError("The number of workers must be greater than or equal to 1.")
+        if args.workers > MAX_PROCESSES:
+            raise ValueError(
+                f"The number of workers specified ({args.workers}) exceeds "
+                f"the number of available CPU cores ({MAX_PROCESSES})."
+            )
+
+    if args.width is not None:
+        if args.width <= 0:
+            raise ValueError("Width must be greater than 0.")
+        if args.width > frame_count:
+            raise ValueError(
+                f"Specified width ({args.width}) cannot be greater than the number of frames ({frame_count}) in the "
+                f"video."
+            )
+
+    if frame_count < MIN_FRAME_COUNT:
+        raise ValueError(f"The video must have at least {MIN_FRAME_COUNT} frames.")
+
+    if args.all_methods and args.method is not None:
+        raise ValueError("The --all_methods flag cannot be used with the --method argument.")
+
+
+def get_dominant_color_function(method: str) -> Callable:
+    """
+    Returns the appropriate function to get the dominant color based on the specified method.
+
+    :param str method: The method to use for color extraction ('avg', 'kmeans', 'hsv', or 'bgr').
+    :return: Function to get the dominant color.
+    :raises ValueError: If the method is invalid.
+    """
+    if method == "avg":
+        # return get_dominant_color_avg
+        return get_dominant_color_mean
+    if method == "kmeans":
+        return get_dominant_color_kmeans
+    if method == "hsv":
+        return get_dominant_color_hsv
+    if method == "bgr":
+        return get_dominant_color_bgr
+
+    raise ValueError(f"Invalid method: {method}")
+
+
+def format_time(seconds: float) -> str:
+    """
+    Formats time in seconds to a string representation.
+
+    :param int seconds: The time in seconds.
+    :return: Formatted time string.
+    """
+    hours, remainder = divmod(seconds, 3600)
+    minutes, seconds = divmod(remainder, 60)
+    if hours > 0:
+        return f"{int(hours)}h {int(minutes)}m {int(seconds)}s"
+
+    return f"{int(minutes)}m {int(seconds)}s"
+
+
+def get_video_properties(video: cv2.VideoCapture, args: argparse.Namespace) -> tuple:
+    """
+    Extracts and returns various properties of a video.
+
+    :param cv2.VideoCapture video: The video capture object.
+    :param args: Command line arguments.
+    :return: Tuple containing total frames, FPS, video duration in seconds, and video size in bytes.
+    """
+    total_frames = int(video.get(cv2.CAP_PROP_FRAME_COUNT))
+    fps = video.get(cv2.CAP_PROP_FPS)
+    video_duration = total_frames / fps if fps > 0 else 0  # in seconds
+    video_size = path.getsize(args.input_video_path)  # in bytes
+
+    return total_frames, fps, video_duration, video_size
+
+
+def save_barcode_image(barcode: np.ndarray, base_name: str, args: argparse.Namespace, method: str) -> None:
+    """
+    Saves a generated barcode image to a file.
+
+    :param np.ndarray barcode: The barcode image as a NumPy array.
+    :param str base_name: The base name of the file to save.
+    :param args: Command line arguments.
+    :param str method: The method used for color extraction.
+    """
+    current_dir = path.dirname(path.abspath(__file__))
+    project_root = path.dirname(current_dir)  # Go up one directory to get to the project root
+    # If destination_path isn't specified, construct one based on the video's name
+    if not args.destination_path:
+        barcode_dir = path.join(project_root, "barcodes")
+        ensure_directory(barcode_dir)
+
+        # If an output_name is provided by the user
+        if args.output_name:
+            destination_name = args.output_name + ".png"
+        else:
+            filename_parts = [base_name, method, args.barcode_type]
+            if args.workers:
+                filename_parts.append(f"workers_{str(args.workers)}")
+            destination_name = "_".join(filename_parts) + ".png"
+
+        destination_path = path.join(barcode_dir, destination_name)
+    else:
+        # In case a destination_path is provided, consider appending the method
+        # or managing as per your requirement
+        destination_path = path.join(project_root, args.destination_path)
+
+    if barcode.shape[2] == 4:  # If the image has an alpha channel (RGBA)
+        image = Image.fromarray(barcode, "RGBA")
+    else:  # If the image doesn't have an alpha channel (RGB)
+        image = Image.fromarray(barcode, "RGB")
+
+    image.save(destination_path)
+
+
+def ensure_directory(directory_name: str) -> None:
+    """
+    Ensures that a directory exists. Creates it if it doesn't.
+
+    :param str directory_name: The name of the directory to ensure.
+    """
+    if not path.exists(directory_name):
+        makedirs(directory_name)

src/video_processing.py

@@ -0,0 +1,149 @@
+from multiprocessing import Pool
+from typing import Callable, List, Optional
+
+import cv2
+import numpy as np
+from tqdm import tqdm
+
+
+def load_video(video_path: str) -> tuple:
+    """
+    Load a video file and return its properties.
+
+    :param str video_path: The path to the video file.
+    :return: Tuple containing the video capture object, frame count, frame width, and frame height.
+    """
+    video = cv2.VideoCapture(video_path)
+
+    if not video.isOpened():
+        raise ValueError(f"Could not open the video file: {video_path}")
+
+    frame_count = int(video.get(cv2.CAP_PROP_FRAME_COUNT))
+    frame_width = int(video.get(cv2.CAP_PROP_FRAME_WIDTH))
+    frame_height = int(video.get(cv2.CAP_PROP_FRAME_HEIGHT))
+
+    if frame_count <= 0:
+        raise ValueError(f"The video file {video_path} has no frames.")
+
+    if frame_width <= 0 or frame_height <= 0:
+        raise ValueError(f"The video file {video_path} has invalid dimensions.")
+
+    return video, frame_count, frame_width, frame_height
+
+
+def parallel_extract_colors(
+    video_path: str,
+    frame_count: int,
+    color_extractor: Callable,
+    workers: int,
+    target_frames: Optional[int] = None,
+) -> list:
+    """
+    Extract dominant colors from frames in a video file using parallel processing.
+
+    :param str video_path: The path to the video file.
+    :param int frame_count: The total number of frames in the video.
+    :param Callable color_extractor: A function to extract the dominant color from a frame.
+    :param int workers: Number of parallel workers.
+    :param Optional[int] target_frames: The total number of frames to sample.
+    :return: List of dominant colors for the frames in the video.
+    """
+    if target_frames is None:
+        target_frames = frame_count
+
+    frames_per_worker = frame_count // workers
+    target_frames_per_worker = target_frames // workers
+
+    with Pool(workers) as pool:
+        args = [
+            (
+                video_path,
+                i * frames_per_worker,
+                (i + 1) * frames_per_worker - 1,
+                color_extractor,
+                target_frames_per_worker,
+            )
+            for i in range(workers)
+        ]
+
+        if frame_count % workers != 0 or target_frames % workers != 0:
+            args[-1] = (
+                video_path,
+                args[-1][1],
+                frame_count - 1,
+                color_extractor,
+                target_frames - (workers - 1) * target_frames_per_worker,
+            )
+
+        results = pool.starmap(extract_colors, args)
+
+    # Concatenate results from all workers
+    final_colors = [color for colors in results for color in colors]
+
+    return final_colors
+
+
+def extract_colors(
+    video_path: str,
+    start_frame: int,
+    end_frame: int,
+    color_extractor: Callable,
+    target_frames: Optional[int] = None,
+) -> List:
+    """
+    Extracts dominant colors from frames in a video file.
+
+    :param str video_path: The video capture object.
+    :param int start_frame: The index of the first frame to process.
+    :param int end_frame: The index of the last frame to process.
+    :param Callable color_extractor: A function to extract the dominant color from a frame.
+    :param Optional[int] target_frames: The total number of frames to sample.
+    :return: List of dominant colors from the sampled frames.
+    """
+    video = cv2.VideoCapture(video_path)
+    video.set(cv2.CAP_PROP_POS_FRAMES, start_frame)
+
+    # Calculate frame_skip based on target_frames
+    total_frames = end_frame - start_frame + 1
+    if target_frames:
+        frame_skip = total_frames // target_frames
+    else:
+        frame_skip = 1
+
+    colors = []
+
+    for _ in tqdm(range(target_frames or total_frames), desc="Processing frames"):
+        ret, frame = video.read()  # Read the first or next frame
+        if ret:
+            dominant_color = color_extractor(frame)
+            colors.append(dominant_color)
+        for _ in range(frame_skip - 1):
+            video.grab()  # Skip frames
+
+    video.release()
+
+    return colors
+
+
+def crop_black_borders(frame: np.ndarray, threshold: int = 30) -> np.ndarray:
+    """
+    Crop out black borders from a frame.
+
+    :param np.ndarray frame: Input frame.
+    :param int threshold: Threshold below which a pixel is considered 'black'.
+    :return np.ndarray: Cropped frame.
+    """
+    gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
+
+    _, binary = cv2.threshold(gray, threshold, 255, cv2.THRESH_BINARY)
+
+    contours, _ = cv2.findContours(binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+    if not contours:
+        return frame
+
+    cnt = max(contours, key=cv2.contourArea)
+    x, y, w, h = cv2.boundingRect(cnt)
+
+    cropped_frame = frame[y : y + h, x : x + w]
+
+    return cropped_frame

tests/test_barcode_generation.py

@@ -0,0 +1,58 @@
+import unittest
+import numpy as np
+from src import (
+    barcode_generation,
+)  # Adjust the import as per your project structure and naming
+
+
+class TestBarcodeGeneration(unittest.TestCase):
+    """
+    Test the barcode generation functions.
+    """
+
+    def setUp(self) -> None:
+        """
+        Set up the test case.
+        :return: None
+        """
+        self.colors = [np.array([255, 0, 0]), np.array([0, 255, 0])]
+        self.frame_height = 2
+        self.frame_count = 2
+        self.img_size = 100
+        self.width = 1
+
+    def test_generate_barcode_default(self) -> None:
+        """
+        Test the generate_barcode function.
+        :return: None
+        """
+        barcode = barcode_generation.generate_barcode(self.colors, self.frame_height, self.frame_count)
+
+        self.assertIsInstance(barcode, np.ndarray)
+        self.assertEqual(
+            barcode.shape, (self.frame_height, self.frame_count, 3)
+        )  # Should match the input frame dimensions
+
+    def test_generate_barcode_with_width(self) -> None:
+        """
+        Test the generate_barcode function with a specified width.
+        :return: None
+        """
+        barcode = barcode_generation.generate_barcode(self.colors, self.frame_height, self.frame_count, self.width)
+
+        self.assertIsInstance(barcode, np.ndarray)
+        self.assertEqual(barcode.shape, (self.frame_height, self.width, 3))
+
+    def test_generate_circular_barcode(self) -> None:
+        """
+        Test the generate_circular_barcode function.
+        :return: None
+        """
+        barcode = barcode_generation.generate_circular_barcode(self.colors, self.img_size)
+
+        self.assertIsInstance(barcode, np.ndarray)
+        self.assertEqual(barcode.shape, (self.img_size, self.img_size, 4))  # RGBA image
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_color_extraction.py

@@ -0,0 +1,39 @@
+import unittest
+import numpy as np
+from src import (
+    color_extraction,
+)  # Adjust the import based on your project structure and naming
+
+
+class TestColorExtraction(unittest.TestCase):
+    """
+    Test the color extraction functions.
+    """
+
+    def setUp(self) -> None:
+        self.frame = np.array([[[255, 0, 0], [0, 255, 0]], [[0, 0, 255], [255, 255, 255]]], dtype=np.uint8)
+
+    def test_get_dominant_color_mean(self) -> None:
+        """
+        Test the get_dominant_color_mean function.
+        :return: None
+        """
+        # Call your function with the sample frame
+        dominant_color = color_extraction.get_dominant_color_mean(self.frame)
+
+        # Assert the expected result
+        self.assertEqual(dominant_color.tolist(), [127.5, 127.5, 127.5])
+
+    def test_get_dominant_color_kmeans(self) -> None:
+        """
+        Test the get_dominant_color_kmeans function.
+        :return: None
+        """
+        dominant_color = color_extraction.get_dominant_color_kmeans(self.frame)
+
+        self.assertIsInstance(dominant_color, np.ndarray)
+        self.assertEqual(dominant_color.shape, (3,))  # Should be a 3-element array representing a color
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_integration.py

@@ -0,0 +1,88 @@
+import argparse
+import unittest
+import os
+import glob
+
+from src import main
+
+
+class TestIntegration(unittest.TestCase):
+    def setUp(self) -> None:
+        """
+        Set up the test case.
+        """
+        # Recreate the parser with minimal necessary setup
+        self.parser = argparse.ArgumentParser(description="Test Parser")
+        self.parser.add_argument("--input_video_path", type=str)
+        self.parser.add_argument("--destination_path", type=str, nargs="?", default=None)
+        self.parser.add_argument("--barcode_type", choices=["horizontal", "circular"], default="horizontal")
+        self.parser.add_argument("--method", choices=["avg", "kmeans", "hsv", "bgr"], default="avg")
+        self.parser.add_argument("--workers", type=int, default=None)
+        self.parser.add_argument("--width", type=int, default=None)
+        self.parser.add_argument("--output_name", type=str, nargs="?", default=None)
+        self.parser.add_argument("--all_methods", type=bool, default=False)
+
+        self.input_video_path = "tests/sample.mp4"
+
+    def _run_test(self, barcode_type, workers, width=None):
+        args = [
+            "--input_video_path",
+            self.input_video_path,
+            "--barcode_type",
+            barcode_type,
+            "--method",
+            "avg",
+            "--workers",
+            str(workers),
+        ]
+        if width is not None:
+            args.extend(["--width", str(width)])
+
+        # Parse args using the recreated parser
+        parsed_args = self.parser.parse_args(args)
+
+        # Execute the program with the parsed Namespace object
+        main.main(parsed_args)
+
+    def test_horizontal_1_worker_default_width(self):
+        self._run_test("horizontal", 1)
+
+    def test_horizontal_1_worker_defined_width(self):
+        self._run_test("horizontal", 1, 90)
+
+    def test_horizontal_2_workers_default_width(self):
+        self._run_test("horizontal", 2)
+
+    def test_horizontal_2_workers_defined_width(self):
+        self._run_test("horizontal", 2, 90)
+
+    def test_circular_1_worker_default_width(self):
+        self._run_test("circular", 1)
+
+    def test_circular_1_worker_defined_width(self):
+        self._run_test("circular", 1, 90)
+
+    def test_circular_2_workers_default_width(self):
+        self._run_test("circular", 2)
+
+    def test_circular_2_workers_defined_width(self):
+        self._run_test("circular", 2, 90)
+
+    @classmethod
+    def tearDownClass(cls) -> None:
+        """
+        Clean up by deleting the generated images that start with "sample_"
+        after all tests in this class have run.
+        """
+        # Specify the path to the barcode images that start with "sample_"
+        barcode_images_path = "barcodes/sample_*"
+        # Use glob.glob to find all matching files
+        for img_file in glob.glob(barcode_images_path):
+            try:
+                os.remove(img_file)  # Remove the file
+            except OSError as e:
+                print(f"Error: {img_file} : {e.strerror}")
+
+
+if __name__ == "__main__":
+    unittest.main()