Spaces:
Paused
Paused
| import cv2 | |
| import numpy as np | |
| import tensorflow as tf | |
| from PIL import Image | |
| import os | |
| TFLITE_PATH: str = os.path.join(os.path.dirname(os.path.realpath(__file__)), "models", "model_mobilenet_v2.tflite") | |
| IMAGE_SIZE: tuple[int, int] = (160, 160) | |
| CLASS_NAMES: list[str] = [ | |
| "A", "B", "C", "D", "E", | |
| "F", "G", "H", "I", "J", | |
| "K", "L", "M", "N", "O", | |
| "P", "Q", "R", "S", "T", | |
| "U", "V", "W", "X", "Y", | |
| "Z", "del", "space", | |
| ] | |
| TARGET_FRAME_COUNT: int = 3 | |
| TARGET_CONSECUTIVE_PREDICTIONS: int = 4 | |
| TARGET_PREDICTION_SCORE: float = 0.92 | |
| def load_model(): | |
| interpreter = tf.lite.Interpreter(model_path=TFLITE_PATH) | |
| # print(interpreter.get_signature_list()) | |
| classify_lite = interpreter.get_signature_runner("serving_default") | |
| return classify_lite | |
| def get_image_array(image_data): | |
| img_array = tf.keras.utils.img_to_array(image_data) | |
| img_array = tf.expand_dims(img_array, 0) # Create a batch | |
| return img_array | |
| def predict(classify_lite, image_array): | |
| score_lite = classify_lite(input_2=image_array)["outputs"] | |
| predicted_char = CLASS_NAMES[np.argmax(score_lite)] | |
| prediction_score = np.max(score_lite) | |
| return predicted_char, prediction_score | |
| def max_predicted(predictions: dict[str, int]) -> tuple[str, int]: | |
| return max(predictions.items(), key=lambda k: k[1]) | |
| # Load the model once globally to prevent massive memory leaks and CPU spikes | |
| GLOBAL_CLASSIFY_LITE = load_model() | |
| def sign_detection(img): | |
| classify_lite = GLOBAL_CLASSIFY_LITE | |
| x1, y1 = 100, 100 | |
| x2, y2 = (x1 + IMAGE_SIZE[0]), (y1 + IMAGE_SIZE[1]) | |
| previous_predictions: dict[str, int] = {letter: 0 for letter in CLASS_NAMES} | |
| text: str = "" | |
| img = cv2.flip(img, 1) | |
| predicted_char: str = "" | |
| prediction_score: float = 0.0 | |
| img_cropped = img[y1:y2, x1:x2] | |
| image_data = Image.fromarray(img_cropped) | |
| image_array = get_image_array(image_data) | |
| predicted_char, prediction_score = predict(classify_lite, image_array) | |
| if (prediction_score >= TARGET_PREDICTION_SCORE): | |
| previous_predictions[predicted_char] += 1 | |
| letter, count = max_predicted(previous_predictions) | |
| if (count >= TARGET_CONSECUTIVE_PREDICTIONS): | |
| previous_predictions = {letter: 0 for letter in CLASS_NAMES} | |
| if letter == "space": | |
| text += " " | |
| elif letter == "del": | |
| text = text[:-1] | |
| else: | |
| text += letter | |
| cv2.putText( | |
| img, | |
| predicted_char.upper(), | |
| (5, 100), | |
| cv2.FONT_HERSHEY_SIMPLEX, | |
| 4, | |
| (255, 255, 255), | |
| 4, | |
| ) | |
| cv2.putText( | |
| img, | |
| f"(score = {prediction_score:.2f})", | |
| (50, 70), | |
| cv2.FONT_HERSHEY_SIMPLEX, | |
| 1, | |
| (255, 255, 255), | |
| ) | |
| cv2.rectangle(img, (x1, y1), (x2, y2), (255, 0, 0), 2) | |
| # Removed cv2.imshow because it crashes on headless servers | |
| return img, predicted_char, str(prediction_score) | |