utils/preprocess.py

import cv2
import numpy as np


def convert_to_png(image_path):
    image = cv2.imread(str(image_path))
    png_path = str(image_path).rsplit('.', 1)[0] + '.png'
    cv2.imwrite(png_path, image)
    return png_path


OUTPUT_SIZE = 256
edge_threshold = 30


def fill_background_and_crop(image_path, edge_threshold=30):
    image = cv2.imread(str(image_path))
    image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

    # 縦と横の小さい方が256になるようにリサイズ
    height, width = image_rgb.shape[:2]
    scale = OUTPUT_SIZE / min(height, width)
    new_height, new_width = int(height * scale), int(width * scale)
    resized_image = cv2.resize(
        image_rgb, (new_width, new_height), interpolation=cv2.INTER_AREA)

    # リサイズ後の画像から中央を256x256で切り取る
    center_x, center_y = new_width // 2, new_height // 2
    half_size = OUTPUT_SIZE // 2
    cropped_image = resized_image[
        max(0, center_y - half_size):min(center_y + half_size, new_height),
        max(0, center_x - half_size):min(center_x + half_size, new_width)
    ]

    cv2.imwrite(str(image_path), cv2.cvtColor(
        cropped_image, cv2.COLOR_RGB2BGR))