Spaces:
Running
on
Zero
Running
on
Zero
Update app.py
Browse files
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 |
-
|
287 |
-
|
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 |
-
|
323 |
-
|
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 |
# 介面部分
|