Prepare background canvas, pure white if not set.
| 276 | |
| 277 | |
| 278 | class _CavasProducer: |
| 279 | """Prepare background canvas, pure white if not set.""" |
| 280 | def __init__(self, |
| 281 | frame_list, |
| 282 | resolution, |
| 283 | kp2d=None, |
| 284 | image_array=None, |
| 285 | default_scale=1.5): |
| 286 | """Initialize a canvas writer.""" |
| 287 | # check the origin background frames |
| 288 | if frame_list is not None: |
| 289 | _check_frame_path(frame_list) |
| 290 | self.frame_list = frame_list |
| 291 | else: |
| 292 | self.frame_list = [] |
| 293 | self.resolution = resolution |
| 294 | self.kp2d = kp2d |
| 295 | |
| 296 | # with numpy array frames |
| 297 | self.image_array = image_array |
| 298 | |
| 299 | if self.resolution is None: |
| 300 | if self.image_array is not None: |
| 301 | self.auto_resolution = self.image_array.shape[1:3] |
| 302 | elif len(self.frame_list) > 1 and \ |
| 303 | check_path_existence( |
| 304 | self.frame_list[0], 'file') == Existence.FileExist: |
| 305 | tmp_image_array = cv2.imread(self.frame_list[0]) |
| 306 | self.auto_resolution = tmp_image_array.shape[:2] |
| 307 | else: |
| 308 | |
| 309 | self.auto_resolution = [ |
| 310 | int(np.max(kp2d) * default_scale), |
| 311 | int(np.max(kp2d) * default_scale) |
| 312 | ] |
| 313 | self.len = kp2d.shape[0] |
| 314 | |
| 315 | if self.image_array is None: |
| 316 | self.len_frame = len(self.frame_list) |
| 317 | else: |
| 318 | self.len_frame = self.image_array.shape[0] |
| 319 | |
| 320 | def __getitem__(self, frame_index): |
| 321 | """Get frame data from frame_list of image_array.""" |
| 322 | # frame file exists, resolution not set |
| 323 | if frame_index < self.len_frame and self.resolution is None: |
| 324 | if self.image_array is not None: |
| 325 | canvas = self.image_array[frame_index] |
| 326 | else: |
| 327 | canvas = cv2.imread(self.frame_list[frame_index]) |
| 328 | if self.kp2d is None: |
| 329 | kp2d_frame = None |
| 330 | else: |
| 331 | kp2d_frame = self.kp2d[frame_index] |
| 332 | # no frame file, resolution has been set |
| 333 | elif frame_index >= self.len_frame and self.resolution is not None: |
| 334 | canvas = np.ones((self.resolution[0], self.resolution[1], 3), |
| 335 | dtype=np.uint8) * 255 |