Update app.py

app.py

@@ -6,6 +6,7 @@ import numpy as np
 import os
 import shutil
 from huggingface_hub import snapshot_download
+import cv2
 
 MAL_CLASSES = ['Adialer.C', 'Agent.FYI', 'Allaple.A', 'Allaple.L', 'Alueron.gen!J', 'Autorun.K', 'C2LOP.P', 'C2LOP.gen!g', 'Dialplatform.B', 'Dontovo.A', 'Fakerean', 'Instantaccess', 'Lolyda.AA1', 'Lolyda.AA2', 'Lolyda.AA3', 'Lolyda.AT', 'Malex.gen!J', 'Obfuscator.AD', 'Rbot!gen', 'Skintrim.N', 'Swizzor.gen!E', 'Swizzor.gen!I', 'VB.AT', 'Wintrim.BX', 'Yuner.A']
 UPLOAD_DIR = "uploads"
@@ -61,3 +62,51 @@ def analyse(file_name: str):
     predicted_label = model.predict(img_array)
     
     return {"result": MAL_CLASSES[np.argmax(predicted_label)]}
+
+def convert_binary_to_grayscale_image(binary_file_path, output_image_path, height=None, width=None):
+    with open(binary_file_path, "rb") as bin_file:
+        binary_data = bin_file.read()
+
+    grayscale_image = np.frombuffer(binary_data, dtype=np.uint8)
+
+    total_pixels = len(grayscale_image)
+    if height is None and width is None:
+        raise ValueError("Either height or width must be specified.")
+    elif height is None:
+        height = total_pixels // width
+    elif width is None:
+        width = total_pixels // height
+
+    if height * width != total_pixels:
+        raise ValueError(
+            f"The binary file size ({total_pixels}) is not compatible with the specified dimensions ({height}x{width})."
+        )
+
+    grayscale_image = grayscale_image.reshape((height, width))
+
+    # Save
+    cv2.imwrite(output_image_path, grayscale_image)
+
+    print(f"Grayscale image saved to {output_image_path} with dimensions ({height}, {width})")
+    return grayscale_image
+
+@app.get("/analysebin/{file_name}")
+def analyse_bin(file_name: str):
+    download_dir = snapshot_download("GranularFireplace/malware")
+    print(download_dir)
+    print(os.path.join(download_dir, 'model_v2_with_weight.keras'))
+
+    convert_binary_to_grayscale_image(os.path.join(UPLOAD_DIR, file_name), "image.png", 256)
+    
+    img = keras.preprocessing.image.load_img(
+        "image.png", target_size=(64, 64)
+    )
+    
+    img_array = keras.preprocessing.image.img_to_array(tf.image.rgb_to_grayscale(img))
+    
+    img_array = tf.expand_dims(img_array, 0) # Create a batch
+    
+    model = tf.keras.models.load_model(os.path.join(download_dir, 'model_v2_with_weight.keras'))
+    predicted_label = model.predict(img_array)
+    
+    return {"result": MAL_CLASSES[np.argmax(predicted_label)]}