diff --git a/api/models.py b/api/models.py index 12afb59..b39b176 100644 --- a/api/models.py +++ b/api/models.py @@ -18,9 +18,9 @@ class Comic(models.Model): @classmethod def create_from_nparray(cls, nparray_file, video): + if nparray_file.max() <= 1: + nparray_file = (nparray_file * 255).astype(int) tmp_name = uuid.uuid4() - if not os.path.exists(f"{settings.TMP_DIR}"): - os.makedirs(f"{settings.TMP_DIR}") cv2.imwrite(f"{settings.TMP_DIR}{tmp_name}.png", nparray_file) with open(f"{settings.TMP_DIR}{tmp_name}.png", mode="rb") as tmp_file: comic_image = File(tmp_file, name=f"{tmp_name}.png") diff --git a/comic_layout/comic_layout.py b/comic_layout/comic_layout.py index 9a9842e..af478a1 100644 --- a/comic_layout/comic_layout.py +++ b/comic_layout/comic_layout.py @@ -14,10 +14,10 @@ class LayoutGenerator(): second_row = cv2.resize(second_row, (first_row.shape[1], - int(second_row.shape[0] * first_row.shape[1] / second_row.shape[1]))) + (second_row.shape[0] * first_row.shape[1]) // second_row.shape[1])) fourth_row = cv2.resize(fourth_row, (first_row.shape[1], - int(fourth_row.shape[0] * first_row.shape[1] / fourth_row.shape[1]))) + (fourth_row.shape[0] * first_row.shape[1]) // fourth_row.shape[1])) return np.vstack([first_row, second_row, third_row, fourth_row]) @@ -25,6 +25,6 @@ class LayoutGenerator(): def _pad_images(frames): padded_result_imgs = [] for img in frames: - padded_img = cv2.copyMakeBorder(img, 10, 10, 10, 10, cv2.BORDER_CONSTANT, value=(255, 255, 255)) + padded_img = cv2.copyMakeBorder(img, 5, 5, 5, 5, cv2.BORDER_CONSTANT, value=(1, 1, 1)) padded_result_imgs.append(padded_img) return padded_result_imgs diff --git a/keyframes/keyframes.py b/keyframes/keyframes.py index d12695b..99830f9 100644 --- a/keyframes/keyframes.py +++ b/keyframes/keyframes.py @@ -39,10 +39,9 @@ class KeyFramesExtractor(): @staticmethod def _get_frames_with_interval(interval, all_keyframes): chosen_frames = [] - chosen_frames_tmp_dir = uuid.uuid4() - os.mkdir(jj(f"{settings.TMP_DIR}", f"{chosen_frames_tmp_dir}")) for i in range(settings.NUMBERS_OF_FRAMES_TO_SHOW): frame = cv2.imread(all_keyframes[(i + 1) * interval]) chosen_frames.append(frame) + return chosen_frames diff --git a/style_transfer/style_transfer.py b/style_transfer/style_transfer.py index 6c30aff..c0b2fc9 100644 --- a/style_transfer/style_transfer.py +++ b/style_transfer/style_transfer.py @@ -18,7 +18,7 @@ class StyleTransfer(): @staticmethod def _cartoon_gan_stylize(frames, gpu=True, **kwargs): - style = kwargs.get("style", "Hosoda") + style = kwargs.get("style", "Hayao") resize = kwargs.get("resize", 450) # load pretrained model @@ -38,10 +38,7 @@ class StyleTransfer(): else: w = resize h = int(w * ratio) - input_image = cv2.resize(img, (w, h), interpolation=cv2.INTER_CUBIC) - - # RGB -> BGR - input_image = input_image[:, :, [2, 1, 0]] + input_image = cv2.resize(img, (w, h)) input_image = transforms.ToTensor()(input_image).unsqueeze(0) # preprocess, (-1, 1) @@ -52,13 +49,13 @@ class StyleTransfer(): output_image = model(input_image) output_image = output_image[0] - # BGR -> RGB - output_image = output_image[[2, 1, 0], :, :] - # deprocess, (0, 1) - output_image = output_image.data.cpu().float() * 0.5 + 0.5 + output_image = (output_image.data.cpu().float() * 0.5 + 0.5).numpy() - # switch channels and append image to result images - stylized_imgs.append(np.rollaxis(output_image.numpy(), 0, 3)) + # switch channels -> (c, h, w) -> (h, w, c) + output_image = np.rollaxis(output_image, 0, 3) + + # append image to result images + stylized_imgs.append(output_image) return stylized_imgs