add more to valid list. Allow top3. fix confidence percentages.

app.py

@@ -1,10 +1,11 @@
 import os
-os.environ['TF_USE_LEGACY_KERAS']='1'
 import gradio as gr
 import tensorflow as tf
 import numpy as np
 import requests
 import torch
+import heapq
+from PIL import Image
 
 from huggingface_hub import snapshot_download
 from huggingface_hub import hf_hub_download
@@ -28,7 +29,11 @@ detector_path = hf_hub_download(repo_id= "eshieh2/jaguarhead",
                                              filename = "jaguarheadv5.pt")
 detector = torch.hub.load('ultralytics/yolov5', 'custom', path = detector_path)
 
+topk = 3
+
 def classify_image(in_image):
+  if in_image is None:
+    return None
   width,height = in_image.size
   heads = detector(in_image)
   masks = [] # tuple of box coords and string
@@ -37,7 +42,7 @@ def classify_image(in_image):
     w = x2 - x
     h = y2 - y
     inp = in_image.crop((x,y,x2,y2))
-    inp = inp.resize((480,480))
+    inp = inp.resize((480,480),Image.BILINEAR)
     inp = np.array(inp)
     inp = np.reshape(inp,(-1, 480, 480, 3)).astype(np.float32)
     inp = np.divide(inp,255.0)
@@ -45,12 +50,24 @@ def classify_image(in_image):
     prediction = tf.squeeze(prediction)
     pred = {labels[i]: float(prediction[i]) for i in range(label_count)}
     #print(pred)
-    max_key = max(pred, key=pred.get)
     rect = (int(x),int(y),int(x2),int(y2))
-    if show_all or max_key.lower() in valid:
-      masks.append((rect,f"{max_key}:{pct}"))
+    if topk is not None:
+      top = heapq.nlargest(topk,pred,key=pred.get)
+      label = ''
+      for t in top:
+        if show_all or t.lower() in valid:
+          if len(label) != 0:
+            label += ", "
+          label += f"{t}:{pred[t]:.3f}"
+      if len(label)==0:
+        label = 'unknown'
+      masks.append((rect,label))
     else:
-      masks.append((rect,f"unknown",))
+      max_key = max(pred, key=pred.get)
+      if show_all or max_key.lower() in valid:
+            masks.append((rect,f"{max_key}:{pred[max_key]}"))
+      else:
+            masks.append((rect,f"unknown",))
   return (in_image,masks)
 
 image = gr.Image(type='pil')

valid.txt

@@ -1,7 +1,13 @@
 bagua
 guaraci
+kasimir
+manath
+margo
 marcela
 medrosa
+ousado
 oxum
 patricia
+saseka
 ti
+tusk