DawnC commited on
Commit
8f303b0
1 Parent(s): 86d9dde

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +4 -171
app.py CHANGED
@@ -283,94 +283,18 @@ def _predict_single_dog(image):
283
  topk_probs_percent = [f"{prob.item() * 100:.2f}%" for prob in topk_probs[0]]
284
  return top1_prob, topk_breeds, topk_probs_percent
285
 
286
- # async def detect_multiple_dogs(image, conf_threshold=0.3):
287
- # # 調整 YOLO 模型的置信度閾值
288
- # return await asyncio.to_thread(_detect_multiple_dogs, image, conf_threshold)
289
-
290
- # def _detect_multiple_dogs(image, conf_threshold):
291
- # results = model_yolo(image, conf=conf_threshold)
292
- # dogs = []
293
- # for result in results:
294
- # for box in result.boxes:
295
- # if box.cls == 16: # COCO 資料集中狗的類別是 16
296
- # xyxy = box.xyxy[0].tolist()
297
- # confidence = box.conf.item()
298
- # if confidence >= conf_threshold: # 只保留置信度高於閾值的框
299
- # cropped_image = image.crop((xyxy[0], xyxy[1], xyxy[2], xyxy[3]))
300
- # dogs.append((cropped_image, confidence, xyxy))
301
- # return dogs
302
-
303
- # async def detect_multiple_dogs(image, conf_threshold=0.2, iou_threshold=0.5):
304
- # results = model_yolo(image, conf=conf_threshold, iou=iou_threshold)[0]
305
- # dogs = []
306
- # for box in results.boxes:
307
- # if box.cls == 16: # COCO 資料集中狗的類別是 16
308
- # xyxy = box.xyxy[0].tolist()
309
- # confidence = box.conf.item()
310
- # cropped_image = image.crop((xyxy[0], xyxy[1], xyxy[2], xyxy[3]))
311
- # dogs.append((cropped_image, confidence, xyxy))
312
- # return dogs
313
- # 此為如果後面調不好 使用的版本
314
-
315
- async def detect_multiple_dogs(image, conf_threshold=0.3, iou_threshold=0.5):
316
  results = model_yolo(image, conf=conf_threshold, iou=iou_threshold)[0]
317
  dogs = []
318
  for box in results.boxes:
319
  if box.cls == 16: # COCO 資料集中狗的類別是 16
320
  xyxy = box.xyxy[0].tolist()
321
  confidence = box.conf.item()
322
- area = (xyxy[2] - xyxy[0]) * (xyxy[3] - xyxy[1])
323
- if area > 1000: # 過濾掉太小的檢測框
324
- cropped_image = image.crop((xyxy[0], xyxy[1], xyxy[2], xyxy[3]))
325
- dogs.append((cropped_image, confidence, xyxy))
326
-
327
- # 合併重疊的檢測框
328
- dogs = merge_overlapping_boxes(dogs, iou_threshold=0.6)
329
-
330
  return dogs
331
 
332
- def merge_overlapping_boxes(dogs, iou_threshold=0.6):
333
- merged_dogs = []
334
- while dogs:
335
- base = dogs.pop(0)
336
- i = 0
337
- while i < len(dogs):
338
- if calculate_iou(base[2], dogs[i][2]) > iou_threshold:
339
- # 合併重疊的框
340
- base = merge_boxes(base, dogs.pop(i))
341
- else:
342
- i += 1
343
- merged_dogs.append(base)
344
- return merged_dogs
345
-
346
- def merge_boxes(box1, box2):
347
- xyxy1, conf1, _ = box1
348
- xyxy2, conf2, _ = box2
349
- merged_xyxy = [
350
- min(xyxy1[0], xyxy2[0]),
351
- min(xyxy1[1], xyxy2[1]),
352
- max(xyxy1[2], xyxy2[2]),
353
- max(xyxy1[3], xyxy2[3])
354
- ]
355
- merged_conf = max(conf1, conf2)
356
- merged_image = Image.new('RGB', (int(merged_xyxy[2] - merged_xyxy[0]), int(merged_xyxy[3] - merged_xyxy[1])))
357
- merged_image.paste(box1[0], (0, 0))
358
- return (merged_image, merged_conf, merged_xyxy)
359
-
360
- def calculate_iou(box1, box2):
361
- # 計算兩個邊界框的交集面積
362
- x1 = max(box1[0], box2[0])
363
- y1 = max(box1[1], box2[1])
364
- x2 = min(box1[2], box2[2])
365
- y2 = min(box1[3], box2[3])
366
-
367
- intersection = max(0, x2 - x1) * max(0, y2 - y1)
368
-
369
- area1 = (box1[2] - box1[0]) * (box1[3] - box1[1])
370
- area2 = (box2[2] - box2[0]) * (box2[3] - box2[1])
371
-
372
- iou = intersection / float(area1 + area2 - intersection)
373
- return iou
374
 
375
 
376
  # async def predict(image):
@@ -484,97 +408,6 @@ def calculate_iou(box1, box2):
484
  # if __name__ == "__main__":
485
  # iface.launch()
486
 
