Spaces:
Sleeping
Sleeping
Update tasks/image.py
Browse files- tasks/image.py +14 -21
tasks/image.py
CHANGED
@@ -18,7 +18,7 @@ router = APIRouter()
|
|
18 |
DESCRIPTION = "YOLO Smoke Detection"
|
19 |
ROUTE = "/image"
|
20 |
|
21 |
-
yolo_model = YOLO("
|
22 |
|
23 |
def parse_boxes(annotation_string):
|
24 |
"""Parse multiple boxes from a single annotation string.
|
@@ -104,57 +104,50 @@ async def evaluate_image(request: ImageEvaluationRequest):
|
|
104 |
#--------------------------------------------------------------------------------------------
|
105 |
predictions = []
|
106 |
true_labels = []
|
107 |
-
pred_boxes = []
|
108 |
-
true_boxes_list = []
|
109 |
|
110 |
for example in test_dataset:
|
111 |
# Extract image and annotations
|
112 |
image = example["image"]
|
113 |
annotation = example.get("annotations", "").strip()
|
114 |
|
115 |
-
|
116 |
has_smoke = len(annotation) > 0
|
117 |
true_labels.append(1 if has_smoke else 0)
|
118 |
|
119 |
-
|
120 |
if has_smoke:
|
121 |
image_true_boxes = parse_boxes(annotation)
|
122 |
if image_true_boxes:
|
123 |
true_boxes_list.append(image_true_boxes)
|
124 |
else:
|
125 |
-
true_boxes_list.append([])
|
126 |
else:
|
127 |
-
true_boxes_list.append([])
|
128 |
|
129 |
-
#
|
130 |
-
results = yolo_model .predict(image, verbose=False)
|
131 |
|
132 |
-
# Extract predicted box if predictions exist
|
133 |
if len(results[0].boxes):
|
134 |
pred_box = results[0].boxes.xywhn[0].cpu().numpy().tolist()
|
135 |
-
predictions.append(1)
|
136 |
pred_boxes.append(pred_box)
|
137 |
else:
|
138 |
-
predictions.append(0)
|
139 |
-
pred_boxes.append([])
|
140 |
|
141 |
-
# Filter out entries with empty boxes
|
142 |
filtered_true_boxes_list = []
|
143 |
filtered_pred_boxes = []
|
144 |
|
145 |
-
for true_boxes, pred_boxes_entry in zip(true_boxes_list, pred_boxes):
|
146 |
-
if true_boxes and pred_boxes_entry:
|
147 |
filtered_true_boxes_list.append(true_boxes)
|
148 |
filtered_pred_boxes.append(pred_boxes_entry)
|
149 |
|
150 |
-
|
151 |
-
# Replace the original lists with the filtered ones
|
152 |
true_boxes_list = filtered_true_boxes_list
|
153 |
pred_boxes = filtered_pred_boxes
|
154 |
|
155 |
-
|
156 |
-
|
157 |
-
|
158 |
#--------------------------------------------------------------------------------------------
|
159 |
# YOUR MODEL INFERENCE STOPS HERE
|
160 |
#--------------------------------------------------------------------------------------------
|
|
|
18 |
DESCRIPTION = "YOLO Smoke Detection"
|
19 |
ROUTE = "/image"
|
20 |
|
21 |
+
yolo_model = YOLO("yolo11_tr_20012025_frugalAIchal.pt")
|
22 |
|
23 |
def parse_boxes(annotation_string):
|
24 |
"""Parse multiple boxes from a single annotation string.
|
|
|
104 |
#--------------------------------------------------------------------------------------------
|
105 |
predictions = []
|
106 |
true_labels = []
|
107 |
+
pred_boxes = []
|
108 |
+
true_boxes_list = []
|
109 |
|
110 |
for example in test_dataset:
|
111 |
# Extract image and annotations
|
112 |
image = example["image"]
|
113 |
annotation = example.get("annotations", "").strip()
|
114 |
|
115 |
+
|
116 |
has_smoke = len(annotation) > 0
|
117 |
true_labels.append(1 if has_smoke else 0)
|
118 |
|
119 |
+
|
120 |
if has_smoke:
|
121 |
image_true_boxes = parse_boxes(annotation)
|
122 |
if image_true_boxes:
|
123 |
true_boxes_list.append(image_true_boxes)
|
124 |
else:
|
125 |
+
true_boxes_list.append([])
|
126 |
else:
|
127 |
+
true_boxes_list.append([])
|
128 |
|
129 |
+
results = yolo_model .predict(image, verbose=False) # INFERENCE - prediction
|
|
|
130 |
|
|
|
131 |
if len(results[0].boxes):
|
132 |
pred_box = results[0].boxes.xywhn[0].cpu().numpy().tolist()
|
133 |
+
predictions.append(1)
|
134 |
pred_boxes.append(pred_box)
|
135 |
else:
|
136 |
+
predictions.append(0)
|
137 |
+
pred_boxes.append([])
|
138 |
|
|
|
139 |
filtered_true_boxes_list = []
|
140 |
filtered_pred_boxes = []
|
141 |
|
142 |
+
for true_boxes, pred_boxes_entry in zip(true_boxes_list, pred_boxes): # Only see when annotation(s) is/are both on true label and prediction
|
143 |
+
if true_boxes and pred_boxes_entry:
|
144 |
filtered_true_boxes_list.append(true_boxes)
|
145 |
filtered_pred_boxes.append(pred_boxes_entry)
|
146 |
|
147 |
+
|
|
|
148 |
true_boxes_list = filtered_true_boxes_list
|
149 |
pred_boxes = filtered_pred_boxes
|
150 |
|
|
|
|
|
|
|
151 |
#--------------------------------------------------------------------------------------------
|
152 |
# YOUR MODEL INFERENCE STOPS HERE
|
153 |
#--------------------------------------------------------------------------------------------
|