|
|
import torch |
|
|
import torch.nn as nn |
|
|
import torch.nn.functional as F |
|
|
from torchvision import transforms |
|
|
from PIL import Image |
|
|
|
|
|
class RiceLeafValidator(nn.Module): |
|
|
def __init__(self): |
|
|
super(RiceLeafValidator, self).__init__() |
|
|
self.conv1 = nn.Conv2d(3, 16, kernel_size=3, padding=1) |
|
|
self.conv2 = nn.Conv2d(16, 32, kernel_size=3, padding=1) |
|
|
self.conv3 = nn.Conv2d(32, 64, kernel_size=3, padding=1) |
|
|
self.pool = nn.MaxPool2d(kernel_size=2, stride=2) |
|
|
self.fc1 = nn.Linear(64 * 28 * 28, 128) |
|
|
self.fc2 = nn.Linear(128, 2) |
|
|
self.dropout = nn.Dropout(0.3) |
|
|
|
|
|
def forward(self, x): |
|
|
x = self.pool(F.relu(self.conv1(x))) |
|
|
x = self.pool(F.relu(self.conv2(x))) |
|
|
x = self.pool(F.relu(self.conv3(x))) |
|
|
x = x.view(-1, 64 * 28 * 28) |
|
|
x = F.relu(self.fc1(x)) |
|
|
x = self.dropout(x) |
|
|
x = self.fc2(x) |
|
|
return x |
|
|
|
|
|
|
|
|
data_transforms = transforms.Compose([ |
|
|
transforms.Resize((224, 224)), |
|
|
transforms.ToTensor(), |
|
|
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) |
|
|
]) |
|
|
|
|
|
def is_rice_leaf(image_path, model, device, threshold=0.7): |
|
|
model.eval() |
|
|
image = Image.open(image_path).convert('RGB') |
|
|
image_tensor = data_transforms(image).unsqueeze(0).to(device) |
|
|
with torch.no_grad(): |
|
|
outputs = model(image_tensor) |
|
|
probabilities = F.softmax(outputs, dim=1)[0] |
|
|
rice_leaf_prob = probabilities[1].item() |
|
|
return rice_leaf_prob >= threshold, rice_leaf_prob |
|
|
|
