Spaces:
Running
on
Zero
Running
on
Zero
Update app.py
Browse files
app.py
CHANGED
@@ -654,7 +654,7 @@ def preprocess_image(image):
|
|
654 |
return transform(image).unsqueeze(0)
|
655 |
|
656 |
@spaces.GPU
|
657 |
-
|
658 |
"""
|
659 |
Predicts the dog breed using only the classifier.
|
660 |
Args:
|
@@ -687,7 +687,7 @@ async def predict_single_dog(image):
|
|
687 |
return probabilities[0], breeds[:3], relative_probs
|
688 |
|
689 |
@spaces.GPU
|
690 |
-
|
691 |
|
692 |
results = model_manager.yolo_model(image, conf=conf_threshold,
|
693 |
iou=iou_threshold)[0]
|
@@ -772,7 +772,7 @@ def create_breed_comparison(breed1: str, breed2: str) -> dict:
|
|
772 |
return comparison_data
|
773 |
|
774 |
|
775 |
-
|
776 |
"""
|
777 |
Main prediction function that handles both single and multiple dog detection.
|
778 |
|
@@ -791,7 +791,7 @@ async def predict(image):
|
|
791 |
image = Image.fromarray(image)
|
792 |
|
793 |
# Detect dogs in the image
|
794 |
-
dogs =
|
795 |
color_scheme = get_color_scheme(len(dogs) == 1)
|
796 |
|
797 |
# Prepare for annotation
|
@@ -828,7 +828,7 @@ async def predict(image):
|
|
828 |
draw.text((label_x, label_y), label, fill=color, font=font)
|
829 |
|
830 |
# Predict breed
|
831 |
-
top1_prob, topk_breeds, relative_probs =
|
832 |
combined_confidence = detection_confidence * top1_prob
|
833 |
|
834 |
# Format results based on confidence with error handling
|
|
|
654 |
return transform(image).unsqueeze(0)
|
655 |
|
656 |
@spaces.GPU
|
657 |
+
def predict_single_dog(image):
|
658 |
"""
|
659 |
Predicts the dog breed using only the classifier.
|
660 |
Args:
|
|
|
687 |
return probabilities[0], breeds[:3], relative_probs
|
688 |
|
689 |
@spaces.GPU
|
690 |
+
def detect_multiple_dogs(image, conf_threshold=0.3, iou_threshold=0.55):
|
691 |
|
692 |
results = model_manager.yolo_model(image, conf=conf_threshold,
|
693 |
iou=iou_threshold)[0]
|
|
|
772 |
return comparison_data
|
773 |
|
774 |
|
775 |
+
def predict(image):
|
776 |
"""
|
777 |
Main prediction function that handles both single and multiple dog detection.
|
778 |
|
|
|
791 |
image = Image.fromarray(image)
|
792 |
|
793 |
# Detect dogs in the image
|
794 |
+
dogs = detect_multiple_dogs(image)
|
795 |
color_scheme = get_color_scheme(len(dogs) == 1)
|
796 |
|
797 |
# Prepare for annotation
|
|
|
828 |
draw.text((label_x, label_y), label, fill=color, font=font)
|
829 |
|
830 |
# Predict breed
|
831 |
+
top1_prob, topk_breeds, relative_probs = predict_single_dog(cropped_image)
|
832 |
combined_confidence = detection_confidence * top1_prob
|
833 |
|
834 |
# Format results based on confidence with error handling
|