22 lines
730 B
Python
22 lines
730 B
Python
from django.db import models
|
|
|
|
# Create your models here.
|
|
|
|
class MalwarePrediction(models.Model):
|
|
MODEL_TYPE_CHOICES = [
|
|
(1, 'KNeighborsClassifier'),
|
|
(2, 'RandomForestClassifier'),
|
|
(3, 'XGBClassifier'),
|
|
(4, 'SGDClassifier'),
|
|
]
|
|
|
|
process_name = models.CharField(max_length=255)
|
|
process_class = models.CharField(max_length=50)
|
|
probability_of_malware = models.FloatField()
|
|
predicted_malware = models.CharField(max_length=50)
|
|
model_type = models.IntegerField(choices=MODEL_TYPE_CHOICES, default=1) # New field to store model type
|
|
|
|
def __str__(self):
|
|
return f"{self.process_name} - {self.predicted_malware} - {self.get_model_type_display()}"
|
|
|