Spaces:
Sleeping
Sleeping
File size: 3,196 Bytes
a18e884 bc51538 a18e884 bc51538 96d6137 bc51538 a18e884 bc51538 96d6137 a18e884 bc51538 a18e884 96d6137 a18e884 bc51538 a18e884 bc51538 a18e884 bc51538 a18e884 96d6137 a18e884 96d6137 bc51538 37bdfa4 6745426 a18e884 bc51538 a18e884 96d6137 a18e884 bc51538 a18e884 96d6137 a18e884 bc51538 a18e884 bc51538 a18e884 96d6137 bc51538 a18e884 96d6137 | 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 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 | import argparse
import numpy as np
from pathlib import Path
from PIL import Image
import tensorflow as tf
import numpy as np
import torch
from torchvision import transforms
from model_pytorch import SaraCNN
CLASS_NAMES = ["buildings", "forest", "glacier", "mountain", "sea", "street"]
IMAGE_SIZE = (150, 150)
MEAN = [0.485, 0.456, 0.406]
STD = [0.229, 0.224, 0.225]
# PyTorch Prediction
def predict_pytorch(image_path: str, model_path: str = "sara_model.pth") -> dict:
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
checkpoint = torch.load(model_path, map_location=device)
class_names = checkpoint.get("class_names", CLASS_NAMES)
model = SaraCNN(num_classes=len(class_names))
model.load_state_dict(checkpoint["model_state"])
model.to(device)
model.eval()
transform = transforms.Compose([
transforms.Resize(IMAGE_SIZE),
transforms.ToTensor(),
transforms.Normalize(MEAN, STD),
])
img = Image.open(image_path).convert("RGB")
tensor = transform(img).unsqueeze(0).to(device)
with torch.no_grad():
logits = model(tensor)
probs = torch.softmax(logits, dim=1).squeeze().cpu().numpy()
idx = int(np.argmax(probs))
pred_class = class_names[idx]
confidence = float(probs[idx])
return {
"predicted_class": pred_class,
"confidence": round(confidence * 100, 2),
"all_probabilities": {
cls: round(float(p) * 100, 2)
for cls, p in zip(class_names, probs)
},
}
# TensorFlow Prediction
def predict_tensorflow(image_path: str, model_path: str = "sara_model.keras") -> dict:
model = tf.keras.models.load_model(model_path, compile=False)
img = tf.keras.utils.load_img(image_path, target_size=(150, 150))
arr = tf.keras.utils.img_to_array(img) / 255.0
arr = np.expand_dims(arr, axis=0)
probs = model.predict(arr, verbose=0)[0]
idx = int(np.argmax(probs))
pred_class = CLASS_NAMES[idx]
confidence = float(probs[idx])
return {
"predicted_class": pred_class,
"confidence": round(confidence * 100, 2),
"all_probabilities": {
cls: round(float(p) * 100, 2)
for cls, p in zip(CLASS_NAMES, probs)
},
}
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--model", required=True, choices=["pytorch", "tensorflow"])
parser.add_argument("--image", required=True)
parser.add_argument("--model_path", default=None)
args = parser.parse_args()
if args.model == "pytorch":
path = args.model_path or "sara_model.pth"
result = predict_pytorch(args.image, model_path=path)
else:
path = args.model_path or "sara_model.keras"
result = predict_tensorflow(args.image, model_path=path)
print(f"\n Classe : {result['predicted_class']}")
print(f" Confidence : {result['confidence']}%")
print("\n All probabilities :")
for cls, prob in sorted(result["all_probabilities"].items(), key=lambda x: -x[1]):
bar = " " * int(prob / 5)
print(f" {cls:<12} {prob:6.2f}% {bar}")
|