search

thongtruongvietsol / yolov7_int8

0 stars 0 forks source link

hihi #2

Open thongtruongvietsol opened 5 months ago

thongtruongvietsol commented 5 months ago

import numpy as np import torch from PIL import Image from utils.datasets import letterbox from utils.plots import plot_one_box

def preprocess_image(img0: np.ndarray): """ Preprocess image according to YOLOv7 input requirements. Takes image in np.array format, resizes it to specific size using letterbox resize, converts color space from BGR (default in OpenCV) to RGB and changes data layout from HWC to CHW.

Parameters:
  img0 (np.ndarray): image for preprocessing
Returns:
  img (np.ndarray): image after preprocessing
  img0 (np.ndarray): original image
"""
# resize
img = letterbox(img0, auto=False)[0]

# Convert
img = img.transpose(2, 0, 1)
img = np.ascontiguousarray(img)
return img, img0

def prepare_input_tensor(image: np.ndarray): """ Converts preprocessed image to tensor format according to YOLOv7 input requirements. Takes image in np.array format with unit8 data in [0, 255] range and converts it to torch.Tensor object with float data in [0, 1] range

Parameters:
  image (np.ndarray): image for conversion to tensor
Returns:
  input_tensor (torch.Tensor): float tensor ready to use for YOLOv7 inference
"""
input_tensor = image.astype(np.float32)  # uint8 to fp16/32
input_tensor /= 255.0  # 0 - 255 to 0.0 - 1.0

if input_tensor.ndim == 3:
    input_tensor = np.expand_dims(input_tensor, 0)
return input_tensor

label names for visualization

DEFAULT_NAMES = [ "person", "bicycle", "car", "motorcycle", "airplane", "bus", "train", "truck", "boat", "traffic light", "fire hydrant", "stop sign", "parking meter", "bench", "bird", "cat", "dog", "horse", "sheep", "cow", "elephant", "bear", "zebra", "giraffe", "backpack", "umbrella", "handbag", "tie", "suitcase", "frisbee", "skis", "snowboard", "sports ball", "kite", "baseball bat", "baseball glove", "skateboard", "surfboard", "tennis racket", "bottle", "wine glass", "cup", "fork", "knife", "spoon", "bowl", "banana", "apple", "sandwich", "orange", "broccoli", "carrot", "hot dog", "pizza", "donut", "cake", "chair", "couch", "potted plant", "bed", "dining table", "toilet", "tv", "laptop", "mouse", "remote", "keyboard", "cell phone", "microwave", "oven", "toaster", "sink", "refrigerator", "book", "clock", "vase", "scissors", "teddy bear", "hair drier", "toothbrush", ]

colors for visualization

COLORS = {name: [np.random.randint(0, 255) for _ in range(3)] for i, name in enumerate(DEFAULT_NAMES)}