update since gradio was updated

app.py

@@ -1,5 +1,3 @@
-# import io
-# import base64
 from PIL import Image
 from fastai.vision.all import load_learner
 from binary2image import get_size, save_file, get_binary_data
@@ -13,18 +11,30 @@ import pickle
 
 ## Loading the models
 entropy_classifier = pickle.load(open('entropy_tester_classifier.pkl', 'rb'))
-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"
+model_NonObf = load_learner("model_non_obfuscated.pkl", cpu=True)  # change to "model_NonObf.pkl" if needed
+model_Obf = load_learner("resnet18_obfuscated_samples.pkl", cpu=True)  # change to "model_Shikata.pkl" if needed
+
 
 def entropy_tester(bin_data):
-    entropy = stats.entropy(bin_data,base=2)
+    entropy = stats.entropy(bin_data, base=2)
     pred = entropy_classifier.predict(
-        pd.DataFrame(data = entropy.reshape(1,-1), columns=['Entropy'])
-        )
+        pd.DataFrame(data=entropy.reshape(1, -1), columns=['Entropy'])
+    )
     return pred[0]
 
+
 def process_file(file):
-    greyscale_data = get_binary_data(file.name)
+    # Gradio 4 with type="filepath" returns a string path.
+    # Older Gradio could return a file object with `.name`.
+    if file is None:
+        return "No file", "Please upload a file.", None, None
+
+    if isinstance(file, str):
+        file_path = file
+    else:
+        file_path = file.name
+
+    greyscale_data = get_binary_data(file_path)
 
     pred_entropy = entropy_tester(greyscale_data)
 
@@ -38,34 +48,50 @@ def process_file(file):
         prediction, _, probas = model_NonObf.predict(converted_filename)
     elif pred_entropy == "Obfuscated":
         prediction, _, probas = model_Obf.predict(converted_filename)
+    else:
+        # If some unexpected label appears, return a safe message
+        return (
+            "Unknown",
+            "Could not determine obfuscation type.",
+            None,
+            Image.open(converted_filename),
+        )
 
     message = f"Your file is {prediction}!!!"
-    
-    if pred_entropy == "NonObf": pred_entropy = "Non-Obfuscated"
-    
+
+    if pred_entropy == "NonObf":
+        pred_entropy = "Non-Obfuscated"
+
     # Convert probas to percentages
     probas_percentage = [float(prob) * 100 for prob in probas]
-    
-    fig = px.bar(x=["Goodware", "Malware"], 
-                 y=probas_percentage, labels={'x':'Type', 'y':'Probability (%)'},
-                 height=300)
 
-    
+    fig = px.bar(
+        x=["Goodware", "Malware"],
+        y=probas_percentage,
+        labels={'x': 'Type', 'y': 'Probability (%)'},
+        height=300,
+    )
+
     return pred_entropy, message, fig, Image.open("tempfile.png")
 
+
 # Define the layout using Blocks, Row, and Column
 with gr.Blocks() as demo:
     with gr.Column(scale=1):
-        file = gr.File(label="Upload Executable File")
+        file = gr.File(label="Upload Executable File", type="filepath")
     with gr.Row():
-        with gr.Column():            
+        with gr.Column():
             text1 = gr.Textbox(label="Prediction")
             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, width=300)
-    
+
     button = gr.Button(value="Process File")
-    button.click(process_file, inputs=[file], outputs=[text0, text1, prob_dist_img, converted_img])
+    button.click(
+        process_file,
+        inputs=[file],
+        outputs=[text0, text1, prob_dist_img, converted_img],
+    )
 
 demo.launch()