hb-setosys commited on
Commit
459da81
·
verified ·
1 Parent(s): 68b2596

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +40 -42
app.py CHANGED
@@ -2,6 +2,7 @@ import gradio as gr
2
  import cv2
3
  import numpy as np
4
  import os
 
5
 
6
  # Load YOLO model
7
  net = cv2.dnn.readNet('yolov3.weights', 'yolov3.cfg')
@@ -51,12 +52,26 @@ def count_people_in_frame(frame):
51
 
52
  return len(indexes)
53
 
54
- def count_people_video(video_path):
 
 
55
  """
56
- Process video and count people per frame.
 
 
 
 
 
 
 
 
57
  """
 
 
 
 
58
  if not os.path.exists(video_path):
59
- return "Error: Video file not found."
60
 
61
  cap = cv2.VideoCapture(video_path)
62
  if not cap.isOpened():
@@ -78,47 +93,30 @@ def count_people_video(video_path):
78
 
79
  cap.release()
80
 
81
- # Generate analytics
82
  return f"Max People Detected in Video: {max(people_per_frame) if people_per_frame else 0}"
83
 
84
- def analyze_video(video_file):
85
- # Extract video path from uploaded file
86
- video_path = video_file.name # No need for `type="file"`
87
-
88
- # Ensure path exists
89
- if not os.path.exists(video_path):
90
- return "Error: Video file could not be loaded."
91
-
92
- result = count_people_video(video_path)
93
- return result
94
-
95
- def analyze_image(image):
96
- image_cv = np.array(image) # Convert PIL image to NumPy array
97
- people_count = count_people_in_frame(image_cv)
98
- return image, f"People in Image: {people_count}"
99
-
100
- # Gradio Interface for Image Processing
101
- image_interface = gr.Interface(
102
- fn=analyze_image,
103
- inputs=gr.Image(label="Upload Image"),
104
- outputs=[gr.Image(label="Processed Image"), gr.Textbox(label="People Counting Results")],
105
- title="YOLO People Counter (Image)",
106
- description="Upload an image to detect and count people using YOLOv3."
107
- )
108
-
109
- # Gradio Interface for Video Processing
110
- video_interface = gr.Interface(
111
- fn=analyze_video,
112
- inputs=gr.Video(label="Upload Video"), # Remove `type="file"`
113
- outputs=gr.Textbox(label="People Counting Results"),
114
- title="YOLO People Counter (Video)",
115
- description="Upload a video to detect and count people using YOLOv3."
116
- )
117
-
118
- # Combine both interfaces into tabs
119
- app = gr.TabbedInterface(
120
- [image_interface, video_interface],
121
- tab_names=["Image Mode", "Video Mode"]
122
  )
123
 
124
  # Launch app
 
2
  import cv2
3
  import numpy as np
4
  import os
5
+ from PIL import Image
6
 
7
  # Load YOLO model
8
  net = cv2.dnn.readNet('yolov3.weights', 'yolov3.cfg')
 
52
 
53
  return len(indexes)
54
 
55
+ def analyze_image(image):
56
+ """
57
+ Processes an image and detects people.
58
  """
59
+ if isinstance(image, np.ndarray):
60
+ image_cv = image # Already a NumPy array
61
+ else:
62
+ image_cv = np.array(image) # Convert PIL image to NumPy array
63
+
64
+ people_count = count_people_in_frame(image_cv)
65
+ return image, f"People in Image: {people_count}"
66
+
67
+ def analyze_video(video_file):
68
  """
69
+ Processes a video and detects people in each frame.
70
+ """
71
+ video_path = video_file.name
72
+
73
  if not os.path.exists(video_path):
74
+ return "Error: Video file could not be loaded."
75
 
76
  cap = cv2.VideoCapture(video_path)
77
  if not cap.isOpened():
 
93
 
94
  cap.release()
95
 
 
96
  return f"Max People Detected in Video: {max(people_per_frame) if people_per_frame else 0}"
97
 
98
+ def process_input(input_file):
99
+ """
100
+ Determines if the input is an image or a video and calls the appropriate function.
101
+ """
102
+ file_path = input_file.name
103
+ file_extension = os.path.splitext(file_path)[1].lower()
104
+
105
+ if file_extension in [".jpg", ".jpeg", ".png", ".bmp"]:
106
+ image = Image.open(file_path)
107
+ return analyze_image(image)
108
+ elif file_extension in [".mp4", ".avi", ".mov", ".mkv"]:
109
+ return analyze_video(input_file)
110
+ else:
111
+ return "Error: Unsupported file format."
112
+
113
+ # Gradio Interface for Image and Video Processing
114
+ app = gr.Interface(
115
+ fn=process_input,
116
+ inputs=gr.File(label="Upload Image or Video"), # Use File to handle both types
117
+ outputs=[gr.Textbox(label="People Counting Results")],
118
+ title="YOLO People Counter (Image & Video)",
119
+ description="Upload an image or video to detect and count people using YOLOv3."
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
120
  )
121
 
122
  # Launch app