Spaces:
Sleeping
Sleeping
Update services/plantation/plant_count.py
Browse files
services/plantation/plant_count.py
CHANGED
@@ -15,11 +15,12 @@ def process_plants(frame):
|
|
15 |
frame: Input frame (numpy array)
|
16 |
Returns:
|
17 |
list: List of detected plants
|
18 |
-
numpy array:
|
19 |
"""
|
20 |
results = model(frame)
|
21 |
|
22 |
detections = []
|
|
|
23 |
|
24 |
for r in results:
|
25 |
for box in r.boxes:
|
@@ -33,10 +34,21 @@ def process_plants(frame):
|
|
33 |
xyxy = box.xyxy[0].cpu().numpy()
|
34 |
x_min, y_min, x_max, y_max = map(int, xyxy)
|
35 |
|
|
|
|
|
36 |
detections.append({
|
37 |
"type": label,
|
|
|
38 |
"confidence": conf,
|
39 |
"coordinates": [x_min, y_min, x_max, y_max]
|
40 |
})
|
41 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
42 |
return detections, frame
|
|
|
15 |
frame: Input frame (numpy array)
|
16 |
Returns:
|
17 |
list: List of detected plants
|
18 |
+
numpy array: Annotated frame with numbered labels
|
19 |
"""
|
20 |
results = model(frame)
|
21 |
|
22 |
detections = []
|
23 |
+
line_counter = 1 # Initialize counter for numbered labels
|
24 |
|
25 |
for r in results:
|
26 |
for box in r.boxes:
|
|
|
34 |
xyxy = box.xyxy[0].cpu().numpy()
|
35 |
x_min, y_min, x_max, y_max = map(int, xyxy)
|
36 |
|
37 |
+
# Add numbered label
|
38 |
+
detection_label = f"Line {line_counter} - {label.capitalize()} (Conf: {conf:.2f})"
|
39 |
detections.append({
|
40 |
"type": label,
|
41 |
+
"label": detection_label,
|
42 |
"confidence": conf,
|
43 |
"coordinates": [x_min, y_min, x_max, y_max]
|
44 |
})
|
45 |
|
46 |
+
# Draw bounding box and label
|
47 |
+
color = (34, 139, 34) # ForestGreen for plants
|
48 |
+
cv2.rectangle(frame, (x_min, y_min), (x_max, y_max), color, 2)
|
49 |
+
cv2.putText(frame, detection_label, (x_min, y_min - 10),
|
50 |
+
cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
|
51 |
+
|
52 |
+
line_counter += 1
|
53 |
+
|
54 |
return detections, frame
|