tom-b974 commited on
Commit
ea3aa78
·
verified ·
1 Parent(s): 497ffcb

Update tasks/image.py

Browse files
Files changed (1) hide show
  1. 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("best.pt")
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 = [] # Flattened list of predicted boxes
108
- true_boxes_list = [] # Flattened list of ground truth boxes
109
 
110
  for example in test_dataset:
111
  # Extract image and annotations
112
  image = example["image"]
113
  annotation = example.get("annotations", "").strip()
114
 
115
- # Determine if ground truth smoke is present
116
  has_smoke = len(annotation) > 0
117
  true_labels.append(1 if has_smoke else 0)
118
 
119
- # Parse ground truth boxes if smoke is present
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([]) # Add empty list if parsing fails
126
  else:
127
- true_boxes_list.append([]) # Add empty list for no ground truth smoke
128
 
129
- # Perform YOLO inference
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) # Predicted smoke
136
  pred_boxes.append(pred_box)
137
  else:
138
- predictions.append(0) # No smoke predicted
139
- pred_boxes.append([]) # Add empty list if no prediction
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: # Keep only if neither is empty
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
  #--------------------------------------------------------------------------------------------