datascientist22 commited on
Commit
761772a
·
verified ·
1 Parent(s): d6266f2

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +34 -26
app.py CHANGED
@@ -1,42 +1,46 @@
1
- import streamlit as st
2
  import cv2
3
- import tempfile
4
  from ultralytics import YOLO
5
- import numpy as np
 
6
 
7
- # Title and Description
 
 
 
8
  st.title("🔍 YOLOv8 Object Detection on Video")
9
- st.write("Upload a video file to detect objects using the YOLOv8 model. You can download the processed video with bounding boxes around detected objects.")
10
 
11
- # Sidebar for video upload
12
- uploaded_file = st.sidebar.file_uploader("Upload a video", type=["mp4", "avi", "mov"])
13
 
14
  if uploaded_file is not None:
15
  # Save the uploaded video to a temporary file
16
- temp_input_file = tempfile.NamedTemporaryFile(delete=False, suffix=".mp4")
17
  temp_input_file.write(uploaded_file.read())
18
- temp_input_file.close()
19
 
20
  # Display the uploaded video
21
  st.video(temp_input_file.name)
22
 
23
- # Process video button
24
- if st.sidebar.button("Submit"):
25
- # Load YOLOv8 model
26
- model = YOLO("yolov8n.pt")
27
 
28
- # Open the input video file
29
- cap = cv2.VideoCapture(temp_input_file.name)
30
 
 
 
 
 
31
  # Get video properties
32
  frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
33
  frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
34
  fps = cap.get(cv2.CAP_PROP_FPS)
35
 
36
  # Define codec and create VideoWriter object
37
- temp_output_file = tempfile.NamedTemporaryFile(delete=False, suffix=".mp4")
38
- out = cv2.VideoWriter(temp_output_file.name, cv2.VideoWriter_fourcc(*'mp4v'), fps, (frame_width, frame_height))
39
 
 
40
  while cap.isOpened():
41
  ret, frame = cap.read()
42
  if not ret:
@@ -45,7 +49,7 @@ if uploaded_file is not None:
45
  # Perform object detection
46
  results = model(frame)
47
 
48
- # Access detection results and draw bounding boxes
49
  if results:
50
  for result in results:
51
  boxes = result.boxes # Access boxes attribute
@@ -65,14 +69,18 @@ if uploaded_file is not None:
65
  # Release resources
66
  cap.release()
67
  out.release()
 
 
 
68
 
69
- # Display the processed video
70
  st.video(temp_output_file.name)
71
 
72
- # Provide download link for processed video
73
- st.sidebar.download_button(
74
- label="Download Processed Video",
75
- data=open(temp_output_file.name, "rb"),
76
- file_name="processed_video.mp4",
77
- mime="video/mp4"
78
- )
 
 
 
1
  import cv2
 
2
  from ultralytics import YOLO
3
+ import streamlit as st
4
+ from tempfile import NamedTemporaryFile
5
 
6
+ # Load YOLOv8 model
7
+ model = YOLO("yolov8n.pt")
8
+
9
+ # Streamlit UI
10
  st.title("🔍 YOLOv8 Object Detection on Video")
 
11
 
12
+ # Upload video file
13
+ uploaded_file = st.file_uploader("Upload Video", type=["mp4", "avi", "mov"])
14
 
15
  if uploaded_file is not None:
16
  # Save the uploaded video to a temporary file
17
+ temp_input_file = NamedTemporaryFile(delete=False)
18
  temp_input_file.write(uploaded_file.read())
19
+ temp_input_file.flush()
20
 
21
  # Display the uploaded video
22
  st.video(temp_input_file.name)
23
 
24
+ # Define the output video file path
25
+ temp_output_file = NamedTemporaryFile(delete=False, suffix='.mp4')
 
 
26
 
27
+ # Open the input video file
28
+ cap = cv2.VideoCapture(temp_input_file.name)
29
 
30
+ # Check if the video was opened successfully
31
+ if not cap.isOpened():
32
+ st.error("Error: Could not open video file.")
33
+ else:
34
  # Get video properties
35
  frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
36
  frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
37
  fps = cap.get(cv2.CAP_PROP_FPS)
38
 
39
  # Define codec and create VideoWriter object
40
+ fourcc = cv2.VideoWriter_fourcc(*'mp4v') # Codec for .mp4 files
41
+ out = cv2.VideoWriter(temp_output_file.name, fourcc, fps, (frame_width, frame_height))
42
 
43
+ # Process the video frame by frame
44
  while cap.isOpened():
45
  ret, frame = cap.read()
46
  if not ret:
 
49
  # Perform object detection
50
  results = model(frame)
51
 
52
+ # Access detection results
53
  if results:
54
  for result in results:
55
  boxes = result.boxes # Access boxes attribute
 
69
  # Release resources
70
  cap.release()
71
  out.release()
72
+ cv2.destroyAllWindows()
73
+
74
+ st.success("Video processing complete!")
75
 
76
+ # Display the processed video in the browser
77
  st.video(temp_output_file.name)
78
 
79
+ # Provide a download link for the processed video
80
+ with open(temp_output_file.name, 'rb') as file:
81
+ btn = st.download_button(
82
+ label="Download Processed Video",
83
+ data=file,
84
+ file_name="processed_video.mp4",
85
+ mime="video/mp4"
86
+ )