garbage2 / app.py
pranit144's picture
Update app.py
251c3b1 verified
import streamlit as st
import os
import tensorflow as tf
from tensorflow.keras.models import load_model
import cv2
import numpy as np
# Load the model manually using Keras
model = load_model('Model/keras_model.h5', compile=False) # Load model without compilation
# Labels directly added to the code
labels = [
"Nothing",
"Zip-top cans",
"Newspaper",
"Old shoes",
"Watercolor pen",
"Disinfectant",
"Battery",
"Vegetable leaf",
"Apple"
]
# Classification dictionary (modified to skip image-based overlays)
classDic = {0: None,
1: 0,
2: 0,
3: 3,
4: 3,
5: 1,
6: 1,
7: 2,
8: 2}
# Streamlit page configuration
st.title('Object Detection with Webcam')
st.write("This application uses your webcam to detect objects in real-time.")
# Start the webcam feed for Streamlit
cap = cv2.VideoCapture(0) # Manually set to camera index 0 (or change to 1, 2 if necessary)
if not cap.isOpened():
st.error("Error: No camera found!")
exit()
stframe = st.empty() # Placeholder for the webcam feed
# Add the 'Exit' button outside the loop so it's only displayed once
exit_button = st.button('Exit', key="exit_button")
# Loop to keep updating webcam feed
while True:
ret, img = cap.read()
if not ret or img is None:
st.error("Error: Failed to capture image!")
continue
imgResize = cv2.resize(img, (454, 340)) # Resize captured frame to a specific size
# Perform prediction using the loaded model
img_input = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) # Convert to RGB
img_input = cv2.resize(img_input, (224, 224)) # Resize to expected input size for your model
img_input = img_input / 255.0 # Normalize the image
prediction = model.predict(tf.convert_to_tensor(img_input[None, ...])) # Make prediction
classID = prediction.argmax() # Get the predicted class ID
print(f"Predicted class ID: {classID}")
# Handle the class ID (can be used to trigger different actions based on class)
if 0 < classID <= len(labels):
st.write(f"Detected class: {labels[classID - 1]}")
else:
st.write(f"Invalid classID: {classID}, setting to default.")
# Display the resized image in Streamlit
img_resized = cv2.cvtColor(imgResize, cv2.COLOR_BGR2RGB) # Convert image to RGB
stframe.image(img_resized, channels="RGB", caption="Real-time Detection", use_container_width=True)
# Exit condition based on button press
if exit_button:
break
# Release resources and close the window
cap.release()