Spaces:
Sleeping
Sleeping
File size: 1,481 Bytes
84d0c9e 57d41d5 84d0c9e |
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 |
import torch
from torchvision import transforms
from PIL import Image
import os
from src.model import TrashNetClassifier
from src import config
def load_model(model_path, device):
model = TrashNetClassifier()
model.load_state_dict(torch.load(model_path, map_location=device))
model.eval()
model.to(device)
return model
def preprocess_image(image_path, image_size):
transform = transforms.Compose([
transforms.Resize((image_size, image_size)),
transforms.ToTensor(),
transforms.Normalize([0.5]*3, [0.5]*3)
])
image = Image.open(image_path).convert("RGB")
return transform(image).unsqueeze(0)
def predict_image(model, image_tensor, class_names, device):
image_tensor = image_tensor.to(device)
with torch.no_grad():
outputs = model(image_tensor)
probs = torch.softmax(outputs, dim=1)
pred_index = torch.argmax(probs, dim=1).item()
pred_label = class_names[pred_index]
confidence = probs[0][pred_index].item()
return pred_label, confidence
def run_inference(image_path):
device = config.DEVICE
class_names = sorted(os.listdir(os.path.join(config.DATA_DIR, "train")))
model = load_model(config.MODEL_SAVE_PATH, device)
image_tensor = preprocess_image(image_path, config.IMAGE_SIZE)
label, confidence = predict_image(model, image_tensor, class_names, device)
print(f"Prediction: {label} ({confidence*100:.2f}%)")
return label, confidence
|