update to last methodology of paper

app.py

@@ -5,6 +5,7 @@ from fastai.vision.all import load_learner
 from binary2image import get_size, save_file, get_binary_data
 import gradio as gr
 import numpy as np
+import pandas as pd
 import plotly.express as px
 
 from scipy import stats
@@ -12,14 +13,14 @@ import pickle
 
 ## Loading the models
 entropy_classifier = pickle.load(open('entropy_tester_classifier.pkl', 'rb'))
-model_NonObf = load_learner("model.pkl", cpu=True) # change to "model_NonObf.pkl"
-model_Shikata = load_learner("model.pkl", cpu=True) # change to "model_Shikata.pkl"
-model_XOR = load_learner("model.pkl", cpu=True) # change to "model_XOR.pkl"
+model_NonObf = load_learner("model_non_obfuscated.pkl", cpu=True) # change to "model_NonObf.pkl"
+model_Obf = load_learner("resnet18_obfuscated_samples.pkl", cpu=True) # change to "model_Shikata.pkl"
 
 def entropy_tester(bin_data):
     entropy = stats.entropy(bin_data,base=2)
     pred = entropy_classifier.predict(
-        np.array(entropy).reshape(1, -1))
+        pd.DataFrame(data = entropy.reshape(1,-1), columns=['Entropy'])
+        )
     return pred[0]
 
 def process_file(file):
@@ -35,10 +36,8 @@ def process_file(file):
     # Make prediction depending on the type of obfuscation
     if pred_entropy == "NonObf":
         prediction, _, probas = model_NonObf.predict(converted_filename)
-    elif pred_entropy == "Shikata":
-        prediction, _, probas = model_Shikata.predict(converted_filename)
-    elif pred_entropy == "XOR":
-        prediction, _, probas = model_XOR.predict(converted_filename)
+    elif pred_entropy == "Obfuscated":
+        prediction, _, probas = model_Obf.predict(converted_filename)
 
     message = f"Your file is {prediction}!!!"
     
@@ -64,7 +63,7 @@ with gr.Blocks() as demo:
             prob_dist_img = gr.Plot(label="Probability Distribution")
         with gr.Column():
             text0 = gr.Textbox(label="Type of Obfuscation")
-            converted_img = gr.Image(label="Converted Image", height=300)
+            converted_img = gr.Image(label="Converted Image", height=300, width=300)
     
     button = gr.Button(value="Process File")
     button.click(process_file, inputs=[file], outputs=[text0, text1, prob_dist_img, converted_img])

entropy_tester_classifier.pkl

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:dbb214565a0e5969901fb4707f41369040ab882d8b72531060c404e2d1cc4c80
-size 1895142
+oid sha256:b0f23a2b4a5404b6f86f896854c128b54ed69edefbdaf4b77167b03f464e0c91
+size 2008273

resnet18_obfuscated_samples.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b3a57dd13ebbcda60d1f393d5aed2f442861a41def4fc052ec4e52a6e2bd5330
+size 47955320