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
7 changes: 1 addition & 6 deletions src/movie_barcodes/utility.py
Original file line number Diff line number Diff line change
Expand Up @@ -40,7 +40,7 @@ def validate_args(args: argparse.Namespace, frame_count: int, MAX_PROCESSES: int

# Check if the destination path is writable
if args.destination_path is not None:
destination_dir = path.dirname(args.destination_path)
destination_dir = path.dirname(args.destination_path) or "."
if not access(destination_dir, W_OK):
raise PermissionError(f"The specified destination path '{args.destination_path}' is not writable.")

Expand All @@ -62,15 +62,10 @@ def validate_args(args: argparse.Namespace, frame_count: int, MAX_PROCESSES: int
if args.height is not None:
if args.height <= 0:
raise ValueError("Height must be greater than 0.")
if args.height > frame_count:
raise ValueError("Height must be less than or equal to the number of frames.")

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:
"""
Expand Down
64 changes: 35 additions & 29 deletions src/movie_barcodes/video_processing.py
Original file line number Diff line number Diff line change
Expand Up @@ -51,33 +51,45 @@ def parallel_extract_colors(
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 = [
# Cap workers to available work to avoid empty tasks
active_workers = max(1, min(workers, frame_count, target_frames))

# Evenly distribute frame ranges across active workers
base_frames = frame_count // active_workers
remainder_frames = frame_count % active_workers
frame_ranges = []
start = 0
for i in range(active_workers):
length = base_frames + (1 if i < remainder_frames else 0)
end = start + length - 1
frame_ranges.append((start, end))
start = end + 1

# Evenly distribute target samples ensuring total equals target_frames
base_samples = target_frames // active_workers
remainder_samples = target_frames % active_workers
samples_per_worker = [base_samples + (1 if i < remainder_samples else 0) for i in range(active_workers)]

# Build only tasks that have at least one sample (avoid passing 0 -> falsy)
task_args = []
for i in range(active_workers):
if samples_per_worker[i] <= 0:
continue
start_frame, end_frame = frame_ranges[i]
if end_frame < start_frame:
continue
task_args.append(
(
video_path,
i * frames_per_worker,
(i + 1) * frames_per_worker - 1,
color_extractor,
target_frames_per_worker,
False, # disable worker progress bars
)
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,
start_frame,
end_frame,
color_extractor,
target_frames - (workers - 1) * target_frames_per_worker,
False,
samples_per_worker[i],
)
)

results = pool.starmap(extract_colors, args)
with Pool(active_workers) as pool:
results = pool.starmap(extract_colors, task_args)

# Concatenate results from all workers
final_colors = [color for colors in results for color in colors]
Expand All @@ -91,7 +103,6 @@ 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 @@ -101,7 +112,6 @@ 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 @@ -116,11 +126,7 @@ def extract_colors(

colors = []

iterator = range(target_frames or total_frames)
if show_progress:
iterator = tqdm(iterator, desc="Processing frames")

for _ in iterator:
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)
Expand Down
11 changes: 5 additions & 6 deletions tests/test_utility.py
Original file line number Diff line number Diff line change
Expand Up @@ -114,9 +114,9 @@ def test_invalid_height(self) -> None:
with self.assertRaises(ValueError):
utility.validate_args(self.args, self.frame_count, self.MAX_PROCESSES, self.MIN_FRAME_COUNT)

self.args.height = self.frame_count + 1 # Testing for > frame_count
with self.assertRaises(ValueError):
utility.validate_args(self.args, self.frame_count, self.MAX_PROCESSES, self.MIN_FRAME_COUNT)
# Heights greater than frame_count are now allowed (no exception expected)
self.args.height = self.frame_count + 1
utility.validate_args(self.args, self.frame_count, self.MAX_PROCESSES, self.MIN_FRAME_COUNT)

def test_frame_count_too_low(self) -> None:
"""
Expand Down Expand Up @@ -144,13 +144,12 @@ def test_frame_count_too_low_specific_branch(self, mock_exists: MagicMock) -> No

def test_all_methods_and_method_error(self) -> None:
"""
Test that validate_args raises a ValueError when the --all_methods flag is used with the --method argument.
When --all_methods is provided, it should no longer raise even if --method is set.
:return: None
"""
self.args.all_methods = True
self.args.method = "avg" # Explicitly setting method to simulate conflict
with self.assertRaises(ValueError):
utility.validate_args(self.args, self.frame_count, self.MAX_PROCESSES, self.MIN_FRAME_COUNT)
utility.validate_args(self.args, self.frame_count, self.MAX_PROCESSES, self.MIN_FRAME_COUNT)

def test_no_error_raised(self) -> None:
"""
Expand Down