Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 6 additions & 7 deletions src/movie_barcodes/barcode_generation.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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
8 changes: 6 additions & 2 deletions src/movie_barcodes/utility.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
10 changes: 9 additions & 1 deletion src/movie_barcodes/video_processing.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
]
Expand All @@ -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)
Expand All @@ -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.
Expand All @@ -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)
Expand All @@ -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)
Expand Down