Spaces:
Runtime error
Runtime error
File size: 2,692 Bytes
58762a0 4629cf6 58762a0 | 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 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 | import os
import torch
import torch.nn as nn
from torchvision import models as torchvision_models
import tensorflow as tf
def build_tree_vs_nontree():
model = torchvision_models.resnet50(weights=None)
in_features = model.fc.in_features
model.fc = nn.Sequential(
nn.Linear(in_features, 256),
nn.ReLU(),
nn.Dropout(0.3),
nn.Linear(256, 2)
)
return model
def build_mango_stage():
model = torchvision_models.mobilenet_v2(weights=None)
in_features = model.classifier[1].in_features
model.classifier = nn.Sequential(
nn.Dropout(0.5),
nn.Linear(in_features, 256),
nn.ReLU(),
nn.Dropout(0.3),
nn.Linear(256, 128),
nn.ReLU(),
nn.Dropout(0.2),
nn.Linear(128, 4)
)
return model
def build_gum_stage():
model = torchvision_models.mobilenet_v2(weights=None)
in_features = model.classifier[1].in_features
model.classifier = nn.Sequential(
nn.Dropout(0.3),
nn.Linear(in_features, 128),
nn.ReLU(),
nn.Dropout(0.2),
nn.Linear(128, 4)
)
return model
def load_pytorch_model(build_fn, weights_path, device):
if not os.path.exists(weights_path):
raise FileNotFoundError(f"Model not found: {weights_path}")
model = build_fn()
state_dict = torch.load(weights_path, map_location=device)
if isinstance(state_dict, dict) and "model_state" in state_dict:
state_dict = state_dict["model_state"]
model.load_state_dict(state_dict)
model.to(device)
model.eval()
return model
def load_keras_model(weights_path):
if not os.path.exists(weights_path):
raise FileNotFoundError(f"Model not found: {weights_path}")
try:
return tf.keras.models.load_model(weights_path, compile=False)
except Exception:
import h5py
with h5py.File(weights_path, 'r') as f:
model_config = f.attrs.get('model_config')
model = tf.keras.models.model_from_json(
model_config,
custom_objects=None
)
model.load_weights(weights_path)
return model
def load_all_models():
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
return {
"tree_vs_nontree": load_pytorch_model(build_tree_vs_nontree, "models/tree_vs_nontree.pt", device),
"species": load_keras_model("models/best_mobilenetv2_model.h5"),
"mango_stage": load_pytorch_model(build_mango_stage, "models/mobilenetv2_tree_classifier.pth", device),
"gum_stage": load_pytorch_model(build_gum_stage, "models/gum_stage.pth", device),
"device": device
}
|