Spaces:
Sleeping
Sleeping
| import numpy as np | |
| import tflite_runtime.interpreter as tflite | |
| from sklearn.metrics.pairwise import cosine_distances | |
| import streamlit as st | |
| import time | |
| MODEL_PATHS = { | |
| "MobileNet": "./models/mobileNet.tflite", | |
| "ResNet": "./models/resNet.tflite", | |
| } | |
| #@st.cache_resource | |
| def load_identification_model(name="MobileNet"): | |
| model = tflite.Interpreter(model_path=MODEL_PATHS[name]) | |
| return model | |
| def inference(imgs, model): | |
| if len(imgs) > 0: | |
| imgs = np.asarray(imgs).astype(np.float32) / 255 | |
| model.resize_tensor_input(model.get_input_details()[0]["index"], imgs.shape) | |
| model.allocate_tensors() | |
| model.set_tensor(model.get_input_details()[0]["index"], imgs) | |
| model.invoke() | |
| embs = [model.get_tensor(elem["index"]) for elem in model.get_output_details()] | |
| return embs[0] | |
| else: | |
| return [] | |
| def identify(embs_src, embs_gal, labels_gal, imgs_gal, thresh=None): | |
| all_dists = cosine_distances(embs_src, embs_gal) | |
| ident_names, ident_dists, ident_imgs = [], [], [] | |
| for dists in all_dists: | |
| idx_min = np.argmin(dists) | |
| if thresh and dists[idx_min] > thresh: | |
| dist = dists[idx_min] | |
| pred = None | |
| else: | |
| dist = dists[idx_min] | |
| pred = idx_min | |
| ident_names.append(labels_gal[pred] if pred is not None else "Unknown") | |
| ident_dists.append(dist) | |
| ident_imgs.append(imgs_gal[pred] if pred is not None else None) | |
| return ident_names, ident_dists, ident_imgs | |