Update app.py

app.py

@@ -1,111 +1,111 @@
-import streamlit as st
-import numpy as np
-import torch
-import cv2
-from tensorflow.keras.models import load_model
-from tensorflow.keras.preprocessing import image as keras_image
-from PIL import Image
-
-# Load the pre-trained MobileNetV2 model from the .h5 file
-@st.cache_resource
-def load_mobilenet_model():
-    return load_model('mobilenetv2_classifier.h5')  # Replace with your .h5 file path
-
-mobilenet_model = load_mobilenet_model()
-
-# Load the YOLOv5 model (pre-trained)
-@st.cache_resource
-def load_yolov5_model():
-    return torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True)
-
-yolov5_model = load_yolov5_model()
-
-# Function to preprocess and predict image using MobileNetV2 with class messages
-def predict_image(model, img, target_size=(224, 224)):
-    st.image(img, caption="Uploaded Image", use_column_width=True)
-    img = img.resize(target_size)
-    img_array = keras_image.img_to_array(img)
-    img_array = np.expand_dims(img_array, axis=0)
-    img_array /= 255.0  # Normalize to match the training preprocessing
-    predictions = model.predict(img_array)[0]
-
-    # Get indices sorted by prediction confidence (in descending order)
-    sorted_indices = np.argsort(predictions)[::-1]
-    primary_class_index = sorted_indices[0]
-    primary_confidence = predictions[primary_class_index]
-
-    # Check if the primary predicted label is 1 and has a confidence below 0.7
-    if primary_class_index != 1 and primary_confidence < 0.7:
-        secondary_class_index = sorted_indices[1]
-        secondary_confidence = predictions[secondary_class_index]
-        return secondary_class_index, secondary_confidence
-
-    return primary_class_index, primary_confidence
-
-# Map class indices to specific messages
-class_messages = {
-    0: "Face is not visible",
-    1: "Looking at the system",
-    2: "Looking down",
-    3: "Looking left or right",
-    4: "Loking up"
-}
-
-# Function to count people using YOLOv5
-def count_people(img):
-    img_array = np.array(img)
-    results = yolov5_model(img_array)
-    detections = results.pandas().xyxy[0]
-    person_detections = detections[detections['name'] == 'person']
-    num_people = len(person_detections)
-    return num_people, results
-
-# Function to detect communication devices using YOLOv5
-def detect_device(img):
-    img_array = np.array(img)
-    results = yolov5_model(img_array)
-    detections = results.pandas().xyxy[0]
-    device_detections = detections[detections['name'].str.contains('phone|tablet|computer', case=False, na=False)]
-    return device_detections, results
-
-# Streamlit UI setup
-st.title("Image Analysis App")
-
-# Upload image
-uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
-
-if uploaded_file is not None:
-    img = Image.open(uploaded_file)
-    
-    # Display the uploaded image
-    st.image(img, caption="Uploaded Image", use_column_width=True)
-    
-    # Run classification with MobileNetV2
-    st.subheader("Image Classification")
-    predicted_index, confidence = predict_image(mobilenet_model, img)
-    
-    if predicted_index in class_messages:
-        st.write(f"Prediction: {class_messages[predicted_index]}")
-        st.write(f"Confidence score: {confidence:.2f}")
-    else:
-        st.write("Prediction not recognized.")
-    
-    # Run person counting with YOLOv5
-    st.subheader("People Counting")
-    num_people, people_results = count_people(img)
-    st.write(f"Number of people detected: {num_people}")
-    people_results.show()  # Show image with bounding boxes
-
-    # Run device detection with YOLOv5
-    st.subheader("Communication Device Detection")
-    device_detections, device_results = detect_device(img)
-    
-    if not device_detections.empty:
-        st.write("Communication devices detected:")
-        st.write(device_detections[['name', 'confidence']])
-        device_results.show()  # Show image with bounding boxes
-    else:
-        st.write("No communication devices detected in the image.")
-else:
-    st.write("Please upload an image to analyze.")
-
+import streamlit as st
+import numpy as np
+import torch
+import cv2
+from tensorflow.keras.models import load_model
+from tensorflow.keras.preprocessing import image as keras_image
+from PIL import Image
+
+# Load the pre-trained MobileNetV2 model from the .h5 file
+@st.cache_resource
+def load_mobilenet_model():
+    return load_model('mobilenetv2_classifier.h5')  # Replace with your .h5 file path
+
+mobilenet_model = load_mobilenet_model()
+
+# Load the YOLOv5 model (pre-trained)
+@st.cache_resource
+def load_yolov5_model():
+    return torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True)
+
+yolov5_model = load_yolov5_model()
+
+# Function to preprocess and predict image using MobileNetV2 with class messages
+def predict_image(model, img, target_size=(224, 224)):
+    st.image(img, caption="Uploaded Image", use_column_width=True)
+    img = img.resize(target_size)
+    img_array = keras_image.img_to_array(img)
+    img_array = np.expand_dims(img_array, axis=0)
+    img_array /= 255.0  # Normalize to match the training preprocessing
+    predictions = model.predict(img_array)[0]
+
+    # Get indices sorted by prediction confidence (in descending order)
+    sorted_indices = np.argsort(predictions)[::-1]
+    primary_class_index = sorted_indices[0]
+    primary_confidence = predictions[primary_class_index]
+
+    # Check if the primary predicted label is 1 and has a confidence below 0.7
+    if primary_class_index != 1 and primary_confidence < 0.7:
+        secondary_class_index = sorted_indices[1]
+        secondary_confidence = predictions[secondary_class_index]
+        return secondary_class_index, secondary_confidence
+
+    return primary_class_index, primary_confidence
+
+# Map class indices to specific messages
+class_messages = {
+    0: "Face is not visible",
+    1: "Looking at the system",
+    2: "Looking down",
+    3: "Looking left or right",
+    4: "Looking up"
+}
+
+# Function to count people using YOLOv5
+def count_people(img):
+    img_array = np.array(img)
+    results = yolov5_model(img_array)
+    detections = results.pandas().xyxy[0]
+    person_detections = detections[detections['name'] == 'person']
+    num_people = len(person_detections)
+    return num_people, results
+
+# Function to detect communication devices using YOLOv5
+def detect_device(img):
+    img_array = np.array(img)
+    results = yolov5_model(img_array)
+    detections = results.pandas().xyxy[0]
+    device_detections = detections[detections['name'].str.contains('phone|tablet|computer', case=False, na=False)]
+    return device_detections, results
+
+# Streamlit UI setup
+st.title("Image Analysis App")
+
+# Upload image
+uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
+
+if uploaded_file is not None:
+    img = Image.open(uploaded_file)
+    
+    # Display the uploaded image
+    st.image(img, caption="Uploaded Image", use_column_width=True)
+    
+    # Run classification with MobileNetV2
+    st.subheader("Image Classification")
+    predicted_index, confidence = predict_image(mobilenet_model, img)
+    
+    if predicted_index in class_messages:
+        st.write(f"Prediction: {class_messages[predicted_index]}")
+        st.write(f"Confidence score: {confidence:.2f}")
+    else:
+        st.write("Prediction not recognized.")
+    
+    # Run person counting with YOLOv5
+    st.subheader("People Counting")
+    num_people, people_results = count_people(img)
+    st.write(f"Number of people detected: {num_people}")
+    people_results.show()  # Show image with bounding boxes
+
+    # Run device detection with YOLOv5
+    st.subheader("Communication Device Detection")
+    device_detections, device_results = detect_device(img)
+    
+    if not device_detections.empty:
+        st.write("Communication devices detected:")
+        st.write(device_detections[['name', 'confidence']])
+        device_results.show()  # Show image with bounding boxes
+    else:
+        st.write("No communication devices detected in the image.")
+else:
+    st.write("Please upload an image to analyze.")
+