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()