Spaces:
Runtime error
Runtime error
Harsimran19
commited on
Commit
•
08d6e3c
1
Parent(s):
61b13b4
Upload 7 files
Browse files- app.py +74 -0
- config.py +43 -0
- requirements.txt +7 -0
- utils.py +168 -0
- weights/Detection/Fabric.pt +3 -0
- weights/Detection/Rust.pt +3 -0
- weights/Detection/Scratch.pt +3 -0
app.py
ADDED
@@ -0,0 +1,74 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
#!/usr/bin/env python
|
2 |
+
# -*- coding: utf-8 -*-
|
3 |
+
"""
|
4 |
+
-------------------------------------------------
|
5 |
+
@File Name: app.py
|
6 |
+
@Author: yash mohite
|
7 |
+
@Date: 2023/5/15
|
8 |
+
@Description:
|
9 |
+
-------------------------------------------------
|
10 |
+
"""
|
11 |
+
from pathlib import Path
|
12 |
+
from PIL import Image
|
13 |
+
import streamlit as st
|
14 |
+
|
15 |
+
import config
|
16 |
+
from utils import load_model, infer_uploaded_image, infer_uploaded_video, infer_uploaded_webcam
|
17 |
+
|
18 |
+
# setting page layout
|
19 |
+
st.set_page_config(
|
20 |
+
page_title="Defect Detection with YOLOv8",
|
21 |
+
page_icon="🤖",
|
22 |
+
layout="wide",
|
23 |
+
initial_sidebar_state="expanded"
|
24 |
+
)
|
25 |
+
|
26 |
+
# main page heading
|
27 |
+
st.title("Defect Detection with YOLOv8")
|
28 |
+
|
29 |
+
# sidebar
|
30 |
+
st.sidebar.header("DL Model Config")
|
31 |
+
|
32 |
+
# model options
|
33 |
+
task_type = st.sidebar.selectbox(
|
34 |
+
"Select Task",
|
35 |
+
["Rust","Scratch","Fabric"]
|
36 |
+
)
|
37 |
+
|
38 |
+
model_type = None
|
39 |
+
model_type = st.sidebar.selectbox(
|
40 |
+
"Select Model",
|
41 |
+
config.DETECTION_MODEL_LIST
|
42 |
+
)
|
43 |
+
|
44 |
+
confidence = float(st.sidebar.slider(
|
45 |
+
"Select Model Confidence", 30, 100, 50)) / 100
|
46 |
+
|
47 |
+
model_path = ""
|
48 |
+
if model_type:
|
49 |
+
model_path = Path(config.DETECTION_MODEL_DIR, str(model_type))
|
50 |
+
else:
|
51 |
+
st.error("Please Select Model in Sidebar")
|
52 |
+
|
53 |
+
# load pretrained DL model
|
54 |
+
try:
|
55 |
+
model = load_model(model_path)
|
56 |
+
except Exception as e:
|
57 |
+
st.error(f"Unable to load model. Please check the specified path: {model_path}")
|
58 |
+
|
59 |
+
# image/video options
|
60 |
+
st.sidebar.header("Image/Video Config")
|
61 |
+
source_selectbox = st.sidebar.selectbox(
|
62 |
+
"Select Source",
|
63 |
+
config.SOURCES_LIST
|
64 |
+
)
|
65 |
+
|
66 |
+
source_img = None
|
67 |
+
if source_selectbox == config.SOURCES_LIST[0]: # Image
|
68 |
+
infer_uploaded_image(confidence, model)
|
69 |
+
elif source_selectbox == config.SOURCES_LIST[1]: # Video
|
70 |
+
infer_uploaded_video(confidence, model)
|
71 |
+
elif source_selectbox == config.SOURCES_LIST[2]: # Webcam
|
72 |
+
infer_uploaded_webcam(confidence, model)
|
73 |
+
else:
|
74 |
+
st.error("Currently only 'Image' and 'Video' source are implemented")
|
config.py
ADDED
@@ -0,0 +1,43 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
from pathlib import Path
|
2 |
+
import sys
|
3 |
+
|
4 |
+
|
5 |
+
file_path = Path(__file__).resolve()
|
6 |
+
|
7 |
+
# Get the parent directory of the current file
|
8 |
+
root_path = file_path.parent
|
9 |
+
|
10 |
+
# Add the root path to the sys.path list if it is not already there
|
11 |
+
if root_path not in sys.path:
|
12 |
+
sys.path.append(str(root_path))
|
13 |
+
|
14 |
+
# Get the relative path of the root directory with respect to the current working directory
|
15 |
+
ROOT = root_path.relative_to(Path.cwd())
|
16 |
+
|
17 |
+
|
18 |
+
# Source
|
19 |
+
SOURCES_LIST = ["Image", "Video", "Webcam"]
|
20 |
+
|
21 |
+
|
22 |
+
# DL model config
|
23 |
+
DETECTION_MODEL_DIR = ROOT / 'weights'/'Detection'
|
24 |
+
#YOLOv8n = DETECTION_MODEL_DIR / "yolov8n.pt"
|
25 |
+
#YOLOv8s = DETECTION_MODEL_DIR / "yolov8s.pt"
|
26 |
+
#YOLOv8m = DETECTION_MODEL_DIR / "yolov8m.pt"
|
27 |
+
#YOLOv8l = DETECTION_MODEL_DIR / "yolov8l.pt"
|
28 |
+
#YOLOv8x = DETECTION_MODEL_DIR / "yolov8x.pt"
|
29 |
+
YOLOV8custom1 = DETECTION_MODEL_DIR /"Rust.pt"
|
30 |
+
YOLOV8custom2 = DETECTION_MODEL_DIR / "Scratch.pt"
|
31 |
+
YOLOV8custom3 = DETECTION_MODEL_DIR / "Fabric.pt"
|
32 |
+
|
33 |
+
# If You Want To Use This Pre Train Model Then UnComment It and uplode in your weights/detection dir and Use Them All
|
34 |
+
|
35 |
+
DETECTION_MODEL_LIST = [
|
36 |
+
#"yolov8n.pt",
|
37 |
+
#"yolov8s.pt",
|
38 |
+
#"yolov8m.pt",
|
39 |
+
#"yolov8l.pt",
|
40 |
+
#"yolov8x.pt",
|
41 |
+
"Rust.pt",
|
42 |
+
"Scratch.pt",
|
43 |
+
"Fabric.pt"]
|
requirements.txt
ADDED
@@ -0,0 +1,7 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
streamlit
|
2 |
+
ultralytics
|
3 |
+
pandas
|
4 |
+
numpy
|
5 |
+
opencv-python
|
6 |
+
# opencv-python-headless==4.7.0.72
|
7 |
+
Pillow
|
utils.py
ADDED
@@ -0,0 +1,168 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
#!/usr/bin/env python
|
2 |
+
# -*- coding: utf-8 -*-
|
3 |
+
"""
|
4 |
+
-------------------------------------------------
|
5 |
+
@File Name: utils.py
|
6 |
+
@Author: yash mohite
|
7 |
+
@Date: 2023/5/16
|
8 |
+
@Description:
|
9 |
+
-------------------------------------------------
|
10 |
+
"""
|
11 |
+
|
12 |
+
|
13 |
+
from ultralytics import YOLO
|
14 |
+
import streamlit as st
|
15 |
+
import cv2
|
16 |
+
from PIL import Image
|
17 |
+
import tempfile
|
18 |
+
|
19 |
+
|
20 |
+
def _display_detected_frames(conf, model, st_frame, image):
|
21 |
+
"""
|
22 |
+
Display the detected objects on a video frame using the YOLOv8 model.
|
23 |
+
:param conf (float): Confidence threshold for object detection.
|
24 |
+
:param model (YOLOv8): An instance of the `YOLOv8` class containing the YOLOv8 model.
|
25 |
+
:param st_frame (Streamlit object): A Streamlit object to display the detected video.
|
26 |
+
:param image (numpy array): A numpy array representing the video frame.
|
27 |
+
:return: None
|
28 |
+
"""
|
29 |
+
# Resize the image to a standard size
|
30 |
+
image = cv2.resize(image, (720, int(720 * (9 / 16))))
|
31 |
+
|
32 |
+
# Predict the objects in the image using YOLOv8 model
|
33 |
+
res = model.predict(image, conf=conf)
|
34 |
+
|
35 |
+
# Plot the detected objects on the video frame
|
36 |
+
res_plotted = res[0].plot()
|
37 |
+
st_frame.image(res_plotted,
|
38 |
+
caption='Detected Video',
|
39 |
+
channels="BGR",
|
40 |
+
use_column_width=True
|
41 |
+
)
|
42 |
+
|
43 |
+
|
44 |
+
@st.cache_resource
|
45 |
+
def load_model(model_path):
|
46 |
+
"""
|
47 |
+
Loads a YOLO object detection model from the specified model_path.
|
48 |
+
|
49 |
+
Parameters:
|
50 |
+
model_path (str): The path to the YOLO model file.
|
51 |
+
|
52 |
+
Returns:
|
53 |
+
A YOLO object detection model.
|
54 |
+
"""
|
55 |
+
model = YOLO(model_path)
|
56 |
+
return model
|
57 |
+
|
58 |
+
|
59 |
+
def infer_uploaded_image(conf, model):
|
60 |
+
"""
|
61 |
+
Execute inference for uploaded image
|
62 |
+
:param conf: Confidence of YOLOv8 model
|
63 |
+
:param model: An instance of the `YOLOv8` class containing the YOLOv8 model.
|
64 |
+
:return: None
|
65 |
+
"""
|
66 |
+
source_img = st.sidebar.file_uploader(
|
67 |
+
label="Choose an image...",
|
68 |
+
type=("jpg", "jpeg", "png", 'bmp', 'webp')
|
69 |
+
)
|
70 |
+
|
71 |
+
col1, col2 = st.columns(2)
|
72 |
+
|
73 |
+
with col1:
|
74 |
+
if source_img:
|
75 |
+
uploaded_image = Image.open(source_img)
|
76 |
+
# adding the uploaded image to the page with caption
|
77 |
+
st.image(
|
78 |
+
image=source_img,
|
79 |
+
caption="Uploaded Image",
|
80 |
+
use_column_width=True
|
81 |
+
)
|
82 |
+
|
83 |
+
if source_img:
|
84 |
+
if st.button("Execution"):
|
85 |
+
with st.spinner("Running..."):
|
86 |
+
res = model.predict(uploaded_image,
|
87 |
+
conf=conf)
|
88 |
+
boxes = res[0].boxes
|
89 |
+
res_plotted = res[0].plot()[:, :, ::-1]
|
90 |
+
|
91 |
+
with col2:
|
92 |
+
st.image(res_plotted,
|
93 |
+
caption="Detected Image",
|
94 |
+
use_column_width=True)
|
95 |
+
try:
|
96 |
+
with st.expander("Detection Results"):
|
97 |
+
for box in boxes:
|
98 |
+
st.write(box.xywh)
|
99 |
+
except Exception as ex:
|
100 |
+
st.write("No image is uploaded yet!")
|
101 |
+
st.write(ex)
|
102 |
+
|
103 |
+
|
104 |
+
def infer_uploaded_video(conf, model):
|
105 |
+
"""
|
106 |
+
Execute inference for uploaded video
|
107 |
+
:param conf: Confidence of YOLOv8 model
|
108 |
+
:param model: An instance of the `YOLOv8` class containing the YOLOv8 model.
|
109 |
+
:return: None
|
110 |
+
"""
|
111 |
+
source_video = st.sidebar.file_uploader(
|
112 |
+
label="Choose a video..."
|
113 |
+
)
|
114 |
+
|
115 |
+
if source_video:
|
116 |
+
st.video(source_video)
|
117 |
+
|
118 |
+
if source_video:
|
119 |
+
if st.button("Execution"):
|
120 |
+
with st.spinner("Running..."):
|
121 |
+
try:
|
122 |
+
tfile = tempfile.NamedTemporaryFile()
|
123 |
+
tfile.write(source_video.read())
|
124 |
+
vid_cap = cv2.VideoCapture(
|
125 |
+
tfile.name)
|
126 |
+
st_frame = st.empty()
|
127 |
+
while (vid_cap.isOpened()):
|
128 |
+
success, image = vid_cap.read()
|
129 |
+
if success:
|
130 |
+
_display_detected_frames(conf,
|
131 |
+
model,
|
132 |
+
st_frame,
|
133 |
+
image
|
134 |
+
)
|
135 |
+
else:
|
136 |
+
vid_cap.release()
|
137 |
+
break
|
138 |
+
except Exception as e:
|
139 |
+
st.error(f"Error loading video: {e}")
|
140 |
+
|
141 |
+
|
142 |
+
def infer_uploaded_webcam(conf, model):
|
143 |
+
"""
|
144 |
+
Execute inference for webcam.
|
145 |
+
:param conf: Confidence of YOLOv8 model
|
146 |
+
:param model: An instance of the `YOLOv8` class containing the YOLOv8 model.
|
147 |
+
:return: None
|
148 |
+
"""
|
149 |
+
try:
|
150 |
+
flag = st.button(
|
151 |
+
label="Stop running"
|
152 |
+
)
|
153 |
+
vid_cap = cv2.VideoCapture(0) # local camera
|
154 |
+
st_frame = st.empty()
|
155 |
+
while not flag:
|
156 |
+
success, image = vid_cap.read()
|
157 |
+
if success:
|
158 |
+
_display_detected_frames(
|
159 |
+
conf,
|
160 |
+
model,
|
161 |
+
st_frame,
|
162 |
+
image
|
163 |
+
)
|
164 |
+
else:
|
165 |
+
vid_cap.release()
|
166 |
+
break
|
167 |
+
except Exception as e:
|
168 |
+
st.error(f"Error loading video: {str(e)}")
|
weights/Detection/Fabric.pt
ADDED
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
1 |
+
version https://git-lfs.github.com/spec/v1
|
2 |
+
oid sha256:962b0ed2a4c82119abb3b8b81e88ece8eb7d2d87b4dd3abf98f26bc0f743ffee
|
3 |
+
size 6248494
|
weights/Detection/Rust.pt
ADDED
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
1 |
+
version https://git-lfs.github.com/spec/v1
|
2 |
+
oid sha256:87413aa004e8bbc3df24ce38802d347397ff405c8283437fb66042f3e0a78190
|
3 |
+
size 23841602
|
weights/Detection/Scratch.pt
ADDED
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
1 |
+
version https://git-lfs.github.com/spec/v1
|
2 |
+
oid sha256:b9a4761c203cb71ae05cb2df203b5ff2a0f4690450dfa3f858b2c33b63d89c7e
|
3 |
+
size 23843768
|