Spaces:
Sleeping
Sleeping
File size: 2,617 Bytes
347ff72 e262f0b 42d7ac3 dd0f7c7 3f1243a 347ff72 7723f84 6620c2e 347ff72 42d7ac3 18a25c4 972d723 158eed5 f1f8687 158eed5 329bdeb 158eed5 329bdeb 347ff72 158eed5 347ff72 279303c 9645511 279303c 49d4d61 e262f0b 49d4d61 279303c dd019a8 42d7ac3 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 |
"""
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")
|