Spaces:
Sleeping
Sleeping
| import tensorflow as tf | |
| tf.config.set_visible_devices([], 'GPU') | |
| # gpu_devices = tf.config.experimental.list_physical_devices('GPU') | |
| # if gpu_devices: | |
| # tf.config.experimental.set_memory_growth(gpu_devices[0], True) | |
| # else: | |
| # print(f"TensorFlow device: {gpu_devices}") | |
| import os | |
| import numpy as np | |
| import keras | |
| from PIL import Image | |
| from keras_cv_attention_models import beit | |
| import matplotlib.pyplot as plt | |
| import tensorflow as tf | |
| from tensorflow.keras.layers import Dense, GlobalAveragePooling2D | |
| from typing import Tuple | |
| #from huggingface_hub import snapshot_download | |
| from labels import lookup_140 | |
| def get_triplet_model_beit(input_shape = (600, 600, 3), | |
| embedding_units = 256, | |
| embedding_depth = 2, | |
| n_classes = 19,backbone_name ='Beit'): | |
| backbone_class = beit.BeitBasePatch16(input_shape=input_shape, pretrained="imagenet21k-ft1k") | |
| backbone_class = tf.keras.Model(backbone_class.input, backbone_class.layers[-2].output) | |
| #features = GlobalAveragePooling2D()(backbone_class.output) | |
| embedding_head = backbone_class.output | |
| for embed_i in range(embedding_depth): | |
| embedding_head = Dense(embedding_units, activation="relu" if embed_i < embedding_depth-1 else "linear")(embedding_head) | |
| embedding_head = tf.nn.l2_normalize(embedding_head, -1, epsilon=1e-5) | |
| logits_head = Dense(n_classes)(backbone_class.output) | |
| model = tf.keras.Model(backbone_class.input, [embedding_head, logits_head]) | |
| model.compile(loss='cce',metrics=['accuracy']) | |
| #model.summary() | |
| return model | |
| load_size = 600 | |
| crop_size = 600 | |
| def _clever_crop(img: tf.Tensor, | |
| target_size: Tuple[int]=(128,128), | |
| grayscale: bool=False | |
| ) -> tf.Tensor: | |
| """[summary] | |
| Args: | |
| img (tf.Tensor): [description] | |
| target_size (Tuple[int], optional): [description]. Defaults to (128,128). | |
| grayscale (bool, optional): [description]. Defaults to False. | |
| Returns: | |
| tf.Tensor: [description] | |
| """ | |
| maxside = tf.math.maximum(tf.shape(img)[0],tf.shape(img)[1]) | |
| minside = tf.math.minimum(tf.shape(img)[0],tf.shape(img)[1]) | |
| new_img = img | |
| if tf.math.divide(maxside,minside) > 1.2: | |
| repeating = tf.math.floor(tf.math.divide(maxside,minside)) | |
| new_img = img | |
| if tf.math.equal(tf.shape(img)[1],minside): | |
| for _ in range(int(repeating)): | |
| new_img = tf.concat((new_img, img), axis=1) | |
| if tf.math.equal(tf.shape(img)[0],minside): | |
| for _ in range(int(repeating)): | |
| new_img = tf.concat((new_img, img), axis=0) | |
| new_img = tf.image.rot90(new_img) | |
| else: | |
| new_img = img | |
| repeating = 0 | |
| img = tf.image.resize(new_img, target_size) | |
| if grayscale: | |
| img = tf.image.rgb_to_grayscale(img) | |
| img = tf.image.grayscale_to_rgb(img) | |
| return img,repeating | |
| def preprocess(img,size=384): | |
| img = np.array(img, np.float32) / 255.0 | |
| img = tf.image.resize(img, (size, size)) | |
| return np.array(img, np.float32) | |
| def select_top_n(preds,n=10): | |
| top_n = np.argsort(preds)[-n:][::-1] | |
| return top_n | |
| def parse_results(top_n,logits): | |
| results = {} | |
| for n in top_n: | |
| label = lookup_140[n] | |
| results[label] = float(logits[n]) | |
| return results | |
| def inference_resnet_embedding_beit(x,model,size=576,n_classes=142,n_top=10): | |
| cropped = _clever_crop(x,(size,size))[0] | |
| prep = preprocess(cropped,size=size) | |
| embedding = model.predict(np.array([prep]))[0][0] | |
| return embedding | |
| def inference_resnet_finer_beit(x,model,size=576,n_classes=142,n_top=10): | |
| cropped = _clever_crop(x,(size,size))[0] | |
| prep = preprocess(cropped,size=size) | |
| logits = tf.nn.softmax(model.predict(np.array([prep]))[1][0]).cpu().numpy() | |
| top_n = select_top_n(logits,n=n_top) | |
| return parse_results(top_n,logits) |