JaredBailey/BirdOrSquirrelV1

1import streamlit as st2import torch 3import torchvision4from torchvision import transforms5from PIL import Image6 7#####8###9# Initialization10###11#####12if 'generate_result' not in st.session_state:13    st.session_state['generate_result'] = 014if 'show_result' not in st.session_state:15    st.session_state['show_result'] = 016if 'number_of_files' not in st.session_state:17    st.session_state['number_of_files'] = 018if 'upload_choice' not in st.session_state:19    st.session_state['upload_choice'] = 'file_up'20 21 22#####23###24# Used to show either the file_uploader or the webcam25###26#####27def change_state():28    if st.session_state['upload_choice'] == 'file_up':29        st.session_state['upload_choice'] = 'webcam'30    else:31        st.session_state['upload_choice'] = 'file_up'32 33# User toggle for file_uploader vs webcam34st.toggle(label="Webcam",  help="Click on to use webcam, off to upload a file", on_change=change_state)   35 36# Use state to know whether to show file_uploader or webcam37if st.session_state['upload_choice'] == 'file_up':38    img = st.file_uploader(label="Upload a photo of a squirrel or bird", type=['png', 'jpg'])39    if img is not None:40        st.session_state['number_of_files'] = 141    else:42        st.session_state['number_of_files'] = 043else:44    img = st.camera_input(label="Webcam")45    if img is not None:46        st.session_state['number_of_files'] = 147    else:48        st.session_state['number_of_files'] = 049 50 51#####52###53# Load the image and apply transformations54###55#####56def predict_image(image_path, model):57 58    image = Image.open(image_path).convert('RGB')59    transform = transforms.Compose([60        transforms.Resize((224, 224)),61        transforms.ToTensor(),62        transforms.Normalize(mean=[0.485, 0.456, 0.406],63                             std=[0.229, 0.224, 0.225])64    ])65    input_image = transform(image).unsqueeze(0)  # Add batch dimension66 67    # Move input tensor to the device (GPU if available)68    input_image = input_image.to('cpu')69 70    # Perform inference71    model.eval()72    with torch.no_grad():73        output = model(input_image)74 75    # Get predicted class probabilities and class index76    probabilities = torch.softmax(output, dim=1)[0]77    predicted_class_index = torch.argmax(probabilities).item()78 79    # Map class index to class label80    class_labels = ["Bird", "Squirrel"]81    predicted_class_label = class_labels[predicted_class_index]82 83    return predicted_class_label84#     print("Class probabilities:")85#     for i, prob in enumerate(probabilities):86#         print(f"{class_labels[i]}: {prob:.4f}")87 88 89#####90###91# Load model and prepare for inference92###93#####94model_loaded = torchvision.models.resnet18(pretrained=False)  # Initialize ResNet18 without pretraining 95model_loaded.fc = torch.nn.Linear(model_loaded.fc.in_features, 2)  # Modify the fully connected layer96model_loaded = model_loaded.to('cpu')  # Move the model to the appropriate device (GPU or CPU)97 98# Load the saved state dictionary into the model99model_path = 'resnet18_custom_model.pth'100model_loaded.load_state_dict(torch.load(model_path, map_location='cpu'))101 102# Set the model to evaluation mode103model_loaded.eval()104 105 106#####107###108# Toggle view of model output in UI109###110#####111if st.session_state['upload_choice'] == 'file_up' and st.session_state['number_of_files'] == 1:112    st.session_state['generate_result'] = 1113    st.session_state['show_result'] = 1114elif st.session_state['upload_choice'] == 'webcam' and st.session_state['number_of_files'] == 1:115    st.session_state['generate_result'] = 1116    st.session_state['show_result'] = 1   117else:118    st.session_state['generate_result'] = 0119    st.session_state['show_result'] = 0120 121 122 123if st.session_state['generate_result'] != 0:124    if img is not None:125        result = predict_image(image_path=img, model=model_loaded)126    st.session_state['generate_result'] = 0127 128if st.session_state['show_result'] != 0:129    if result == 'Bird':130        st.markdown("""131            <style>132            .centered {133                text-align: center;134            }135            </style>136            <div class="centered">137                🐦🐦🐦🐦🐦🐦🐦🐦🐦🐦🐦🐦🐦🐦🐦🐦138            </div>139            """, unsafe_allow_html=True)140        st.markdown("""141            <style>142            .big-font {143                font-size:30px !important;144                text-align: center;145            }146            </style>147            <div class="big-font">148            That's a Bird149            </div>150        """, unsafe_allow_html=True)151        st.markdown("""152            <style>153            .centered {154                text-align: center;155            }156            </style>157            <div class="centered">158                🐦🐦🐦🐦🐦🐦🐦🐦🐦🐦🐦🐦🐦🐦🐦🐦159            </div>160            """, unsafe_allow_html=True)161        if st.session_state['upload_choice'] == 'file_up':162            st.image(img)163    else:164        st.markdown("""165            <style>166            .centered {167                text-align: center;168            }169            </style>170            <div class="centered">171                🐿️🐿️🐿️🐿️🐿️🐿️🐿️🐿️🐿️🐿️🐿️🐿️🐿️🐿️🐿️🐿️172            </div>173            """, unsafe_allow_html=True)174        st.markdown("""175            <style>176            .big-font {177                font-size:30px !important;178                text-align: center;179            }180            </style>181            <div class="big-font">182            That's a Squirrel183            </div>184        """, unsafe_allow_html=True)185        st.markdown("""186            <style>187            .centered {188                text-align: center;189            }190            </style>191            <div class="centered">192                🐿️🐿️🐿️🐿️🐿️🐿️🐿️🐿️🐿️🐿️🐿️🐿️🐿️🐿️🐿️🐿️193            </div>194            """, unsafe_allow_html=True)195        if st.session_state['upload_choice'] == 'file_up':196            st.image(img)