import cv2
import numpy as np


def preprocess(img_path, with_meta=False, input_size=(640, 640)):
    img = cv2.imread(img_path)
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    h, w = img.shape[:2]
    img = cv2.resize(img, input_size[::-1], interpolation=1)
    img = img.astype(np.float32) / 255.
    mean = np.array([0.485, 0.456, 0.406], dtype=np.float32)
    std = np.array([0.229, 0.224, 0.225], dtype=np.float32)
    img = (img - mean) / std
    img = np.expand_dims(img, 0)
    scale_factor = np.array([[input_size[0] / h, input_size[1] / w]], dtype=np.float32)
    return img, scale_factor


def postprocess(pred, img_meta=None, conf=0.001):
    # pred: [N, 6], [label, score, x1, y1, x2, y2]
    # labels: 0.: 吸烟， 1.: 其他
    labels, scores, boxes = pred[:, 0].astype(np.int64), pred[:, 1], pred[:, 2:6]
    return boxes, scores, labels




