(expandable_layout, arg_has_frames, input_batch, arg_batch)
| 217 | |
| 218 | |
| 219 | def expand_arg(expandable_layout, arg_has_frames, input_batch, arg_batch): |
| 220 | num_expand = len(expandable_layout) |
| 221 | assert 1 <= num_expand <= 2 |
| 222 | assert all(c in "FC" for c in expandable_layout) |
| 223 | assert len(input_batch) == len(arg_batch) |
| 224 | expanded_batch = [] |
| 225 | for input_sample, arg_sample in zip(input_batch, arg_batch): |
| 226 | if not arg_has_frames or len(arg_sample) == 1: |
| 227 | arg_sample = arg_sample if not arg_has_frames else arg_sample[0] |
| 228 | num_frames = np.prod(input_sample.shape[:num_expand]) |
| 229 | expanded_batch.extend(arg_sample for _ in range(num_frames)) |
| 230 | else: |
| 231 | frame_idx = expandable_layout.find("F") |
| 232 | assert frame_idx >= 0 |
| 233 | assert len(arg_sample) == input_sample.shape[frame_idx] |
| 234 | if num_expand == 1: |
| 235 | expanded_batch.extend(arg_frame for arg_frame in arg_sample) |
| 236 | else: |
| 237 | channel_idx = 1 - frame_idx |
| 238 | assert expandable_layout[channel_idx] == "C" |
| 239 | if channel_idx > frame_idx: |
| 240 | expanded_batch.extend( |
| 241 | frame_arg |
| 242 | for frame_arg in arg_sample |
| 243 | for _ in range(input_sample.shape[channel_idx]) |
| 244 | ) |
| 245 | else: |
| 246 | expanded_batch.extend( |
| 247 | frame_arg |
| 248 | for _ in range(input_sample.shape[channel_idx]) |
| 249 | for frame_arg in arg_sample |
| 250 | ) |
| 251 | return expanded_batch |
| 252 | |
| 253 | |
| 254 | def expand_arg_input(input_data: ArgData, arg_data: ArgData): |
no test coverage detected