Skip to content

Commit

Permalink
Merge pull request #91 from ks6088ts-labs/feature/issue-88_refactor-v…
Browse files Browse the repository at this point in the history
…ideo

support file input source
  • Loading branch information
ks6088ts authored Aug 24, 2024
2 parents 49fa7d0 + 229d0a4 commit 3ab07f8
Show file tree
Hide file tree
Showing 3 changed files with 64 additions and 20 deletions.
2 changes: 2 additions & 0 deletions apps/99_streamlit_examples/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -81,6 +81,8 @@ Access to http://localhost:8501 and select the sample you want to run from the s

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

ref. [data-videos/traffic.mp4](https://github.com/OlafenwaMoses/ImageAI/blob/master/data-videos/traffic.mp4)

## References

- [🎈 Streamlit + LLM Examples App](https://github.com/streamlit/llm-examples)
Expand Down
82 changes: 62 additions & 20 deletions apps/99_streamlit_examples/pages/12_Video_processing.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,11 @@ class ProcessorType(Enum):
YOLOV8 = "yolov8"


class InputSource(Enum):
CAMERA = "camera"
FILE = "file"


class Processor:
def process(
self,
Expand Down Expand Up @@ -79,10 +84,13 @@ def __init__(
self,
model_name: str = "yolov8n.pt",
confidence: float = 0.5,
# https://stackoverflow.com/a/77479465
classes: list[int] = None,
):
# model_name: https://docs.ultralytics.com/models/yolov8/#supported-tasks-and-modes
self.model = YOLO(model_name)
self.confidence = confidence
self.classes = classes

def process(
self,
Expand All @@ -91,7 +99,7 @@ def process(
results = self.model(
frame,
conf=self.confidence,
classes=[0],
classes=self.classes,
)
output_img = results[0].plot(
labels=True,
Expand All @@ -117,32 +125,66 @@ def get_processor(processor_type: ProcessorType) -> Processor:


with st.sidebar:
# device: https://docs.opencv.org/4.10.0/d8/dfe/classcv_1_1VideoCapture.html#a5d5f5dacb77bbebdcbfb341e3d4355c1
device = st.text_input(
label="input your video/camera device",
value="0",
)
if device.isnumeric():
# e.g. "0" -> 0
device = int(device)
processor_type = st.radio(
label="processor type",
options=[
ProcessorType.BLUR,
ProcessorType.CANNY,
ProcessorType.INVERT,
ProcessorType.YOLOV8,
],
index=0,
format_func=lambda x: x.value,
tab_input, tab_mode = st.tabs(
[
"input",
"mode",
]
)
with tab_input:
source = st.radio(
label="input source",
options=[
InputSource.FILE,
InputSource.CAMERA,
],
index=0,
format_func=lambda x: x.value,
)
if source == InputSource.FILE:
file = st.file_uploader(
label="upload video file",
type=[
"mp4",
"mov",
"avi",
],
)
if file:
file_path = f"/tmp/{file.name}"
with open(file_path, "wb") as f:
f.write(file.read())
device = file_path
if source == InputSource.CAMERA:
# device: https://docs.opencv.org/4.10.0/d8/dfe/classcv_1_1VideoCapture.html#a5d5f5dacb77bbebdcbfb341e3d4355c1
device = st.text_input(
label="input your video/camera device",
value="0",
)
if device.isnumeric():
# e.g. "0" -> 0
device = int(device)

with tab_mode:
processor_type = st.radio(
label="processor type",
options=[
ProcessorType.BLUR,
ProcessorType.CANNY,
ProcessorType.INVERT,
ProcessorType.YOLOV8,
],
index=0,
format_func=lambda x: x.value,
)

st.title("Video processing")

start_button = st.button("Start")
stop = st.button("Stop")

image_loc = st.empty()
message_loc = st.empty()
processor = get_processor(processor_type)

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

if not ret:
st.error("Failed to capture image")
message_loc.error("Failed to read frame")
continue

processed_frame = processor.process(
Expand Down
Binary file modified docs/images/99_streamlit_examples.videoprocessing.png
Loading
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.

0 comments on commit 3ab07f8

Please sign in to comment.