nagasurendra commited on
Commit
f5710ab
·
verified ·
1 Parent(s): cee940c

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +12 -9
app.py CHANGED
@@ -1,8 +1,10 @@
 
 
 
1
  import gradio as gr
2
  import cv2
3
  import numpy as np
4
  from ultralytics import YOLO
5
- import os
6
  import tempfile
7
  from moviepy.editor import ImageSequenceClip
8
  from PIL import Image
@@ -16,7 +18,7 @@ def process_video(video_path, model_name, conf_threshold=0.4):
16
  Process the input video frame by frame using the selected YOLO model,
17
  draw bounding boxes, and return the processed video path.
18
  """
19
- # Select model based on user input
20
  model = model_yolo11 if model_name == "YOLO11n" else model_best
21
 
22
  # Open video capture
@@ -37,7 +39,7 @@ def process_video(video_path, model_name, conf_threshold=0.4):
37
  if not ret:
38
  break
39
 
40
- # Perform object detection
41
  results = model.predict(
42
  source=frame,
43
  conf=conf_threshold,
@@ -46,21 +48,21 @@ def process_video(video_path, model_name, conf_threshold=0.4):
46
  show_conf=True
47
  )
48
 
49
- # Draw bounding boxes on the frame
50
  for result in results:
51
- im_array = result.plot() # Plot bounding boxes
52
- processed_frames.append(im_array[..., ::-1]) # Convert BGR to RGB
53
 
54
  cap.release()
55
 
56
- # Save processed frames to a temporary video file
57
  temp_video_path = os.path.join(tempfile.gettempdir(), "output.mp4")
58
  clip = ImageSequenceClip(processed_frames, fps=fps)
59
  clip.write_videofile(temp_video_path, codec='libx264')
60
 
61
  return temp_video_path
62
 
63
- # Define Gradio interface
64
  with gr.Blocks() as app:
65
  gr.HTML("""
66
  <h1 style='text-align: center'>
@@ -99,4 +101,5 @@ with gr.Blocks() as app:
99
  )
100
 
101
  if __name__ == "__main__":
102
- app.launch()
 
 
1
+ ```python
2
+ import os
3
+ os.environ['YOLO_CONFIG_DIR'] = '/tmp/Ultralytics' # Set Ultralytics config path
4
  import gradio as gr
5
  import cv2
6
  import numpy as np
7
  from ultralytics import YOLO
 
8
  import tempfile
9
  from moviepy.editor import ImageSequenceClip
10
  from PIL import Image
 
18
  Process the input video frame by frame using the selected YOLO model,
19
  draw bounding boxes, and return the processed video path.
20
  """
21
+ # Select model to use
22
  model = model_yolo11 if model_name == "YOLO11n" else model_best
23
 
24
  # Open video capture
 
39
  if not ret:
40
  break
41
 
42
+ # Perform detection
43
  results = model.predict(
44
  source=frame,
45
  conf=conf_threshold,
 
48
  show_conf=True
49
  )
50
 
51
+ # Draw bounding boxes
52
  for result in results:
53
+ im_array = result.plot() # Plot boxes
54
+ processed_frames.append(im_array[..., ::-1]) # Convert BGR to RGB
55
 
56
  cap.release()
57
 
58
+ # Save processed frames to temp video
59
  temp_video_path = os.path.join(tempfile.gettempdir(), "output.mp4")
60
  clip = ImageSequenceClip(processed_frames, fps=fps)
61
  clip.write_videofile(temp_video_path, codec='libx264')
62
 
63
  return temp_video_path
64
 
65
+ # Gradio interface
66
  with gr.Blocks() as app:
67
  gr.HTML("""
68
  <h1 style='text-align: center'>
 
101
  )
102
 
103
  if __name__ == "__main__":
104
+ app.launch()
105
+ ```