Skip to content

Commit 229d0a4

Browse files
committed
support input source
1 parent 49fa7d0 commit 229d0a4

File tree

3 files changed

+64
-20
lines changed

3 files changed

+64
-20
lines changed

apps/99_streamlit_examples/README.md

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -81,6 +81,8 @@ Access to http://localhost:8501 and select the sample you want to run from the s
8181

8282
![Video processing](../../docs/images/99_streamlit_examples.videoprocessing.png)
8383

84+
ref. [data-videos/traffic.mp4](https://github.com/OlafenwaMoses/ImageAI/blob/master/data-videos/traffic.mp4)
85+
8486
## References
8587

8688
- [🎈 Streamlit + LLM Examples App](https://github.com/streamlit/llm-examples)

apps/99_streamlit_examples/pages/12_Video_processing.py

Lines changed: 62 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,11 @@ class ProcessorType(Enum):
1616
YOLOV8 = "yolov8"
1717

1818

19+
class InputSource(Enum):
20+
CAMERA = "camera"
21+
FILE = "file"
22+
23+
1924
class Processor:
2025
def process(
2126
self,
@@ -79,10 +84,13 @@ def __init__(
7984
self,
8085
model_name: str = "yolov8n.pt",
8186
confidence: float = 0.5,
87+
# https://stackoverflow.com/a/77479465
88+
classes: list[int] = None,
8289
):
8390
# model_name: https://docs.ultralytics.com/models/yolov8/#supported-tasks-and-modes
8491
self.model = YOLO(model_name)
8592
self.confidence = confidence
93+
self.classes = classes
8694

8795
def process(
8896
self,
@@ -91,7 +99,7 @@ def process(
9199
results = self.model(
92100
frame,
93101
conf=self.confidence,
94-
classes=[0],
102+
classes=self.classes,
95103
)
96104
output_img = results[0].plot(
97105
labels=True,
@@ -117,32 +125,66 @@ def get_processor(processor_type: ProcessorType) -> Processor:
117125

118126

119127
with st.sidebar:
120-
# device: https://docs.opencv.org/4.10.0/d8/dfe/classcv_1_1VideoCapture.html#a5d5f5dacb77bbebdcbfb341e3d4355c1
121-
device = st.text_input(
122-
label="input your video/camera device",
123-
value="0",
124-
)
125-
if device.isnumeric():
126-
# e.g. "0" -> 0
127-
device = int(device)
128-
processor_type = st.radio(
129-
label="processor type",
130-
options=[
131-
ProcessorType.BLUR,
132-
ProcessorType.CANNY,
133-
ProcessorType.INVERT,
134-
ProcessorType.YOLOV8,
135-
],
136-
index=0,
137-
format_func=lambda x: x.value,
128+
tab_input, tab_mode = st.tabs(
129+
[
130+
"input",
131+
"mode",
132+
]
138133
)
134+
with tab_input:
135+
source = st.radio(
136+
label="input source",
137+
options=[
138+
InputSource.FILE,
139+
InputSource.CAMERA,
140+
],
141+
index=0,
142+
format_func=lambda x: x.value,
143+
)
144+
if source == InputSource.FILE:
145+
file = st.file_uploader(
146+
label="upload video file",
147+
type=[
148+
"mp4",
149+
"mov",
150+
"avi",
151+
],
152+
)
153+
if file:
154+
file_path = f"/tmp/{file.name}"
155+
with open(file_path, "wb") as f:
156+
f.write(file.read())
157+
device = file_path
158+
if source == InputSource.CAMERA:
159+
# device: https://docs.opencv.org/4.10.0/d8/dfe/classcv_1_1VideoCapture.html#a5d5f5dacb77bbebdcbfb341e3d4355c1
160+
device = st.text_input(
161+
label="input your video/camera device",
162+
value="0",
163+
)
164+
if device.isnumeric():
165+
# e.g. "0" -> 0
166+
device = int(device)
167+
168+
with tab_mode:
169+
processor_type = st.radio(
170+
label="processor type",
171+
options=[
172+
ProcessorType.BLUR,
173+
ProcessorType.CANNY,
174+
ProcessorType.INVERT,
175+
ProcessorType.YOLOV8,
176+
],
177+
index=0,
178+
format_func=lambda x: x.value,
179+
)
139180

140181
st.title("Video processing")
141182

142183
start_button = st.button("Start")
143184
stop = st.button("Stop")
144185

145186
image_loc = st.empty()
187+
message_loc = st.empty()
146188
processor = get_processor(processor_type)
147189

148190
if start_button:
@@ -152,7 +194,7 @@ def get_processor(processor_type: ProcessorType) -> Processor:
152194
ret, frame = capture.read()
153195

154196
if not ret:
155-
st.error("Failed to capture image")
197+
message_loc.error("Failed to read frame")
156198
continue
157199

158200
processed_frame = processor.process(
-604 KB
Loading

0 commit comments

Comments
 (0)