app.py

import streamlit as st
import numpy as np
import tensorflow as tf
from huggingface_hub import from_pretrained_keras
from PIL import Image

burmese_snake_classifer_pretrained = from_pretrained_keras('jojo-ai-mst/burmese_snake_classifier')
snake_nonsnake_classifier_pretrained = from_pretrained_keras('jojo-ai-mst/snake_nonsnake_classifier')

class_names = ['Bungarus fasciatus (Banded Krait)', 'Chrysopelea ornata (Golden Tree Snake)', "Daboia siamensis (Eastern Russell's viper)", 'Fowlea piscator (Checkered Keelback)', 'Laticauda colubrina (Sea Snake)', 'Lycodon aulicus (Wolf Snake)', 'Naja kaouthia(Cobra)', 'Ophiophagus_hannah(King Cobra)', 'Rhadophis helleri (Heller Red necked keelback)', 'Trimeresurus_sp (Asian Palm Pit vipers)']

def is_snake(img_array):
  prediction = snake_nonsnake_classifier_pretrained.predict(img_array)
  score = prediction[0][0]
  if score < 0.2:
    return False
  return True

def classify_snake(img_array):
  predictions = burmese_snake_classifer_pretrained.predict(img_array)
  score = tf.nn.softmax(predictions[0])

  result =  "This image most likely belongs to {} with a {:.2f} percent confidence.".format(class_names[np.argmax(score)], 100 * np.max(score))

  return result

def predict_img(input_img):
  img_array = np.expand_dims(input_img, 0) 
  if is_snake(img_array):
    return classify_snake(img_array)
  else:
    return "This image is not containing snake or poorly containing snake that MMDeepSnake can't detect snake in the image"

st.title("MM DeepSnake")
st.subheader('MM DeepSnake AI classifies snakes species in Myanmar, currently only 10 species')

with st.sidebar:
  st.write("""
    At the moment, we support
    - Trimeresurus_sp (Asian Palm Pit vipers) - မြွေစိမ်းမြီးခြောက်
    - Rhadophis helleri (Heller Red necked keelback) - လည်ပင်းနီမြွေ
    - Lycodon aulicus (Wolf Snake) - မြွေဝံပုလွေ
    - Fowlea piscator (Checkered Keelback) - ရေမြွေဗျောက်မ
    - Daboia siamensis (Eastern Russell's viper) - မြွေပွေး
    - Chrysopelea ornata (Golden Tree Snake) - ထန်းမြွေ
    - Bungarus fasciatus (Banded Krait) - ငန်းတော်ကြား
    - Ophiophagus hannah(King Cobra) - တောကြီးမြွေဟောက်
    - Laticauda colubrina (Sea Snake) - ဂျက်မြွေ
    - Naja kaouthia (Cobra) - မြွေဟောက်
  """)

uploaded_file = st.file_uploader("Upload Snake Image", type=['png', 'jpg','jpeg'], accept_multiple_files=False,label_visibility="visible")

if uploaded_file is not None:
    original = Image.open(uploaded_file)
    st.image(original, use_column_width=True)

    original = original.resize((300,300))

    img= np.array(original.convert('RGB'))

    result = predict_img(img)
    
    st.write(result)