update to use object detection for heads and pass off classification. Use ImageAnnotation result to annotate detections

.gitignore

@@ -1 +1,3 @@
 .DS_Store
+.gradio
+*~

app.py

@@ -2,28 +2,54 @@ import gradio as gr
 import tensorflow as tf
 import numpy as np
 import requests
+import torch
+
 from huggingface_hub import snapshot_download
+from huggingface_hub import hf_hub_download
 
 label_count = 51
+def readLines(filename):
+    with open(filename,'r') as f:
+        return(f.read().splitlines())
+
 labels = ["missing"] * label_count
-with open('labels.txt','r') as f:
-  labels = f.read().splitlines()
+labels = readLines('labels.txt')
+valid =  readLines('valid.txt')
 
-model_path = extractor_path = snapshot_download(repo_id="eshieh2/jaguarid_pantanal")
+model_path = snapshot_download(repo_id="eshieh2/jaguarid_pantanal")
 model = tf.saved_model.load(f"{model_path}/saved_model")
 serving = model.signatures['serving_default']
-#model = tf.keras.models.load_model(f"{model_path}/saved_model")
-def classify_image(inp):
-  inp = inp.resize((480,480))
-  inp = np.array(inp)
-  inp = np.reshape(inp,(-1, 480, 480, 3)).astype(np.float32)
-  inp = np.divide(inp,255.0)
-  #prediction = model.predict(inp).flatten()
-  prediction = serving(tf.convert_to_tensor(inp))['model']
-  prediction = tf.squeeze(prediction)
-  return {labels[i]: float(prediction[i]) for i in range(label_count)}
+
+detector_path = hf_hub_download(repo_id= "eshieh2/jaguarhead",
+                                             filename = "jaguarheadv5.pt")
+detector = torch.hub.load('ultralytics/yolov5', 'custom', path = detector_path)
+
+def classify_image(in_image):
+  width,height = in_image.size
+  heads = detector(in_image)
+  masks = [] # tuple of box coords and string
+  for head in heads.xyxy[0]:
+    x,y,x2,y2,pct,cls = head.numpy()
+    w = x2 - x
+    h = y2 - y
+    inp = in_image.crop((x,y,x2,y2))
+    inp = inp.resize((480,480))
+    inp = np.array(inp)
+    inp = np.reshape(inp,(-1, 480, 480, 3)).astype(np.float32)
+    inp = np.divide(inp,255.0)
+    prediction = serving(tf.convert_to_tensor(inp))['model']
+    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 max_key.lower() in valid:
+      masks.append((rect,f"{max_key}:{pct}"))
+    else:
+      masks.append((rect,f"unknown",))
+  return (in_image,masks)
 
 image = gr.Image(type='pil')
-label = gr.Label(num_top_classes=3)
+output = gr.AnnotatedImage()
 
-gr.Interface(fn=classify_image, inputs=image, outputs=label, examples = [["kasimir.jpg","guaraci.jpg"]]).launch()
+gr.Interface(fn=classify_image, inputs=image, outputs=output, examples = [["kasimir.jpg"],["guaraci.jpg"],["marcela.jpg"]]).launch()

requirements.txt

@@ -1 +1,2 @@
-tensorflow
+tensorflow
+torch

valid.txt

@@ -0,0 +1,8 @@
+guaraci
+marcela
+margo
+ti
+patricia
+bernard
+kasimir
+medrosa