Upload 4 files

Dockerfile

@@ -0,0 +1,9 @@
+FROM python:3.11-slim
+
+WORKDIR /usr/src/app
+COPY . .
+RUN pip install --no-cache-dir -r requirements.txt
+EXPOSE 7860
+ENV GRADIO_SERVER_NAME="0.0.0.0"
+
+CMD ["python", "app.py"]

app.py

@@ -0,0 +1,81 @@
+import os
+import numpy as np
+import torch
+import torch.nn.functional as F
+import torchvision.transforms as T
+import timm
+from PIL import Image
+import gradio as gr
+
+# hyperparameters
+device = torch.device("cpu")
+input_width, input_height = 224, 224
+
+# load ear detector
+ear_detector = torch.hub.load("ultralytics/yolov5", "custom", path=os.path.join(os.path.dirname(__file__), "weights", "ear_YOLOv5_n.pt"))
+ear_detector.to(device)
+
+# initialize model
+model = timm.create_model("hf-hub:BVRA/MegaDescriptor-T-224", pretrained=True, num_classes=0)
+
+# load state dict containing miscellaneous state or just the model weights
+state_dict = torch.load("weights.pt", map_location=device)
+if "optimizer" in state_dict:
+    model.load_state_dict(state_dict["model"])
+else:
+    model.load_state_dict(state_dict)
+
+model.to(device)
+model.eval()
+
+
+transforms = T.Compose([
+    T.Resize([input_height, input_width]),
+    T.ToTensor(),
+    T.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
+])
+
+database = np.load("database.npz", allow_pickle=True)
+features, identities = database["features"], database["identities"]
+features = torch.from_numpy(features).to(device)
+
+
+def strings2ints(a):
+    idx = {v: i for i, v in enumerate(set([*a]))}
+    return torch.Tensor([idx[e] for e in a]).to(dtype=torch.int64), {v: k for k, v in idx.items()}
+
+def predict(image):
+    with torch.inference_mode():
+
+        output = ear_detector(image)
+        n_preds = len(output.pred[0].tolist())
+        if n_preds == 0:
+            return "Error: Unable to detect elephant ears"
+
+        xyxy = output.xyxy[0].tolist()
+        
+        noncenterness = [(image.width - (xyxy[i][0] + xyxy[i][2] / 2)) ** 2 + (image.height - (xyxy[i][1] + xyxy[i][3] / 2)) ** 2 for i in range(n_preds)]
+        centermost_idx = np.argmin(noncenterness)
+        
+        image = image.crop(tuple(output.xyxy[0].tolist()[centermost_idx][:4]))
+
+        if output.pred[0].tolist()[centermost_idx][-1] >= 0.5:
+            image = image.transpose(Image.FLIP_LEFT_RIGHT)
+
+        image = transforms(image).unsqueeze(0).to(device)
+
+        embedding = model(image)
+        similarity = torch.matmul(F.normalize(embedding), F.normalize(features).T)
+
+        similarity_sorted_idx = torch.argsort(similarity[0], descending=True).cpu().numpy().reshape(-1)
+        candidates = identities.reshape(-1)[similarity_sorted_idx].tolist()
+        candidates_similarity = similarity[0, similarity_sorted_idx].tolist()
+
+    return f"We are about {max(0, candidates_similarity[0]):.0%} confident that this elephant is {'_'.join(candidates[0].split('_')[:-1])}"
+
+
+gr.Interface(
+    fn=predict,
+    inputs=gr.Image(type="pil"),
+    outputs=gr.Label(),
+).launch(share=True)

database.npz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1985b902ee2826f7215bced0c330c0a6b0bda13559779071cfb34a9ee0342c03
+size 23072727

weights.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:da82dcf2c50dec71cafefd803cc8d078cdcf9226cbf571b7adde0dd4b14c6e7a
+size 224710970