SURESHBEEKHANI commited on
Commit
7b90fcd
Β·
verified Β·
1 Parent(s): af7bbd2

Upload 4 files

Browse files
Files changed (4) hide show
  1. app.py +123 -0
  2. logo/logo.png +0 -0
  3. model/best.pt +3 -0
  4. requirements.txt +3 -0
app.py ADDED
@@ -0,0 +1,123 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import streamlit as st
2
+ from ultralytics import YOLO
3
+ from PIL import Image
4
+ import torchvision.transforms as transforms
5
+ import base64
6
+ import cv2
7
+ import numpy as np
8
+
9
+ # Set Streamlit Page Configuration
10
+ st.set_page_config(
11
+ page_title="PPE Detect",
12
+ page_icon="logo/logo.png",
13
+ layout="centered"
14
+ )
15
+
16
+ # Cache the YOLO model to optimize performance
17
+ @st.cache_resource()
18
+ def load_model():
19
+ return YOLO("model/best.pt") # Ensure correct model path
20
+
21
+ model = load_model()
22
+
23
+ # Define image transformation pipeline
24
+ transform = transforms.Compose([
25
+ transforms.Resize((640, 640)),
26
+ transforms.ToTensor()
27
+ ])
28
+
29
+ # Function to perform PPE detection on images
30
+ def predict_ppe(image: Image.Image):
31
+ try:
32
+ image_tensor = transform(image).unsqueeze(0) # Add batch dimension
33
+ results = model.predict(image_tensor)
34
+ output_image = results[0].plot() # Overlay predictions
35
+ return Image.fromarray(output_image)
36
+ except Exception as e:
37
+ st.error(f"Prediction Error: {e}")
38
+ return None
39
+
40
+ # Function to encode image to base64 for embedding
41
+ def get_base64_image(image_path):
42
+ try:
43
+ with open(image_path, "rb") as img_file:
44
+ return base64.b64encode(img_file.read()).decode()
45
+ except FileNotFoundError:
46
+ return None
47
+
48
+ # Function for real-time PPE detection using webcam
49
+ def live_ppe_detection():
50
+ st.sidebar.write("Starting live detection...")
51
+ cap = cv2.VideoCapture(0)
52
+ if not cap.isOpened():
53
+ st.sidebar.error("Error: Could not open webcam.")
54
+ return
55
+
56
+ stframe = st.empty()
57
+ stop_button = st.sidebar.button("Stop Live Detection", key="stop_button")
58
+
59
+ while cap.isOpened():
60
+ ret, frame = cap.read()
61
+ if not ret:
62
+ st.sidebar.error("Failed to capture video frame.")
63
+ break
64
+
65
+ results = model.predict(frame)
66
+ output_frame = results[0].plot()
67
+ stframe.image(output_frame, channels="BGR")
68
+
69
+ if stop_button:
70
+ break
71
+
72
+ cap.release()
73
+ cv2.destroyAllWindows()
74
+
75
+ # Display logo
76
+ image_base64 = get_base64_image("logo/logo.png")
77
+ if image_base64:
78
+ st.markdown(
79
+ f'<div style="text-align: center;"><img src="data:image/png;base64,{image_base64}" width="100"></div>',
80
+ unsafe_allow_html=True
81
+ )
82
+
83
+ # UI Customization
84
+ st.markdown("""
85
+ <style>
86
+ [data-testid="stSidebar"] { background-color: #1E1E2F; }
87
+ [data-testid="stSidebar"] h1, [data-testid="stSidebar"] h2 { color: white; }
88
+ h1 { text-align: center; font-size: 36px; font-weight: bold; color: #2C3E50; }
89
+ div.stButton > button { background-color: #3498DB; color: white; font-weight: bold; }
90
+ div.stButton > button:hover { background-color: #2980B9; }
91
+ </style>
92
+ """, unsafe_allow_html=True)
93
+
94
+ # Sidebar - File Upload
95
+ st.sidebar.header("πŸ“€ Upload an Image")
96
+ uploaded_file = st.sidebar.file_uploader("Drag and drop or browse", type=['jpg', 'png', 'jpeg'])
97
+
98
+ # Sidebar - Live Predictions
99
+ st.sidebar.header("πŸ“‘ Live Predictions")
100
+ if st.sidebar.button("Start Live Detection", key="start_button"):
101
+ live_ppe_detection()
102
+
103
+ # Main Page
104
+ st.title("PPE Detect")
105
+ st.markdown("<p style='text-align: center;'>Detect personal protective equipment (PPE) in images.</p>", unsafe_allow_html=True)
106
+
107
+ if uploaded_file:
108
+ image = Image.open(uploaded_file).convert("RGB")
109
+ col1, col2 = st.columns(2)
110
+
111
+ with col1:
112
+ st.image(image, caption="πŸ“· Uploaded Image", use_container_width=True)
113
+
114
+ if st.sidebar.button("πŸ” Predict PPE", key="predict_button"):
115
+ detected_image = predict_ppe(image)
116
+ if detected_image:
117
+ with col2:
118
+ st.image(detected_image, caption="🎯 PPE Detection Result", use_container_width=True)
119
+ else:
120
+ st.error("Detection failed. Please try again.")
121
+
122
+
123
+ st.info("This app uses **YOLO** for PPE detection. Upload an image or start live detection to get started.")
logo/logo.png ADDED
model/best.pt ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:55e526d2cd1861601f0de8660177fe4393f6d773b83e6aedeec4f310a1f080e8
3
+ size 5473235
requirements.txt ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ ultralytics
2
+ Pillow
3
+ opencv-python