Spaces:
Sleeping
Sleeping
File size: 1,606 Bytes
3ca686a b8aa51a 3ca686a b8aa51a 3ca686a b8aa51a d662e61 b8aa51a d662e61 3ca686a b8aa51a 3ca686a b8aa51a d662e61 b8aa51a 3ca686a b8aa51a 3ca686a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 |
import os
import shutil
import logging
import tensorflow as tf
from tensorflow.keras.layers import Layer
from huggingface_hub import snapshot_download
# Model config
REPO_ID = "can-org/AI-VS-HUMAN-IMAGE-classifier"
MODEL_DIR = "./IMG_Models"
WEIGHTS_PATH = os.path.join(MODEL_DIR, "latest-my_cnn_model.h5")
# Device info (for logging)
gpus = tf.config.list_physical_devices("GPU")
device = "cuda" if gpus else "cpu"
# Global model reference
_model_img = None
# Custom layer used in the model
class Cast(Layer):
def call(self, inputs):
return tf.cast(inputs, tf.float32)
def warmup():
global _model_img
download_model_repo()
_model_img = load_model()
logging.info("Image model is ready.")
def download_model_repo():
if os.path.exists(MODEL_DIR) and os.path.isdir(MODEL_DIR):
logging.info("Image model already exists, skipping download.")
return
snapshot_path = snapshot_download(repo_id=REPO_ID)
os.makedirs(MODEL_DIR, exist_ok=True)
shutil.copytree(snapshot_path, MODEL_DIR, dirs_exist_ok=True)
def load_model():
global _model_img
if _model_img is not None:
return _model_img
print(f"{'GPU detected' if device == 'cuda' else 'No GPU detected'}, loading model on {device.upper()}.")
_model_img = tf.keras.models.load_model(
WEIGHTS_PATH, custom_objects={"Cast": Cast}
)
print("Model input shape:", _model_img.input_shape)
return _model_img
def get_model():
global _model_img
if _model_img is None:
download_model_repo()
_model_img = load_model()
return _model_img
|