ium_464863/predict.py

34 lines
847 B
Python

import torch
import os
import pandas as pd
import numpy as np
from NeuralNetwork import NeuralNetwork
# Load model if it exists
if os.path.exists('./models/model.pth'):
# Create model
model = torch.load('./models/model.pth')
# Load test data
test = pd.read_csv('./datasets/test.csv')
# Split data
X_test = test.drop(columns=['id', 'diagnosis']).values
y_test = test['diagnosis'].values
# Convert data to PyTorch tensors
X_test = torch.FloatTensor(X_test)
y_test = torch.FloatTensor(y_test).view(-1, 1)
# Predict
with torch.no_grad():
y_pred = model(X_test)
y_pred = np.where(y_pred >= 0.5, 1, 0)
# Save predictions to CSV
pd.DataFrame(y_pred, columns=['Prediction']).to_csv('predictions.csv', index=False)
else:
raise FileNotFoundError('Model not found')