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")