Spaces:
Sleeping
Sleeping
Update app.py
Browse files
app.py
CHANGED
@@ -37,7 +37,7 @@ model.eval()
|
|
37 |
# transforms.ToTensor()
|
38 |
# ])
|
39 |
transform = transforms.Compose([ # Ensure input is a PIL image
|
40 |
-
transforms.Resize((
|
41 |
transforms.ToTensor()
|
42 |
])
|
43 |
# transform = transforms.Compose([
|
@@ -148,65 +148,7 @@ def detect_objects_in_image(image):
|
|
148 |
|
149 |
# return Image.open(buf)
|
150 |
|
151 |
-
def generate_vehicle_count_graph(object_counts):
|
152 |
-
color_palette = ['#4C9ACD', '#88B8A3', '#7F9C9C', '#D1A3B5', '#A1C6EA', '#FFB6C1', '#F0E68C', '#D3B0D8', '#F8A5D1',
|
153 |
-
'#B8B8D1']
|
154 |
|
155 |
-
fig, ax = plt.subplots(figsize=(8, 5))
|
156 |
-
labels = list(object_counts.keys())
|
157 |
-
values = list(object_counts.values())
|
158 |
-
|
159 |
-
ax.bar(labels, values, color=color_palette[:len(labels)])
|
160 |
-
ax.set_xlabel("Vehicle Categories", fontsize=12, fontweight='bold')
|
161 |
-
ax.set_ylabel("Number of Vehicles", fontsize=12, fontweight='bold')
|
162 |
-
ax.set_title("Detected Vehicles in Image", fontsize=14, fontweight='bold')
|
163 |
-
|
164 |
-
plt.xticks(rotation=45, ha='right', fontsize=10)
|
165 |
-
plt.yticks(fontsize=10)
|
166 |
-
plt.tight_layout()
|
167 |
-
|
168 |
-
buf = BytesIO()
|
169 |
-
plt.savefig(buf, format='png')
|
170 |
-
buf.seek(0)
|
171 |
-
|
172 |
-
plt.close(fig) # ✅ CLOSE THE FIGURE TO FREE MEMORY
|
173 |
-
print("checkl1")
|
174 |
-
return Image.open(buf)
|
175 |
-
|
176 |
-
|
177 |
-
def detect_objects_in_video(video_input):
|
178 |
-
cap = cv2.VideoCapture(video_input)
|
179 |
-
if not cap.isOpened():
|
180 |
-
return "Error: Cannot open video file.", None # Returning a second value (None) to match expected outputs
|
181 |
-
|
182 |
-
frame_width, frame_height, fps = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)), int(
|
183 |
-
cap.get(cv2.CAP_PROP_FRAME_HEIGHT)), int(cap.get(cv2.CAP_PROP_FPS))
|
184 |
-
temp_video_output = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
|
185 |
-
out = cv2.VideoWriter(temp_video_output, cv2.VideoWriter_fourcc(*'mp4v'), fps, (frame_width, frame_height))
|
186 |
-
|
187 |
-
# Initialize the counts for vehicle categories
|
188 |
-
total_counts = {name: 0 for name in ['car', 'truck', 'bus', 'motorcycle', 'bicycle']}
|
189 |
-
|
190 |
-
while cap.isOpened():
|
191 |
-
ret, frame = cap.read()
|
192 |
-
if not ret:
|
193 |
-
break
|
194 |
-
|
195 |
-
|
196 |
-
|
197 |
-
image = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
|
198 |
-
|
199 |
-
# Get frame with detected objects and graph
|
200 |
-
frame_with_boxes, graph_image = detect_objects_in_image(image)
|
201 |
-
|
202 |
-
# Convert image back to OpenCV format for writing video
|
203 |
-
out.write(cv2.cvtColor(np.array(frame_with_boxes), cv2.COLOR_RGB2BGR))
|
204 |
-
|
205 |
-
cap.release()
|
206 |
-
out.release()
|
207 |
-
print("checl2")
|
208 |
-
|
209 |
-
return temp_video_output, graph_image # Return both expected outputs
|
210 |
|
211 |
|
212 |
|
|
|
37 |
# transforms.ToTensor()
|
38 |
# ])
|
39 |
transform = transforms.Compose([ # Ensure input is a PIL image
|
40 |
+
transforms.Resize((960, 540)),
|
41 |
transforms.ToTensor()
|
42 |
])
|
43 |
# transform = transforms.Compose([
|
|
|
148 |
|
149 |
# return Image.open(buf)
|
150 |
|
|
|
|
|
|
|
151 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
152 |
|
153 |
|
154 |
|