smallGroupProject / models /mock_model.py
Mert Yerlikaya
Add feature-rich Gradio UI with mock model
505fc99
raw
history blame
2.87 kB
"""
Mock model for UI development
This allows parallel development while the actual model is being trained
"""
import torch
import torch.nn as nn
import numpy as np
class MockPlantDiseaseModel(nn.Module):
"""
Mock CNN model that mimics the structure of a real plant disease classifier
Returns realistic-looking predictions for UI testing
"""
def __init__(self, num_classes=39):
super(MockPlantDiseaseModel, self).__init__()
# Simple architecture that matches expected input/output
self.features = nn.Sequential(
nn.Conv2d(3, 32, kernel_size=3, padding=1),
nn.ReLU(),
nn.MaxPool2d(2),
nn.Conv2d(32, 64, kernel_size=3, padding=1),
nn.ReLU(),
nn.MaxPool2d(2),
nn.AdaptiveAvgPool2d((1, 1))
)
self.classifier = nn.Sequential(
nn.Flatten(),
nn.Linear(64, num_classes)
)
self.num_classes = num_classes
def forward(self, x):
"""
Forward pass that returns realistic probabilities
"""
x = self.features(x)
x = self.classifier(x)
# Add some controlled randomness to make predictions look realistic
# In a real model, this would be learned weights
return x
def create_mock_predictions(class_names):
"""
Create realistic-looking mock predictions
Returns a dict with class names and probabilities
"""
num_classes = len(class_names)
# Create random probabilities that sum to 1
# Give higher weight to a few "predicted" classes
logits = np.random.randn(num_classes)
logits[np.random.randint(0, num_classes)] += 3 # Make one class likely
logits[np.random.randint(0, num_classes)] += 1.5 # Make another somewhat likely
# Convert to probabilities using softmax
probs = np.exp(logits) / np.sum(np.exp(logits))
# Create prediction dict
predictions = {name: float(prob) for name, prob in zip(class_names, probs)}
return predictions
def get_mock_model():
"""
Returns a mock model instance
"""
model = MockPlantDiseaseModel(num_classes=39)
model.eval() # Set to evaluation mode
return model
if __name__ == "__main__":
# Test the mock model
print("Testing mock model...")
model = get_mock_model()
# Test with random input
test_input = torch.randn(1, 3, 256, 256)
with torch.no_grad():
output = model(test_input)
print(f"Output shape: {output.shape}")
print(f"Sample logits: {output[0][:5]}")
# Test mock predictions
from config import CLASS_NAMES
predictions = create_mock_predictions(CLASS_NAMES)
top_5 = sorted(predictions.items(), key=lambda x: x[1], reverse=True)[:5]
print("\nTop 5 predictions:")
for name, prob in top_5:
print(f" {name}: {prob:.4f}")