fox-robot / utils.py
pixelprotest's picture
Upload 3 files
a69b646 verified
import os
import cv2
import numpy as np
# import sys
import glob
import gradio as gr
import random
# import importlib.util
import datetime
# from tensorflow.lite.python.interpreter import Interpreter
# import matplotlib
import matplotlib.pyplot as plt
### ---------------------------- image utils ---------------------------------
def parse_image_for_detection(img):
"""
if img comes from gradio, it makes sure its a numpy array
if img is a file path, it reads it and converts from BGR to RGB
it also returns the width and height of the image
"""
if isinstance(img, str):
## if its a file path, we read it and convert from BGR to RGB
image = cv2.imread(img)
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
else:
## otherwise assume its a numpy array from Gradio UI.
## but make sure that it actually is.
if not isinstance(img, np.ndarray):
img = np.array(img)
image = img
## lets also get the width and height of the original image
image_height, image_width, _ = image.shape
return image, image_width, image_height
def resize_image(np_image, width, height):
image_resized = cv2.resize(np_image, (width, height))
np_image = np.expand_dims(image_resized, axis=0)
return np_image
def normalize_image(np_image, interpreter):
## check if the model expects a floating point input
is_model_float = (interpreter.get_input_details()[0]['dtype'] == np.float32)
## Normalize pixel values if using a floating model (i.e. if model is non-quantized)
if is_model_float:
input_mean = 256.0 / 2.0
input_std = 256.0 / 2.0
np_image = (np.float32(np_image) - input_mean) / input_std
return np_image
def save_image(image, output_dir, output_width=1600, output_height=1200, dpi=80):
"""
saves the image in the output dir, as a matplotlib figure
"""
## make sure output directory exists
os.makedirs(output_dir, exist_ok=True)
## first get the figsize in inches based on pixel output width, height
figsize = get_figsize_from_pixels(output_width, output_height, dpi=dpi)
## now plot the image
plt.figure(figsize=figsize)
plt.imshow(image)
plt.tight_layout(pad=3)
## generate an output filename with a timestamp
timestamp = datetime.datetime.now().strftime('%Y%m%d_%H%M%S_%f')
output_filename = os.path.join(output_dir, f'img_{timestamp}.png')
## save the figure with the output filename
plt.savefig(output_filename, dpi=dpi)
return output_filename
### ---------------------------- image utils ---------------------------------
### ---------------------------- basic utils ---------------------------------
def get_labels(labels_filepath):
with open(labels_filepath, 'r') as f:
labels = [line.strip() for line in f.readlines()]
return labels
def get_random_images(dirpath, image_count=10):
""" returns a list of random image filepaths from the dirpath """
images = glob.glob(dirpath + '/*.jpg') + \
glob.glob(dirpath + '/*.JPG') + \
glob.glob(dirpath + '/*.png') + \
glob.glob(dirpath + '/*.bmp')
# img_filepaths = random.sample(images, image_count)
# return img_filepaths
return sorted(images)
def get_figsize_from_pixels(width, height, dpi=80):
""" returns the width and height in inches based on the dpi
used for matplotlib figures
"""
width_in = width / dpi
height_in = height / dpi
return (width_in, height_in)
### ---------------------------- basic utils ---------------------------------