Spaces:
Running
on
Zero
Running
on
Zero
Update app.py
Browse files
app.py
CHANGED
|
@@ -26,7 +26,6 @@ from html_templates import (
|
|
| 26 |
format_description_html,
|
| 27 |
format_single_dog_result,
|
| 28 |
format_multiple_breeds_result,
|
| 29 |
-
format_error_message,
|
| 30 |
format_unknown_breed_message,
|
| 31 |
format_not_dog_message,
|
| 32 |
format_warning_html,
|
|
@@ -240,37 +239,6 @@ def predict_single_dog(image):
|
|
| 240 |
|
| 241 |
return probabilities[0], breeds[:3], relative_probs
|
| 242 |
|
| 243 |
-
# @spaces.GPU
|
| 244 |
-
# def detect_multiple_dogs(image, conf_threshold=0.3, iou_threshold=0.55):
|
| 245 |
-
|
| 246 |
-
# results = model_manager.yolo_model(image, conf=conf_threshold,
|
| 247 |
-
# iou=iou_threshold)[0]
|
| 248 |
-
|
| 249 |
-
# dogs = []
|
| 250 |
-
# boxes = []
|
| 251 |
-
# for box in results.boxes:
|
| 252 |
-
# if box.cls == 16: # COCO dataset class for dog is 16
|
| 253 |
-
# xyxy = box.xyxy[0].tolist()
|
| 254 |
-
# confidence = box.conf.item()
|
| 255 |
-
# boxes.append((xyxy, confidence))
|
| 256 |
-
|
| 257 |
-
# if not boxes:
|
| 258 |
-
# dogs.append((image, 1.0, [0, 0, image.width, image.height]))
|
| 259 |
-
# else:
|
| 260 |
-
# nms_boxes = non_max_suppression(boxes, iou_threshold)
|
| 261 |
-
|
| 262 |
-
# for box, confidence in nms_boxes:
|
| 263 |
-
# x1, y1, x2, y2 = box
|
| 264 |
-
# w, h = x2 - x1, y2 - y1
|
| 265 |
-
# x1 = max(0, x1 - w * 0.05)
|
| 266 |
-
# y1 = max(0, y1 - h * 0.05)
|
| 267 |
-
# x2 = min(image.width, x2 + w * 0.05)
|
| 268 |
-
# y2 = min(image.height, y2 + h * 0.05)
|
| 269 |
-
# cropped_image = image.crop((x1, y1, x2, y2))
|
| 270 |
-
# dogs.append((cropped_image, confidence, [x1, y1, x2, y2]))
|
| 271 |
-
|
| 272 |
-
# return dogs
|
| 273 |
-
|
| 274 |
@spaces.GPU
|
| 275 |
def detect_multiple_dogs(image, conf_threshold=0.3, iou_threshold=0.55):
|
| 276 |
"""
|
|
@@ -374,131 +342,13 @@ def create_breed_comparison(breed1: str, breed2: str) -> dict:
|
|
| 374 |
|
| 375 |
return comparison_data
|
| 376 |
|
| 377 |
-
|
| 378 |
-
# def predict(image):
|
| 379 |
-
# """
|
| 380 |
-
# Main prediction function that handles both single and multiple dog detection.
|
| 381 |
-
|
| 382 |
-
# Args:
|
| 383 |
-
# image: PIL Image or numpy array
|
| 384 |
-
|
| 385 |
-
# Returns:
|
| 386 |
-
# tuple: (html_output, annotated_image, initial_state)
|
| 387 |
-
# """
|
| 388 |
-
|
| 389 |
-
# if image is None:
|
| 390 |
-
# return format_warning_html("Please upload an image to start."), None, None
|
| 391 |
-
|
| 392 |
-
# try:
|
| 393 |
-
# if isinstance(image, np.ndarray):
|
| 394 |
-
# image = Image.fromarray(image)
|
| 395 |
-
|
| 396 |
-
# # Detect dogs in the image
|
| 397 |
-
# dogs = detect_multiple_dogs(image)
|
| 398 |
-
# color_scheme = get_color_scheme(len(dogs) == 1)
|
| 399 |
-
|
| 400 |
-
# # Prepare for annotation
|
| 401 |
-
# annotated_image = image.copy()
|
| 402 |
-
# draw = ImageDraw.Draw(annotated_image)
|
| 403 |
-
|
| 404 |
-
# try:
|
| 405 |
-
# font = ImageFont.truetype("arial.ttf", 24)
|
| 406 |
-
# except:
|
| 407 |
-
# font = ImageFont.load_default()
|
| 408 |
-
|
| 409 |
-
# dogs_info = ""
|
| 410 |
-
|
| 411 |
-
# # Process each detected dog
|
| 412 |
-
# for i, (cropped_image, detection_confidence, box) in enumerate(dogs):
|
| 413 |
-
# color = color_scheme if len(dogs) == 1 else color_scheme[i % len(color_scheme)]
|
| 414 |
-
|
| 415 |
-
# # Draw box and label on image
|
| 416 |
-
# draw.rectangle(box, outline=color, width=4)
|
| 417 |
-
# label = f"Dog {i+1}"
|
| 418 |
-
# label_bbox = draw.textbbox((0, 0), label, font=font)
|
| 419 |
-
# label_width = label_bbox[2] - label_bbox[0]
|
| 420 |
-
# label_height = label_bbox[3] - label_bbox[1]
|
| 421 |
-
|
| 422 |
-
# # Draw label background and text
|
| 423 |
-
# label_x = box[0] + 5
|
| 424 |
-
# label_y = box[1] + 5
|
| 425 |
-
# draw.rectangle(
|
| 426 |
-
# [label_x - 2, label_y - 2, label_x + label_width + 4, label_y + label_height + 4],
|
| 427 |
-
# fill='white',
|
| 428 |
-
# outline=color,
|
| 429 |
-
# width=2
|
| 430 |
-
# )
|
| 431 |
-
# draw.text((label_x, label_y), label, fill=color, font=font)
|
| 432 |
-
|
| 433 |
-
# # Predict breed
|
| 434 |
-
# top1_prob, topk_breeds, relative_probs = predict_single_dog(cropped_image)
|
| 435 |
-
# combined_confidence = detection_confidence * top1_prob
|
| 436 |
-
|
| 437 |
-
# # Format results based on confidence with error handling
|
| 438 |
-
# try:
|
| 439 |
-
# if combined_confidence < 0.2:
|
| 440 |
-
# dogs_info += format_error_message(color, i+1)
|
| 441 |
-
# elif top1_prob >= 0.45:
|
| 442 |
-
# breed = topk_breeds[0]
|
| 443 |
-
# description = get_dog_description(breed)
|
| 444 |
-
# # Handle missing breed description
|
| 445 |
-
# if description is None:
|
| 446 |
-
# # 如果沒有描述,創建一個基本描述
|
| 447 |
-
# description = {
|
| 448 |
-
# "Name": breed,
|
| 449 |
-
# "Size": "Unknown",
|
| 450 |
-
# "Exercise Needs": "Unknown",
|
| 451 |
-
# "Grooming Needs": "Unknown",
|
| 452 |
-
# "Care Level": "Unknown",
|
| 453 |
-
# "Good with Children": "Unknown",
|
| 454 |
-
# "Description": f"Identified as {breed.replace('_', ' ')}"
|
| 455 |
-
# }
|
| 456 |
-
# dogs_info += format_single_dog_result(breed, description, color)
|
| 457 |
-
# else:
|
| 458 |
-
# # 修改format_multiple_breeds_result的調用,包含錯誤���理
|
| 459 |
-
# dogs_info += format_multiple_breeds_result(
|
| 460 |
-
# topk_breeds,
|
| 461 |
-
# relative_probs,
|
| 462 |
-
# color,
|
| 463 |
-
# i+1,
|
| 464 |
-
# lambda breed: get_dog_description(breed) or {
|
| 465 |
-
# "Name": breed,
|
| 466 |
-
# "Size": "Unknown",
|
| 467 |
-
# "Exercise Needs": "Unknown",
|
| 468 |
-
# "Grooming Needs": "Unknown",
|
| 469 |
-
# "Care Level": "Unknown",
|
| 470 |
-
# "Good with Children": "Unknown",
|
| 471 |
-
# "Description": f"Identified as {breed.replace('_', ' ')}"
|
| 472 |
-
# }
|
| 473 |
-
# )
|
| 474 |
-
# except Exception as e:
|
| 475 |
-
# print(f"Error formatting results for dog {i+1}: {str(e)}")
|
| 476 |
-
# dogs_info += format_error_message(color, i+1)
|
| 477 |
-
|
| 478 |
-
# # Wrap final HTML output
|
| 479 |
-
# html_output = format_multi_dog_container(dogs_info)
|
| 480 |
-
|
| 481 |
-
# # Prepare initial state
|
| 482 |
-
# initial_state = {
|
| 483 |
-
# "dogs_info": dogs_info,
|
| 484 |
-
# "image": annotated_image,
|
| 485 |
-
# "is_multi_dog": len(dogs) > 1,
|
| 486 |
-
# "html_output": html_output
|
| 487 |
-
# }
|
| 488 |
-
|
| 489 |
-
# return html_output, annotated_image, initial_state
|
| 490 |
-
|
| 491 |
-
# except Exception as e:
|
| 492 |
-
# error_msg = f"An error occurred: {str(e)}\n\nTraceback:\n{traceback.format_exc()}"
|
| 493 |
-
# print(error_msg)
|
| 494 |
-
# return format_warning_html(error_msg), None, None
|
| 495 |
-
|
| 496 |
|
| 497 |
@spaces.GPU
|
| 498 |
def predict(image):
|
| 499 |
"""
|
| 500 |
主要的預測函數,負責處理狗的檢測和品種辨識。
|
| 501 |
它整合了YOLO的物體檢測和專門的品種分類模型。
|
|
|
|
| 502 |
|
| 503 |
Args:
|
| 504 |
image: PIL Image 或 numpy array
|
|
|
|
| 26 |
format_description_html,
|
| 27 |
format_single_dog_result,
|
| 28 |
format_multiple_breeds_result,
|
|
|
|
| 29 |
format_unknown_breed_message,
|
| 30 |
format_not_dog_message,
|
| 31 |
format_warning_html,
|
|
|
|
| 239 |
|
| 240 |
return probabilities[0], breeds[:3], relative_probs
|
| 241 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 242 |
@spaces.GPU
|
| 243 |
def detect_multiple_dogs(image, conf_threshold=0.3, iou_threshold=0.55):
|
| 244 |
"""
|
|
|
|
| 342 |
|
| 343 |
return comparison_data
|
| 344 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 345 |
|
| 346 |
@spaces.GPU
|
| 347 |
def predict(image):
|
| 348 |
"""
|
| 349 |
主要的預測函數,負責處理狗的檢測和品種辨識。
|
| 350 |
它整合了YOLO的物體檢測和專門的品種分類模型。
|
| 351 |
+
實施雙層檢測,非狗會直接忽略.
|
| 352 |
|
| 353 |
Args:
|
| 354 |
image: PIL Image 或 numpy array
|