File size: 3,183 Bytes
bc8984b |
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 |
# Function to make predictions from image
import torch
import os
from PIL import Image
from torchvision import models, transforms
import torch.nn as nn
# Define model at module level
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = None
def load_model():
global model
if model is None:
model = BananaNet().to(device)
model_path = os.path.join(os.path.dirname(os.path.dirname(__file__)), 'best_model.pth')
model.load_state_dict(torch.load(model_path, map_location=device))
model.eval()
# Original function, renamed for flexibility
def predict_from_image_full(model, image_path):
"""
Predict seed count and curvature from a banana image
Args:
model: Trained BananaNet model
image_path: Path to the banana image
Returns:
dict: Predictions for seed count and curvature
"""
image_tensor = preprocess_image(image_path)
model.eval()
with torch.no_grad():
predictions = model(image_tensor)
seed_count = int(round(predictions[0][0].item()))
curvature = round(predictions[0][1].item(), 1)
return {
'seeds': seed_count,
'curvature': curvature
}
# Wrapper for compatibility with app.py
def predict_from_image(image_path):
global model
load_model() # Ensure model is loaded
image_tensor = preprocess_image(image_path).to(device)
with torch.no_grad():
predictions = model(image_tensor)
seed_count = int(round(predictions[0][0].item()))
curvature = round(predictions[0][1].item(), 1)
return {
'seeds': seed_count,
'curvature': curvature
}
def preprocess_image(image_path):
transform = transforms.Compose([
transforms.Resize((224, 224)),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[
0.229, 0.224, 0.225])
])
image = Image.open(image_path).convert('RGB')
return transform(image).unsqueeze(0)
class BananaNet(nn.Module):
def __init__(self):
super(BananaNet, self).__init__()
self.base_model = models.resnet18(pretrained=True)
num_features = self.base_model.fc.in_features
self.base_model.fc = nn.Identity()
self.regression_head = nn.Sequential(
nn.Linear(num_features, 512),
nn.ReLU(),
nn.Dropout(0.3),
nn.Linear(512, 256),
nn.ReLU(),
nn.Dropout(0.3),
nn.Linear(256, 2)
)
def forward(self, x):
features = self.base_model(x)
output = self.regression_head(features)
return output
# Load model and weights
model = BananaNet()
model_path = os.path.join(os.path.dirname(__file__), '..', 'best_model.pth')
if os.path.exists(model_path):
model.load_state_dict(torch.load(model_path, map_location="cpu"))
print("Loaded trained model weights from best_model.pth")
else:
print("Warning: best_model.pth not found, using untrained model.")
model.eval()
|