snake-detection / app.py
AmitPress
color change
04147aa
from ultralytics import YOLO
import streamlit as st
from PIL import Image
import openai
secret_key = st.secrets["openapi"]
openai.api_key = secret_key
model = YOLO('resource/twenty-eight-first.pt')
classes = ['Ahaetulla nasuta', 'Ahaetulla prasina', 'Boiga cynodon', 'Boiga multomaculata', 'Bungarus caeruleus', 'Bungarus fasciatus', 'Chrysopelea ornata', 'Coelognathus radiatus', 'Daboia russelii', 'Dendrelaphis pictus', 'Dendrelaphis tristis', 'Echis carinatus', 'Fowlea piscator', 'Hydrophis platurus', 'Indotyphlops braminus', 'Laticauda colubrina', 'Lycodon aulicus', 'Malayopython reticulatus', 'Naja kaouthia', 'Naja naja', 'Oreocryptophis porphyraceus', 'Psammodynastes pulverulentus', 'Ptyas korros', 'Ptyas mucosa', 'Python bivittatus', 'Rhabdophis subminiatus', 'Trimeresurus albolabris', 'Xenochrophis piscator']
col1, col2 = st.columns([2,1])
col1.title('Snake Detection System')
photo = col2.file_uploader('Upload an Image (JPEG)', type='jpg')
if photo:
col2.success('Successfully Uploaded')
col1.image(photo)
image = Image.open(photo)
try:
bx = model(image)[0].boxes[0]
col1.write(f"Result: {(classes[int(bx.cls[0])])}")
txt = classes[int(bx.cls[0])]
genus, species = txt.split(' ')
venomousity = openai.Completion.create(engine="text-davinci-003", prompt=f"Is {genus} {species} is venomous? if venomous just say venomous otherwise say nonvenomous", max_tokens=50)
desc = openai.Completion.create(engine="text-davinci-003", prompt=f"Give a short description of {genus} {species}", max_tokens=500)
snake_type = venomousity.choices[0].text
with col1:
if snake_type == 'Venomous':
st.markdown(f'<span style="color: red;">{snake_type}</span>', unsafe_allow_html=True)
else:
st.markdown(f'<span style="color: green;">{snake_type}</span>', unsafe_allow_html=True)
st.markdown(desc.choices[0].text)
except:
col1.write(f"Error: There might be some network issue")
else:
col1.write('Please Upload an Image of a Snake')
col1.write('Result will be shown as soon as the image is uploaded')