DawnC commited on
Commit
0cfe815
1 Parent(s): d8ba3f9

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +43 -69
app.py CHANGED
@@ -168,39 +168,10 @@ async def predict_single_dog(image):
168
  return top1_prob, topk_breeds, topk_probs_percent
169
 
170
 
171
- # async def detect_multiple_dogs(image, conf_threshold=0.35, iou_threshold=0.55):
172
- # results = model_yolo(image, conf=conf_threshold, iou=iou_threshold)[0]
173
- # dogs = []
174
- # boxes = []
175
- # for box in results.boxes:
176
- # if box.cls == 16: # COCO dataset class for dog is 16
177
- # xyxy = box.xyxy[0].tolist()
178
- # confidence = box.conf.item()
179
- # boxes.append((xyxy, confidence))
180
-
181
- # if not boxes:
182
- # dogs.append((image, 1.0, [0, 0, image.width, image.height]))
183
- # else:
184
- # nms_boxes = non_max_suppression(boxes, iou_threshold)
185
-
186
- # for box, confidence in nms_boxes:
187
- # x1, y1, x2, y2 = box
188
- # w, h = x2 - x1, y2 - y1
189
- # x1 = max(0, x1 - w * 0.05)
190
- # y1 = max(0, y1 - h * 0.05)
191
- # x2 = min(image.width, x2 + w * 0.05)
192
- # y2 = min(image.height, y2 + h * 0.05)
193
- # cropped_image = image.crop((x1, y1, x2, y2))
194
- # dogs.append((cropped_image, confidence, [x1, y1, x2, y2]))
195
-
196
- # return dogs
197
-
198
-
199
- async def detect_multiple_dogs(image, conf_threshold=0.3, iou_threshold=0.4):
200
  results = model_yolo(image, conf=conf_threshold, iou=iou_threshold)[0]
201
  dogs = []
202
  boxes = []
203
-
204
  for box in results.boxes:
205
  if box.cls == 16: # COCO dataset class for dog is 16
206
  xyxy = box.xyxy[0].tolist()
@@ -213,54 +184,57 @@ async def detect_multiple_dogs(image, conf_threshold=0.3, iou_threshold=0.4):
213
  nms_boxes = non_max_suppression(boxes, iou_threshold)
214
 
215
  for box, confidence in nms_boxes:
216
- x1, y1, x2, y2 = [int(coord) for coord in box]
 
 
 
 
 
217
  cropped_image = image.crop((x1, y1, x2, y2))
218
  dogs.append((cropped_image, confidence, [x1, y1, x2, y2]))
219
 
220
- # 應用過濾器來移除可能的錯誤檢測
221
- dogs = filter_detections(dogs, (image.width, image.height))
222
-
223
  return dogs
224
 
225
- def filter_detections(dogs, image_size):
226
- filtered_dogs = []
227
- image_area = image_size[0] * image_size[1]
 
 
228
 
229
- for dog in dogs:
230
- _, _, box = dog
231
- dog_area = (box[2] - box[0]) * (box[3] - box[1])
232
- area_ratio = dog_area / image_area
 
 
 
 
 
 
233
 
234
- if 0.01 < area_ratio < 0.9: # 過濾掉太小或太大的檢測框
235
- filtered_dogs.append(dog)
 
 
236
 
237
- return filtered_dogs
238
-
239
-
240
- def merge_overlapping_boxes(boxes, overlap_threshold):
241
- merged = []
242
- while boxes:
243
- base_box = boxes.pop(0)
244
- i = 0
245
- while i < len(boxes):
246
- if calculate_iou(base_box[0], boxes[i][0]) > overlap_threshold:
247
- # 合併框,取較大的置信度
248
- merged_box = merge_boxes(base_box[0], boxes[i][0])
249
- merged_conf = max(base_box[1], boxes[i][1])
250
- base_box = (merged_box, merged_conf)
251
- boxes.pop(i)
252
- else:
253
- i += 1
254
- merged.append(base_box)
255
- return merged
256
-
257
- def merge_boxes(box1, box2):
258
- x1 = min(box1[0], box2[0])
259
- y1 = min(box1[1], box2[1])
260
- x2 = max(box1[2], box2[2])
261
- y2 = max(box1[3], box2[3])
262
- return [x1, y1, x2, y2]
263
 
 
 
 
 
 
 
 
 
 
 
 
 
 
264
 
265
 
266
  def non_max_suppression(boxes, iou_threshold):
 
168
  return top1_prob, topk_breeds, topk_probs_percent
169
 
170
 
171
+ async def detect_multiple_dogs(image, conf_threshold=0.35, iou_threshold=0.55):
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
172
  results = model_yolo(image, conf=conf_threshold, iou=iou_threshold)[0]
173
  dogs = []
174
  boxes = []
 
175
  for box in results.boxes:
176
  if box.cls == 16: # COCO dataset class for dog is 16
177
  xyxy = box.xyxy[0].tolist()
 
184
  nms_boxes = non_max_suppression(boxes, iou_threshold)
185
 
186
  for box, confidence in nms_boxes:
187
+ x1, y1, x2, y2 = box
188
+ w, h = x2 - x1, y2 - y1
189
+ x1 = max(0, x1 - w * 0.05)
190
+ y1 = max(0, y1 - h * 0.05)
191
+ x2 = min(image.width, x2 + w * 0.05)
192
+ y2 = min(image.height, y2 + h * 0.05)
193
  cropped_image = image.crop((x1, y1, x2, y2))
194
  dogs.append((cropped_image, confidence, [x1, y1, x2, y2]))
195
 
 
 
 
196
  return dogs
197
 
198
+
199
+ # async def detect_multiple_dogs(image, conf_threshold=0.35, iou_threshold=0.5):
200
+ # results = model_yolo(image, conf=conf_threshold, iou=iou_threshold)[0]
201
+ # dogs = []
202
+ # boxes = []
203
 
204
+ # for box in results.boxes:
205
+ # if box.cls == 16: # COCO dataset class for dog is 16
206
+ # xyxy = box.xyxy[0].tolist()
207
+ # confidence = box.conf.item()
208
+ # boxes.append((xyxy, confidence))
209
+
210
+ # if not boxes:
211
+ # dogs.append((image, 1.0, [0, 0, image.width, image.height]))
212
+ # else:
213
+ # nms_boxes = non_max_suppression(boxes, iou_threshold)
214
 
215
+ # for box, confidence in nms_boxes:
216
+ # x1, y1, x2, y2 = [int(coord) for coord in box]
217
+ # cropped_image = image.crop((x1, y1, x2, y2))
218
+ # dogs.append((cropped_image, confidence, [x1, y1, x2, y2]))
219
 
220
+ # # 應用過濾器來移除可能的錯誤檢測
221
+ # dogs = filter_detections(dogs, (image.width, image.height))
222
+
223
+ # return dogs
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
224
 
225
+ # def filter_detections(dogs, image_size):
226
+ # filtered_dogs = []
227
+ # image_area = image_size[0] * image_size[1]
228
+
229
+ # for dog in dogs:
230
+ # _, _, box = dog
231
+ # dog_area = (box[2] - box[0]) * (box[3] - box[1])
232
+ # area_ratio = dog_area / image_area
233
+
234
+ # if 0.01 < area_ratio < 0.9: # 過濾掉太小或太大的檢測框
235
+ # filtered_dogs.append(dog)
236
+
237
+ # return filtered_dogs
238
 
239
 
240
  def non_max_suppression(boxes, iou_threshold):