diff --git a/src/movie_barcodes/barcode_generation.py b/src/movie_barcodes/barcode_generation.py index 8fede55..33539e2 100644 --- a/src/movie_barcodes/barcode_generation.py +++ b/src/movie_barcodes/barcode_generation.py @@ -35,13 +35,13 @@ def generate_circular_barcode(colors: list, img_size: int, scale_factor: int = 1 ): radius = (idx + 1) * radius_increment - # Handle both simple RGB tuples and smoothed frames + # Handle both simple BGR tuples and smoothed frames if isinstance(color, np.ndarray) and color.ndim > 1: - # For smoothed frames, take the average color + # For smoothed frames, take the average BGR color color = np.mean(color, axis=(0, 1)).astype(int) - # Ensure color is in BGR format for OpenCV - bgr_color = color[::-1] if len(color) == 3 else color[2::-1] + # Colors are BGR throughout the pipeline; draw directly in BGR + bgr_color = color cv2.circle( barcode_high_res, @@ -89,8 +89,7 @@ def generate_barcode( # For single color values color_column = np.full((frame_height, 1, 3), color, dtype=np.uint8) - # Convert to BGR (which will be interpreted as RGB when saved with PIL) - bgr_color = cv2.cvtColor(color_column, cv2.COLOR_RGB2BGR) - barcode[:, i] = bgr_color.reshape(frame_height, 3) + # Keep BGR internally; perform RGB conversion once at save-time + barcode[:, i] = color_column.reshape(frame_height, 3) return barcode diff --git a/src/movie_barcodes/utility.py b/src/movie_barcodes/utility.py index 75c911d..f591e0c 100644 --- a/src/movie_barcodes/utility.py +++ b/src/movie_barcodes/utility.py @@ -160,9 +160,13 @@ def save_barcode_image(barcode: np.ndarray, base_name: str, args: argparse.Names destination_path = path.join(project_root, destination_path) if barcode.shape[2] == 4: # If the image has an alpha channel (RGBA) - image = Image.fromarray(barcode, "RGBA") + # Convert BGRA -> RGBA once at save-time + barcode_to_save = cv2.cvtColor(barcode, cv2.COLOR_BGRA2RGBA) + image = Image.fromarray(barcode_to_save, "RGBA") else: # If the image doesn't have an alpha channel (RGB) - image = Image.fromarray(barcode, "RGB") + # Convert BGR -> RGB once at save-time + barcode_to_save = cv2.cvtColor(barcode, cv2.COLOR_BGR2RGB) + image = Image.fromarray(barcode_to_save, "RGB") image.save(destination_path) logging.info("File saved at '%s'", destination_path) diff --git a/src/movie_barcodes/video_processing.py b/src/movie_barcodes/video_processing.py index 2c7b459..da0d5f4 100644 --- a/src/movie_barcodes/video_processing.py +++ b/src/movie_barcodes/video_processing.py @@ -62,6 +62,7 @@ def parallel_extract_colors( (i + 1) * frames_per_worker - 1, color_extractor, target_frames_per_worker, + False, # disable worker progress bars ) for i in range(workers) ] @@ -73,6 +74,7 @@ def parallel_extract_colors( frame_count - 1, color_extractor, target_frames - (workers - 1) * target_frames_per_worker, + False, ) results = pool.starmap(extract_colors, args) @@ -89,6 +91,7 @@ def extract_colors( end_frame: int, color_extractor: Callable, target_frames: Optional[int] = None, + show_progress: bool = True, ) -> List: """ Extracts dominant colors from frames in a video file. @@ -98,6 +101,7 @@ def extract_colors( :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. + :param bool show_progress: Whether to display a tqdm progress bar during extraction. :return: List of dominant colors from the sampled frames. """ video = cv2.VideoCapture(video_path) @@ -112,7 +116,11 @@ def extract_colors( colors = [] - for _ in tqdm(range(target_frames or total_frames), desc="Processing frames"): + iterator = range(target_frames or total_frames) + if show_progress: + iterator = tqdm(iterator, desc="Processing frames") + + for _ in iterator: ret, frame = video.read() # Read the first or next frame if ret: dominant_color = color_extractor(frame)