Untitled

 avatar
unknown
plain_text
9 months ago
1.8 kB
21
Indexable
import cv2
import numpy as np


def preprocess_yolov3(image_file, norm=True, use_simple_preprocessing=True, use_lance_preprocess=True):
    """
    Preprocess ảnh sao cho giống hệt pipeline letterbox_square + ToFloat(max_value=255.0).
    - Đọc ảnh, chuyển sang RGB
    - Resize theo tỉ lệ, pad để ra ảnh vuông 240x240 (padding đen)
    - Nếu norm=True => chia 255.0
    - Trả về ảnh (1, H, W, 3)
    """
    im = cv2.imread(image_file, cv2.IMREAD_UNCHANGED)
    if im is None:
        raise FileNotFoundError(f"No such image: {image_file}")

    # --- Chuyển sang RGB ---
    if im.ndim == 2:
        rgb = cv2.cvtColor(im, cv2.COLOR_GRAY2RGB)
    elif im.shape[2] == 3:
        rgb = cv2.cvtColor(im, cv2.COLOR_BGR2RGB)
    elif im.shape[2] == 4:
        rgb = cv2.cvtColor(im, cv2.COLOR_BGRA2RGB)
    else:
        raise ValueError(f"Unsupported image format: {im.shape}")

    # --- Luôn resize kiểu letterbox square (giống _letterbox_square) ---
    target_size = 240
    ih, iw = rgb.shape[:2]
    scale = target_size / max(ih, iw)
    nh, nw = int(round(ih * scale)), int(round(iw * scale))
    resized = cv2.resize(rgb, (nw, nh), interpolation=cv2.INTER_CUBIC)

    # pad đen ở giữa
    new_img = np.zeros((target_size, target_size, 3), dtype=np.float32)
    top = (target_size - nh) // 2
    left = (target_size - nw) // 2
    new_img[top:top + nh, left:left + nw] = resized.astype(np.float32)

    # --- Chuẩn hóa giá trị ---
    if norm:
        preprocessed = new_img / 255.0
    else:
        preprocessed = np.round(new_img).astype("uint8")

    # --- Giữ nguyên các tham số và return kiểu cũ ---
    return np.expand_dims(preprocessed, axis=0)


def image_preprocess(img_path, norm=True):
    return preprocess_yolov3(img_path, norm=norm)
Editor is loading...
Leave a Comment