CofeehousePy/services/nsfw_detection/coffeehouse_nsfw/detector_utils.py

62 lines
1.4 KiB
Python

import os
import cv2
import numpy as np
from PIL import Image
def read_image_bgr(path):
""" Read an image in BGR format.
Args
path: Path to the image.
"""
if isinstance(path, str):
image = np.ascontiguousarray(Image.open(path).convert("RGB"))
else:
path = cv2.cvtColor(path, cv2.COLOR_BGR2RGB)
image = np.ascontiguousarray(Image.fromarray(path))
return image[:, :, ::-1]
def _preprocess_image(x, mode="caffe"):
x = x.astype(np.float32)
if mode == "tf":
x /= 127.5
x -= 1.0
elif mode == "caffe":
x -= [103.939, 116.779, 123.68]
return x
def compute_resize_scale(image_shape, min_side=800, max_side=1333):
(rows, cols, _) = image_shape
smallest_side = min(rows, cols)
scale = min_side / smallest_side
largest_side = max(rows, cols)
if largest_side * scale > max_side:
scale = max_side / largest_side
return scale
def resize_image(img, min_side=800, max_side=1333):
scale = compute_resize_scale(img.shape, min_side=min_side, max_side=max_side)
img = cv2.resize(img, None, fx=scale, fy=scale)
return img, scale
def preprocess_image(
image_path, min_side=800, max_side=1333,
):
image = read_image_bgr(image_path)
image = _preprocess_image(image)
image, scale = resize_image(image, min_side=min_side, max_side=max_side)
return image, scale