Spaces:
Sleeping
Sleeping
| """ | |
| IM 417 Classification | |
| """ | |
| import cv2 | |
| import numpy as np | |
| import os | |
| import streamlit as st | |
| import tempfile | |
| from streamlit_drawable_canvas import st_canvas | |
| from tensorflow import keras | |
| from huggingface_hub import from_pretrained_keras | |
| from transformers import AutoModel | |
| # Load the model (using a Streamlit caching mechanism for efficiency) | |
| #@st.cache(allow_output_mutation=True) # Allow model mutation for prediction | |
| def load_model(): | |
| try: | |
| model = from_pretrained_keras("tomeheya/IM-417-128x128-classification", force_download=True) | |
| #model = AutoModel.from_pretrained("tomeheya/IM-417-128x128-classification") | |
| print("Model loaded successfully!") | |
| return model | |
| except Exception as e: | |
| print(f"Error loading model: {e}") | |
| return None | |
| # Function to preprocess the input image | |
| def preprocess_image(image): | |
| image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) | |
| image = cv2.bitwise_not(image) | |
| image = np.expand_dims(image, axis=-1) # Add channel dimension | |
| image = image / 255.0 # Normalize pixel values | |
| return image | |
| # Function to perform inference | |
| def predict(image): | |
| model = load_model() | |
| if not model: | |
| return "Model couldn't be loaded !!!" | |
| preprocessed_image = preprocess_image(image) | |
| preprocessed_image = np.expand_dims(preprocessed_image, axis=0) # Add batch dimension | |
| prediction = model.predict(preprocessed_image) | |
| predicted_label = decode_prediction(prediction) | |
| return predicted_label | |
| # Function to decode the model prediction | |
| def decode_prediction(prediction): | |
| # Decode the prediction based on your model's output format | |
| # For example, if your model outputs class probabilities: | |
| class_idx = np.argmax(prediction) | |
| predicted_label = "IM-417 sign number :: " + str(class_idx + 1) | |
| return predicted_label | |
| # Initialize the canvas | |
| canvas_result = st_canvas( | |
| fill_color="rgba(255, 165, 0, 0.3)", | |
| stroke_width=3, | |
| stroke_color="rgba(0, 0, 0, 1)", | |
| background_color="rgba(255, 255, 255, 1)", | |
| height=128, | |
| width=128, | |
| update_streamlit=True, | |
| key="canvas" | |
| ) | |
| # Add a submit button | |
| if st.button("Submit"): | |
| if canvas_result.json_data: | |
| image_from_canvas = canvas_result.image_data | |
| st.text(f"Original data: {len(image_from_canvas)}") | |
| preprocessed_image = preprocess_image(image_from_canvas) | |
| predicted_label = predict(preprocessed_image) | |
| st.text(f"Predicted label: {predicted_label}") | |
| else: | |
| st.text("Please draw an image on the canvas before submitting.") | |
| # Streamlit app layout | |
| st.title("IM 417 Classification") | |