DawnC commited on
Commit
45473c9
1 Parent(s): c2d5142

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +5 -5
app.py CHANGED
@@ -654,7 +654,7 @@ def preprocess_image(image):
654
  return transform(image).unsqueeze(0)
655
 
656
  @spaces.GPU
657
- async def predict_single_dog(image):
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
- async 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,7 +772,7 @@ def create_breed_comparison(breed1: str, breed2: str) -> dict:
772
  return comparison_data
773
 
774
 
775
- async def predict(image):
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 = await detect_multiple_dogs(image)
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 = await predict_single_dog(cropped_image)
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