-
Notifications
You must be signed in to change notification settings - Fork 6
/
Copy pathapp.py
56 lines (40 loc) · 1.24 KB
/
app.py
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
from flask import Flask, request, jsonify
from flask_cors import CORS, cross_origin
import pickle
import os
from src.pipeline.predict_pipeline import PredictPipeline
app = Flask(__name__)
# get the path to the models folder
model_path = "./artifacts/Best Model"
# load the model
with open(os.path.join(model_path, 'Decision Tree.pkl'), 'rb') as f:
model = pickle.load(f)
# print(model)
pred = PredictPipeline()
# Enable CORS with all origins
cors = CORS(app, resources={r"/*": {"origins": "*"}})
@app.route("/", methods=['GET'])
def index():
return "<h1>Hello World!</h1>"
@app.route('/api/predict', methods=['POST'])
@cross_origin()
def predict():
url = request.json['url']
print("URL: " + url)
transform_url = pred.transformURL(url)
transform_url = transform_url.reshape(1, -1)
# print("transform_url" , transform_url)
prediction = model.predict(transform_url)
# 'benign', 'defacement','phishing','malware'
if(prediction == 0):
res = 'benign'
elif(prediction == 1):
res = 'defacement'
elif(prediction == 2):
res = 'phishing'
else:
res = 'malware'
response = jsonify({'prediction': res})
return response
if __name__ == '__main__':
app.run()