37 lines
867 B
Python
37 lines
867 B
Python
|
import cv2 as cv
|
||
|
|
import numpy as np
|
||
|
|
from skimage.io import imread
|
||
|
|
import os
|
||
|
|
|
||
|
|
|
||
|
|
def get_led_color(robot_image, model):
|
||
|
|
green_color = 1
|
||
|
|
red_color = 2
|
||
|
|
white_color = 3
|
||
|
|
arr_led = []
|
||
|
|
X_led = {}
|
||
|
|
|
||
|
|
resized_image = cv.resize(robot_image, (128, 128), interpolation=cv.INTER_AREA)
|
||
|
|
resized_image = resized_image / 255
|
||
|
|
|
||
|
|
arr_led.append(resized_image)
|
||
|
|
X_led["values"] = np.array(arr_led)
|
||
|
|
image_with_alpha_led = X_led["values"][:, :, :, :3]
|
||
|
|
pred = predict(image_with_alpha_led, model)
|
||
|
|
|
||
|
|
if pred[0] == green_color:
|
||
|
|
return "green"
|
||
|
|
elif pred[0] == red_color:
|
||
|
|
return "red"
|
||
|
|
elif pred[0] == white_color:
|
||
|
|
return "white"
|
||
|
|
else:
|
||
|
|
return "other"
|
||
|
|
|
||
|
|
|
||
|
|
def predict(image, model):
|
||
|
|
predict_x = model.predict(image)
|
||
|
|
classes_x = np.argmax(predict_x, axis=1)
|
||
|
|
print(f'{classes_x}\n')
|
||
|
|
return classes_x
|