Harsimran19 commited on
Commit
08d6e3c
1 Parent(s): 61b13b4

Upload 7 files

Browse files
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