def _update_frame_preview(self, new_chunk_file_path, frame_num, part=1, final=False): num = self.frames.index(frame_num) preview_task_file_path = self._get_preview_task_file_path(num) with handle_image_error(logger), \ handle_none(load_as_pil(new_chunk_file_path), raise_if_none=IOError("load_as_pil failed")) as img: def resize_and_save(img): img_x, img_y = img.size with img.resize((int(round(self.scale_factor * img_x)), int(round(self.scale_factor * img_y))), resample=Image.BILINEAR) as img_resized: img_resized.save(self._get_preview_file_path(num), PREVIEW_EXT) img_resized.save(preview_task_file_path, PREVIEW_EXT) if not final: with self._paste_new_chunk( img, self._get_preview_file_path(num), part, int(self.total_tasks / len(self.frames))) as img_pasted: resize_and_save(img_pasted) else: resize_and_save(img) self.last_preview_path = preview_task_file_path
def _remove_from_preview(self, subtask_id): subtask = self.subtasks_given[subtask_id] empty_color = (0, 0, 0) with handle_image_error(logger), \ self._open_preview() as img: self._mark_task_area(subtask, img, empty_color) img.save(self.preview_file_path, PREVIEW_EXT)
def test_save_image_io_error(self): logger = Mock() with handle_image_error(logger), \ Image.new('RGB', (1, 1)) as image, \ io.BytesIO() as b: b.write = Mock(side_effect=IOError("fail")) image.save(b, 'PNG') assert logger.error.called
def test_save_image_key_error(self): logger = Mock() with handle_image_error(logger), \ Image.new('RGB', (1, 1)) as image, \ io.BytesIO() as b: image.save(b, 'UNKNOWN EXTENSION') assert b.getvalue() == b'' assert logger.error.called
def test_save_image(self): logger = Mock() with handle_image_error(logger), \ Image.new('RGB', (1, 1)) as image, \ io.BytesIO() as b: image.save(b, 'PNG') assert b.getvalue() assert not logger.error.called
def _paste_image(self, final_img, new_part, num): with handle_image_error(logger), \ Image.new("RGB", (self.width, self.height)) as img_offset: offset = self.current_offset _, new_img_res_y = new_part.size self.current_offset += new_img_res_y img_offset.paste(new_part, (0, offset)) result = ImageChops.add(final_img, img_offset) return result
def restart(self): self.chunks = {} self.perfect_match_area_y = 0 self.perfectly_placed_subtasks = 0 if os.path.exists(self.preview_file_path): with handle_image_error(logger), \ Image.new("RGB", (self.preview_res_x, self.preview_res_y)) \ as img: img.save(self.preview_file_path, PREVIEW_EXT)
def verify_img(self, filename): with handle_image_error(logger): with Image.open(filename) as image: img_size = image.size expected = self._task_definition.resolution if tuple(img_size) == tuple(expected): return True logger.warning("Bad resolution\nExpected %sx%s, but got %sx%s", expected[0], expected[1], img_size[0], img_size[1]) return False
def _open_frame_preview(self, preview_file_path): if not os.path.exists(preview_file_path): with handle_image_error(logger), \ Image.new("RGB", (int(round(self.res_x * self.scale_factor)), int(round(self.res_y * self.scale_factor)))) \ as img: img.save(preview_file_path, PREVIEW_EXT) return Image.open(preview_file_path)
def _put_image_together(self): output_file_name = "{}".format(self.output_file, self.output_format) logger.debug('_put_image_together() out: %r', output_file_name) self.collected_file_names = OrderedDict( sorted(self.collected_file_names.items())) if not self._use_outer_task_collector(): collector = CustomCollector(width=self.res_x, height=self.res_y) for file in self.collected_file_names.values(): collector.add_img_file(file) with handle_image_error(logger), \ collector.finalize() as image: image.save_with_extension(output_file_name, self.output_format) else: self._put_collected_files_together( os.path.join(self.tmp_dir, output_file_name), list(self.collected_file_names.values()), "paste")
def update_preview(self, subtask_path, subtask_number): if subtask_number not in self.chunks: self.chunks[subtask_number] = subtask_path with handle_image_error(logger) as handler_result, \ handle_none(load_as_pil(subtask_path), raise_if_none=IOError("load_as_pil failed")) \ as subtask_img: offset = self.get_offset(subtask_number) if subtask_number == self.perfectly_placed_subtasks + 1: _, img_y = subtask_img.size self.perfect_match_area_y += img_y self.perfectly_placed_subtasks += 1 # this is the last task if subtask_number + 1 >= len(self.expected_offsets): height = self.preview_res_y - \ self.expected_offsets[subtask_number] else: height = self.expected_offsets[subtask_number + 1] - \ self.expected_offsets[subtask_number] with subtask_img.resize((self.preview_res_x, height), resample=Image.BILINEAR) \ as subtask_img_resized: def open_or_create_image(): if not os.path.exists(self.preview_file_path) \ or len(self.chunks) == 1: return Image.new( "RGB", (self.preview_res_x, self.preview_res_y)) return Image.open(self.preview_file_path) with open_or_create_image() as preview_img: preview_img.paste(subtask_img_resized, (0, offset)) preview_img.save(self.preview_file_path, PREVIEW_EXT) if not handler_result.success: return if subtask_number == self.perfectly_placed_subtasks and \ (subtask_number + 1) in self.chunks: self.update_preview(self.chunks[subtask_number + 1], subtask_number + 1)
def _open_preview(self, mode="RGB", ext=PREVIEW_EXT): """ If preview file doesn't exist create a new empty one with given mode and extension. Extension should be compatibile with selected mode. """ if self.preview_file_path is None or not os.path.exists( self.preview_file_path): preview_name = "current_preview.{}".format(ext) self.preview_file_path = "{}".format( os.path.join(self.tmp_dir, preview_name)) with handle_image_error(logger), \ Image.new(mode, (int(round(self.res_x * self.scale_factor)), int(round(self.res_y * self.scale_factor)))) \ as img: logger.debug('Saving new preview: %r', self.preview_file_path) img.save(self.preview_file_path, ext) logger.debug('Opening preview: %r, exists?: %s', self.preview_file_path, os.path.exists(self.preview_file_path)) return Image.open(self.preview_file_path)
def _put_frame_together(self, frame_num, num_start): directory = os.path.dirname(self.output_file) output_file_name = os.path.join(directory, self._get_output_name(frame_num)) frame_key = str(frame_num) collected = self.frames_given[frame_key] collected = OrderedDict(sorted(collected.items())) if not self._use_outer_task_collector(): collector = CustomCollector(width=self.res_x, height=self.res_y) for file in collected.values(): collector.add_img_file(file) with handle_image_error(logger), \ collector.finalize() as image: image.save_with_extension(output_file_name, self.output_format) else: self._put_collected_files_together(output_file_name, list(collected.values()), "paste") self.collected_file_names[frame_num] = output_file_name self._update_frame_preview(output_file_name, frame_num, final=True) self._update_frame_task_preview()
def _update_task_preview(self): sent_color = (0, 255, 0) failed_color = (255, 0, 0) preview_name = "current_task_preview.{}".format(PREVIEW_EXT) preview_task_file_path = "{}".format( os.path.join(self.tmp_dir, preview_name)) with handle_image_error(logger), \ self._open_preview() as img_task: subtasks_given = dict(self.subtasks_given) for sub in subtasks_given.values(): if sub['status'].is_active(): self._mark_task_area(sub, img_task, sent_color) if sub['status'] in [ SubtaskStatus.failure, SubtaskStatus.restarted ]: self._mark_task_area(sub, img_task, failed_color) img_task.save(preview_task_file_path, PREVIEW_EXT) self._update_preview_task_file_path(preview_task_file_path)
def _update_frame_preview(self, new_chunk_file_path, frame_num, part=1, final=False): num = self.frames.index(frame_num) if final: with handle_image_error(logger), \ handle_none(load_as_pil(new_chunk_file_path), raise_if_none=IOError("load_as_pil failed")) \ as img, \ img.resize((int(round(self.res_x * self.scale_factor)), int(round(self.res_y * self.scale_factor))), resample=Image.BILINEAR) as scaled: preview_task_file_path = self._get_preview_task_file_path(num) self.last_preview_path = preview_task_file_path scaled.save(preview_task_file_path, PREVIEW_EXT) scaled.save(self._get_preview_file_path(num), PREVIEW_EXT) else: self.preview_updaters[num].update_preview(new_chunk_file_path, part) self._update_frame_task_preview()