487
- async def predict(image):
488
- if image is None:
489
- return "Please upload an image to start.", None, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
490
-
491
- try:
492
- if isinstance(image, np.ndarray):
493
- image = Image.fromarray(image)
494
-
495
- # 嘗試檢測多隻狗,進一步降低閾值以提高檢測率
496
- dogs = await detect_multiple_dogs(image, conf_threshold=0.05) # 降低閾值以檢測更多狗
497
-
498
- if len(dogs) <= 1:
499
- # 單狗情境
500
- return await process_single_dog(image)
501
-
502
- # 多狗情境
503
- color_list = ['#FF0000', '#00FF00', '#0000FF', '#FFFF00', '#00FFFF', '#FF00FF', '#800080', '#FFA500']
504
- explanations = []
505
- buttons = []
506
- annotated_image = image.copy()
507
- draw = ImageDraw.Draw(annotated_image)
508
- font = ImageFont.load_default()
509
-
510
- for i, (cropped_image, _, box) in enumerate(dogs):
511
- top1_prob, topk_breeds, topk_probs_percent = await predict_single_dog(cropped_image)
512
- color = color_list[i % len(color_list)]
513
- draw.rectangle(box, outline=color, width=3)
514
- draw.text((box[0], box[1]), f"Dog {i+1}", fill=color, font=font)
515
-
516
- breed = topk_breeds[0]
517
- if top1_prob >= 0.5:
518
- description = get_dog_description(breed)
519
- formatted_description = format_description(description, breed)
520
- explanations.append(f"Dog {i+1}: {formatted_description}")
521
- elif top1_prob >= 0.2:
522
- dog_explanation = f"Dog {i+1}: Top 3 possible breeds:\n"
523
- dog_explanation += "\n".join([f"{j+1}. **{breed}** ({prob} confidence)" for j, (breed, prob) in enumerate(zip(topk_breeds[:3], topk_probs_percent[:3]))])
524
- explanations.append(dog_explanation)
525
- buttons.extend([gr.update(visible=True, value=f"Dog {i+1}: More about {breed}") for breed in topk_breeds[:3]])
526
- else:
527
- explanations.append(f"Dog {i+1}: The image is unclear or the breed is not in the dataset.")
528
-
529
- final_explanation = "\n\n".join(explanations)
530
- if buttons:
531
- final_explanation += "\n\nClick on a button to view more information about the breed."
532
- return (final_explanation, annotated_image,
533
- buttons[0] if len(buttons) > 0 else gr.update(visible=False),
534
- buttons[1] if len(buttons) > 1 else gr.update(visible=False),
535
- buttons[2] if len(buttons) > 2 else gr.update(visible=False),
536
- gr.update(visible=True)) # 顯示 back 按鈕
537
- else:
538
- return final_explanation, annotated_image, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
539
-
540
- except Exception as e:
541
- return f"An error occurred: {str(e)}", None, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
542
-
543
- async def process_single_dog(image):
544
- top1_prob, topk_breeds, topk_probs_percent = await predict_single_dog(image)
545
- if top1_prob < 0.2:
546
- return "The image is unclear or the breed is not in the dataset. Please upload a clearer image of a dog.", None, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
547
-
548
- breed = topk_breeds[0]
549
- description = get_dog_description(breed)
550
-
551
- if top1_prob >= 0.5:
552
- formatted_description = format_description(description, breed)
553
- return formatted_description, image, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
554
- else:
555
- explanation = (
556
- f"The model couldn't confidently identify the breed. Here are the top 3 possible breeds:\n\n"
557
- f"1. **{topk_breeds[0]}** ({topk_probs_percent[0]} confidence)\n"
558
- f"2. **{topk_breeds[1]}** ({topk_probs_percent[1]} confidence)\n"
559
- f"3. **{topk_breeds[2]}** ({topk_probs_percent[2]} confidence)\n\n"
560
- "Click on a button to view more information about the breed."
561
- )
562
- return (explanation, image,
563
- gr.update(visible=True, value=f"More about {topk_breeds[0]}"),
564
- gr.update(visible=True, value=f"More about {topk_breeds[1]}"),
565
- gr.update(visible=True, value=f"More about {topk_breeds[2]}"),
566
- gr.update(visible=True)) # 顯示 back 按鈕
567
-
568
- def show_details(choice, previous_output):
569
- if not choice:
570
- return previous_output, gr.update(visible=True)
571
-
572
- try:
573
- breed = choice.split("More about ")[-1]
574
- description = get_dog_description(breed)
575
- return format_description(description, breed), gr.update(visible=True)
576
- except Exception as e:
577
- return f"An error occurred while showing details: {e}", gr.update(visible=True)
578
 
579
 
580
  # 介面部分
 
283
  topk_probs_percent = [f"{prob.item() * 100:.2f}%" for prob in topk_probs[0]]
284
  return top1_prob, topk_breeds, topk_probs_percent
285
 
286
+
287
+ async def detect_multiple_dogs(image, conf_threshold=0.2, iou_threshold=0.5):
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
288
  results = model_yolo(image, conf=conf_threshold, iou=iou_threshold)[0]
289
  dogs = []
290
  for box in results.boxes:
291
  if box.cls == 16: # COCO 資料集中狗的類別是 16
292
  xyxy = box.xyxy[0].tolist()
293
  confidence = box.conf.item()
294
+ cropped_image = image.crop((xyxy[0], xyxy[1], xyxy[2], xyxy[3]))
295
+ dogs.append((cropped_image, confidence, xyxy))
 
 
 
 
 
 
296
  return dogs
297
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
298
 
299
 
300
  # async def predict(image):
 
408
  # if __name__ == "__main__":
409
  # iface.launch()
410
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
411
 
412
 
413
  # 介面部分