File size: 2,227 Bytes
c7f0510 |
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 |
import torch
import argparse
from utils.inference_utils import preprocess_image, predict
from utils.train_utils import initialize_model
from utils.interpretability import lime_interpret_image_inference
from utils.data import CLASS_NAMES
def main():
parser = argparse.ArgumentParser(description="Image Inference")
parser.add_argument(
"--model_name",
type=str,
help="Model name (resnet, alexnet, vgg, squeezenet, densenet)",
default="resnet",
)
parser.add_argument(
"--model_weights",
type=str,
help="Path to the model weights",
default="./trained_models/pokemon_resnet.pth",
)
parser.add_argument(
"--image_path",
type=str,
help="Path to the image",
default="./pokemonclassification/PokemonData/Chansey/57ccf27cba024fac9531baa9f619ec62.jpg",
)
parser.add_argument(
"--num_classes", type=int, help="Number of classes", default=150
)
parser.add_argument(
"--lime_interpretability",
action="store_true",
help="Whether to run interpretability or not",
)
parser.add_argument(
"--classify",
action="store_true",
help="Whether to classify the image when saving the lime filter")
args = parser.parse_args()
if args.lime_interpretability:
assert (
args.model_name == "resnet"
), "Interpretability is only supported for ResNet model for now"
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# Initialize the model
model = initialize_model(args.model_name, args.num_classes)
model = model.to(device)
# Load the model weights
model.load_state_dict(torch.load(args.model_weights, map_location=device))
# Preprocess the image
image = preprocess_image(args.image_path, (224, 224)).to(device)
# Perform inference
preds = torch.max(predict(model, image), 1)[1]
print(f"Predicted class: {CLASS_NAMES[preds.item()]}")
if args.lime_interpretability:
lime_interpret_image_inference(args, model, image, device)
if __name__ == "__main__":
main()
|