added all files

app.py

@@ -0,0 +1,12 @@
+"""
+Entry point for Hugging Face Spaces deployment.
+Hugging Face Spaces expects an `app.py` file in the root directory.
+"""
+
+import gradio as gr
+from src.ui.app import app
+
+if __name__ == "__main__":
+    app.launch(
+        theme=gr.themes.Monochrome()
+    )

models/engine_a_model.pkl

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

models/engine_b_classes.json

@@ -0,0 +1 @@
+["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", "VB.AT", "Wintrim.BX", "Yuner.A"]

models/engine_b_model.pth

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

requirements.txt

@@ -0,0 +1,12 @@
+fpdf2>=2.8.7
+gradio>=6.15.2
+lightgbm>=4.6.0
+matplotlib>=3.10.9
+numpy>=2.4.6
+pandas>=3.0.3
+pefile>=2024.8.26
+pillow>=12.2.0
+scikit-learn>=1.8.0
+seaborn>=0.13.2
+torch>=2.12
+torchvision>=0.27

src/engine_a/inference.py

@@ -0,0 +1,207 @@
+"""
+Inference script for Engine A (EMBER).
+
+This module handles loading a pre-trained LightGBM structural classifier
+and parsing uploaded Windows executable binaries (.exe, .dll) using pefile
+to approximate EMBER features. It also supports loading mock JSON profiles.
+"""
+
+import os
+import pickle
+import pefile
+import numpy as np
+import warnings
+import json
+
+
+class EngineAInfer:
+    """
+    Inference Engine for Static Structural Malware Detection.
+
+    Attributes:
+        model: The loaded LightGBM classifier.
+        feature_names (dict): Mapping of extracted array indices to human-readable PE header names.
+    """
+
+    def __init__(self, model_path="models/engine_a_model.pkl"):
+        """
+        Initializes the inference engine by loading the specified LightGBM model.
+
+        Args:
+            model_path (str): Path to the pickled LightGBM model file.
+        """
+        self.model = None
+        if os.path.exists(model_path):
+            with open(model_path, "rb") as f:
+                self.model = pickle.load(f)
+        else:
+            print(
+                f"Warning: {model_path} not found. Engine A will fail unless trained."
+            )
+
+        # Human readable names for the custom extracted pefile features
+        self.feature_names = {
+            0: "Machine Architecture",
+            1: "SizeOfOptionalHeader",
+            2: "Characteristics",
+            3: "MajorLinkerVersion",
+            4: "MinorLinkerVersion",
+            5: "SizeOfCode",
+            6: "SizeOfInitializedData",
+            7: "SizeOfUninitializedData",
+            8: "AddressOfEntryPoint",
+            9: "BaseOfCode",
+            10: "NumberOfSections",
+            11: "NumberOfImports",
+            12: "TotalImportedFunctions",
+            13: "NumberOfExports",
+        }
+
+    def extract_features(self, file_path):
+        """
+        Approximates EMBER features using the pefile library.
+
+        Parses basic PE headers, sections, imports, and exports, and pads
+        the vector to the 2381 features expected by the EMBER dataset format.
+
+        Args:
+            file_path (str): Absolute or relative path to the executable file.
+
+        Returns:
+            np.ndarray: A 2D numpy array of shape (1, 2381) representing the features.
+        """
+        features = np.zeros(2381, dtype=np.float32)
+
+        try:
+            pe = pefile.PE(file_path)
+            # Basic header features
+            features[0] = pe.FILE_HEADER.Machine
+            features[1] = pe.FILE_HEADER.SizeOfOptionalHeader
+            features[2] = pe.FILE_HEADER.Characteristics
+            features[3] = pe.OPTIONAL_HEADER.MajorLinkerVersion
+            features[4] = pe.OPTIONAL_HEADER.MinorLinkerVersion
+            features[5] = pe.OPTIONAL_HEADER.SizeOfCode
+            features[6] = pe.OPTIONAL_HEADER.SizeOfInitializedData
+            features[7] = pe.OPTIONAL_HEADER.SizeOfUninitializedData
+            features[8] = pe.OPTIONAL_HEADER.AddressOfEntryPoint
+            features[9] = pe.OPTIONAL_HEADER.BaseOfCode
+
+            # Number of sections
+            features[10] = pe.FILE_HEADER.NumberOfSections
+
+            # Imports
+            if hasattr(pe, "DIRECTORY_ENTRY_IMPORT"):
+                features[11] = len(pe.DIRECTORY_ENTRY_IMPORT)
+                num_imports = sum(
+                    [len(entry.imports) for entry in pe.DIRECTORY_ENTRY_IMPORT]
+                )
+                features[12] = num_imports
+
+            # Exports
+            if hasattr(pe, "DIRECTORY_ENTRY_EXPORT"):
+                features[13] = len(pe.DIRECTORY_ENTRY_EXPORT.symbols)
+
+        except Exception as e:
+            print(f"PE parsing error (might not be a PE file): {e}")
+
+        return features.reshape(1, -1)
+
+    def predict(self, file_path):
+        """
+        Predicts if a file is malicious based on its structural layout.
+
+        Also calculates SHAP (SHapley Additive exPlanations) values to provide
+        explainability regarding which structural features influenced the decision.
+
+        Args:
+            file_path (str): Path to the target file. Can be a Windows PE binary
+                             or a simulated mock malware profile (.json).
+
+        Returns:
+            dict: Contains 'is_malware' (bool), 'malware_prob' (float), and
+                  'top_features' (list of tuples detailing top contributing factors).
+        """
+        if self.model is None:
+            raise Exception("Engine A model is not loaded. Train the model first.")
+
+        # Handle mock JSON profile logic
+        if file_path.endswith(".json"):
+            try:
+                with open(file_path, "r") as f:
+                    data = json.load(f)
+                if data.get("is_mock_profile"):
+                    features = np.array(data["ember_features"]).reshape(1, -1)
+                else:
+                    features = self.extract_features(file_path)
+            except:
+                features = self.extract_features(file_path)
+        else:
+            features = self.extract_features(file_path)
+
+        # Suppress LGBM missing feature names warning
+        with warnings.catch_warnings():
+            warnings.simplefilter("ignore", UserWarning)
+            probs = self.model.predict_proba(features)[0]
+
+            # Use LightGBM pred_contrib=True to get SHAP values for explainability
+            contributions = self.model.predict(features, pred_contrib=True)[0]
+
+        malware_prob = float(probs[1]) if len(probs) > 1 else float(probs[0])
+
+        # Explainability Logic: Exclude the last element (base expected value)
+        feature_shap = contributions[:-1]
+
+        # If the file is malicious, we want the most positive SHAP values.
+        # If it's benign, we want the most negative SHAP values.
+        is_malicious = malware_prob > 0.5
+
+        if is_malicious:
+            # Sort descending for highest positive push
+            top_indices = np.argsort(feature_shap)[::-1][:3]
+        else:
+            # Sort ascending for lowest negative push
+            top_indices = np.argsort(feature_shap)[:3]
+
+        top_features = []
+        for idx in top_indices:
+            name = self._get_feature_name(idx)
+            contrib = feature_shap[idx]
+            top_features.append((name, float(contrib)))
+
+        return {
+            "is_malware": is_malicious,
+            "malware_prob": malware_prob,
+            "top_features": top_features,
+        }
+
+    def _get_feature_name(self, idx):
+        """
+        Maps an EMBER numerical feature index to a human-readable category.
+
+        EMBER extracts 2381 features using LIEF. Because the parquet file strips
+        the column names, we rely on the known ranges of the EMBER feature specification
+        to explain what the model is looking at.
+        """
+        if idx in self.feature_names and idx <= 13:
+            # Our custom exact mappings for the first 14 features
+            return f"{self.feature_names[idx]} (Index {idx})"
+
+        # Generalized EMBER feature mapping by vector location
+        if 0 <= idx <= 255:
+            return f"Raw Byte Histogram Analysis (Index {idx})"
+        elif 256 <= idx <= 511:
+            return f"Byte Entropy/Complexity Matrix (Index {idx})"
+        elif 512 <= idx <= 615:
+            return f"Embedded Strings Metadata [Paths/URLs/RegKeys] (Index {idx})"
+        elif 616 <= idx <= 626:
+            return f"General PE Structural Info (Index {idx})"
+        elif 627 <= idx <= 688:
+            return f"PE Header Anomaly (Index {idx})"
+        elif 689 <= idx <= 944:
+            return f"PE Sections Characteristics (Index {idx})"
+        elif 945 <= idx <= 2224:
+            return f"Imported Libraries/API Calls (Index {idx})"
+        elif 2225 <= idx <= 2381:
+            return f"Exported Functions/Data Directories (Index {idx})"
+        else:
+            return f"Deep Structural Feature (Index {idx})"

src/engine_b/inference.py

@@ -0,0 +1,156 @@
+"""
+Inference script for Engine B (Malimg).
+
+This module provides functionality to read raw binary files, convert them
+to 2D grayscale image tensors, and classify their malware family using
+a pre-trained Convolutional Neural Network (CNN).
+"""
+
+import os
+import torch
+from torchvision import transforms
+from PIL import Image
+import numpy as np
+import json
+from src.engine_b.train import MalwareCNN
+import math
+
+
+def bytes_to_image(file_path):
+    """
+    Reads a raw binary file and converts its bytes into a square grayscale image.
+
+    If the file is a mock JSON profile, it reads the byte array from the JSON.
+    Otherwise, it reads the raw bytes of the executable and reshapes them into
+    a square 2D matrix, padding the end with zeros if necessary.
+
+    Args:
+        file_path (str): Path to the target binary file.
+
+    Returns:
+        PIL.Image: A grayscale (mode 'L') Image object representing the binary.
+    """
+    # Check if this is our safe mock profile
+    if file_path.endswith(".json"):
+        try:
+            with open(file_path, "r") as f:
+                data = json.load(f)
+            if data.get("is_mock_profile"):
+                byte_array = np.array(data["malimg_bytes"], dtype=np.uint8)
+            else:
+                with open(file_path, "rb") as f:
+                    binary_data = f.read()
+                byte_array = np.frombuffer(binary_data, dtype=np.uint8)
+        except:
+            with open(file_path, "rb") as f:
+                binary_data = f.read()
+            byte_array = np.frombuffer(binary_data, dtype=np.uint8)
+    else:
+        # Read raw binary
+        with open(file_path, "rb") as f:
+            binary_data = f.read()
+        byte_array = np.frombuffer(binary_data, dtype=np.uint8)
+
+    # Calculate image dimensions (square)
+    length = len(byte_array)
+    if length == 0:
+        return Image.new("L", (128, 128), color=0)
+
+    width = int(math.ceil(math.sqrt(length)))
+    height = width
+
+    # Pad array to form a perfect square
+    padded_length = width * height
+    padded_array = np.pad(byte_array, (0, padded_length - length), mode="constant")
+
+    # Reshape and create PIL Image
+    image_2d = padded_array.reshape((height, width))
+    img = Image.fromarray(image_2d, mode="L")
+    return img
+
+
+class EngineBInfer:
+    """
+    Inference Engine for Visual Malware Family Classification.
+
+    Attributes:
+        device (torch.device): CPU or CUDA device for inference.
+        classes (list): Ordered list of malware family class names.
+        model (MalwareCNN): The loaded PyTorch CNN.
+        transform (transforms.Compose): Image preprocessing pipeline.
+    """
+
+    def __init__(
+        self,
+        model_path="models/engine_b_model.pth",
+        classes_path="models/engine_b_classes.json",
+    ):
+        """
+        Initializes the vision inference engine.
+
+        Args:
+            model_path (str): Path to the trained PyTorch state dictionary.
+            classes_path (str): Path to the JSON list mapping indices to class names.
+        """
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.classes = []
+        if os.path.exists(classes_path):
+            with open(classes_path, "r") as f:
+                self.classes = json.load(f)
+        else:
+            # Fallback dummy classes if not trained yet
+            self.classes = [f"Class_{i}" for i in range(24)]
+
+        num_classes = len(self.classes)
+        self.model = MalwareCNN(num_classes=num_classes)
+
+        if os.path.exists(model_path):
+            self.model.load_state_dict(
+                torch.load(model_path, map_location=self.device, weights_only=True)
+            )
+        else:
+            print(f"Warning: {model_path} not found. Using untrained weights.")
+
+        self.model.to(self.device)
+        self.model.eval()
+
+        self.transform = transforms.Compose(
+            [
+                transforms.Resize((128, 128)),
+                transforms.ToTensor(),
+            ]
+        )
+
+    def predict(self, file_path):
+        """
+        Converts the target file to an image and runs CNN inference.
+
+        Args:
+            file_path (str): Path to the target file.
+
+        Returns:
+            dict: Contains 'family' (str), 'confidence' (float),
+                  'all_probabilities' (dict), and 'image' (PIL.Image upscaled).
+        """
+        img = bytes_to_image(file_path)
+        tensor = self.transform(img).unsqueeze(0).to(self.device)
+
+        with torch.no_grad():
+            outputs = self.model(tensor)
+            probabilities = torch.nn.functional.softmax(outputs, dim=1)[0]
+
+            # Create a dictionary of all class probabilities
+            all_probs = {
+                self.classes[i]: probabilities[i].item()
+                for i in range(len(self.classes))
+            }
+
+            top_prob, top_class = torch.max(probabilities, 0)
+            class_name = self.classes[top_class.item()]
+
+            return {
+                "family": class_name,
+                "confidence": top_prob.item(),
+                "all_probabilities": all_probs,
+                "image": img.resize((512, 512), resample=Image.NEAREST),
+            }

src/ui/app.py

@@ -0,0 +1,330 @@
+"""
+Gradio UI Application for the Dual-Engine Malware Analysis Pipeline.
+
+This module provides the frontend interface for uploading executables or simulated profiles,
+routing them to the static structural engine (Engine A) and the visual byte-map engine (Engine B),
+and rendering their ensemble metrics securely in the browser.
+"""
+
+import gradio as gr
+from src.engine_a.inference import EngineAInfer
+from src.engine_b.inference import EngineBInfer
+import os
+import datetime
+import torch
+from fpdf import FPDF
+import tempfile
+
+print("Loading ML models...")
+engine_a = EngineAInfer()
+engine_b = EngineBInfer()
+print("Models loaded successfully.")
+
+
+def analyze_malware(file_path):
+    """
+    Main processing pipeline triggered by the Gradio 'Analyze File' button.
+
+    Args:
+        file_path (str): The path to the uploaded file.
+
+    Returns:
+        tuple: (output_a_text, image_b, prob_table, ensemble_verdict) for the UI components.
+    """
+    if not file_path:
+        return "No file provided.", None, {}, "N/A"
+
+    try:
+        # Engine A Analysis
+        result_a = engine_a.predict(file_path)
+        prob_a = result_a["malware_prob"] * 100
+        status_a = "Malicious" if result_a["is_malware"] else "Benign"
+
+        text_a = f"Status: {status_a}\nConfidence: {prob_a:.2f}%"
+
+        # Build explainability for gr.Label
+        explain_dict = {}
+        total_abs_shap = sum(abs(f[1]) for f in result_a["top_features"])
+        if total_abs_shap == 0:
+            total_abs_shap = 1
+
+        for feature, shap_val in result_a["top_features"]:
+            impact_dir = "Malicious" if shap_val > 0 else "Benign"
+            label_name = f"{feature} [{impact_dir}]"
+            # Normalize to 0-1 so gr.Label can render the horizontal bars properly
+            normalized_impact = float(abs(shap_val) / total_abs_shap)
+            explain_dict[label_name] = normalized_impact
+
+        # Engine B Analysis
+        result_b = engine_b.predict(file_path)
+        all_probs = result_b["all_probabilities"]
+        image_b = result_b["image"]
+
+        # Ensemble Logic
+        ensemble_score = (prob_a + result_b["confidence"] * 100) / 2
+        ensemble_verdict = f"Overall Risk Score: {ensemble_score:.1f}/100"
+
+        # Generate PDF Report
+        pdf_path = generate_pdf_report(
+            file_path=file_path,
+            status_a=status_a,
+            prob_a=prob_a,
+            top_features=result_a["top_features"],
+            family_b=result_b["family"],
+            prob_b=result_b["confidence"] * 100,
+            image_b=image_b,
+            ensemble_score=ensemble_score,
+        )
+
+        return (
+            text_a,
+            explain_dict,
+            image_b,
+            all_probs,
+            ensemble_verdict,
+            gr.update(value=pdf_path, visible=True),
+        )
+
+    except Exception as e:
+        return (
+            f"Error processing file in Engine A: {e}",
+            {},
+            None,
+            {},
+            "Error",
+            gr.update(visible=False),
+        )
+
+
+def generate_pdf_report(
+    file_path, status_a, prob_a, top_features, family_b, prob_b, image_b, ensemble_score
+):
+
+    class PDFReport(FPDF):
+        def header(self):
+            # Header Bar
+            self.set_fill_color(40, 40, 40)
+            self.rect(0, 0, 210, 30, "F")
+            self.set_y(10)
+            self.set_font("Helvetica", style="B", size=24)
+            self.set_text_color(255, 255, 255)
+            self.cell(0, 10, "MALWARE ANALYSIS REPORT", border=0, align="C")
+            self.ln(15)
+
+        def footer(self):
+            self.set_y(-15)
+            self.set_font("Helvetica", style="I", size=8)
+            self.set_text_color(128, 128, 128)
+            self.cell(0, 10, f"Page {self.page_no()}", align="C")
+
+    pdf = PDFReport()
+    pdf.add_page()
+    pdf.set_auto_page_break(auto=True, margin=15)
+    pdf.set_y(40)  # Start below header
+
+    # 1. Overview Section
+    pdf.set_fill_color(230, 230, 230)
+    pdf.set_text_color(0, 0, 0)
+    pdf.set_font("Helvetica", style="B", size=14)
+    pdf.cell(
+        0,
+        10,
+        " 1. Execution Overview",
+        border=0,
+        new_x="LMARGIN",
+        new_y="NEXT",
+        fill=True,
+    )
+    pdf.ln(2)
+
+    pdf.set_font("Helvetica", size=11)
+    pdf.cell(40, 8, "Target File:", border=0)
+    pdf.set_font("Helvetica", style="B", size=11)
+    pdf.cell(
+        0, 8, f"{os.path.basename(file_path)}", border=0, new_x="LMARGIN", new_y="NEXT"
+    )
+
+    pdf.set_font("Helvetica", size=11)
+    pdf.cell(40, 8, "Timestamp:", border=0)
+    pdf.set_font("Helvetica", style="B", size=11)
+    pdf.cell(
+        0,
+        8,
+        f"{datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}",
+        border=0,
+        new_x="LMARGIN",
+        new_y="NEXT",
+    )
+
+    pdf.set_font("Helvetica", size=11)
+    pdf.cell(40, 8, "Risk Score:", border=0)
+    pdf.set_font("Helvetica", style="B", size=12)
+    # Highlight risk score in text
+    pdf.cell(
+        0, 8, f"{ensemble_score:.1f} / 100", border=0, new_x="LMARGIN", new_y="NEXT"
+    )
+    pdf.ln(5)
+
+    # 2. Engine A
+    pdf.set_font("Helvetica", style="B", size=14)
+    pdf.cell(
+        0,
+        10,
+        " 2. Structural Engine (EMBER)",
+        border=0,
+        new_x="LMARGIN",
+        new_y="NEXT",
+        fill=True,
+    )
+    pdf.ln(2)
+
+    pdf.set_font("Helvetica", size=11)
+    pdf.cell(
+        0,
+        8,
+        f"Verdict: {status_a.upper()} (Confidence: {prob_a:.1f}%)",
+        new_x="LMARGIN",
+        new_y="NEXT",
+    )
+    pdf.ln(2)
+
+    pdf.set_font("Helvetica", style="B", size=10)
+    pdf.cell(140, 8, "Top Contributing Feature (SHAP)", border=1, fill=True)
+    pdf.cell(
+        50,
+        8,
+        "Impact Direction",
+        border=1,
+        new_x="LMARGIN",
+        new_y="NEXT",
+        fill=True,
+        align="C",
+    )
+
+    pdf.set_font("Helvetica", size=10)
+    for feature, shap_val in top_features:
+        impact_dir = "MALICIOUS" if shap_val > 0 else "BENIGN"
+        pdf.cell(140, 8, f" {feature}", border=1)
+        pdf.cell(
+            50,
+            8,
+            f"{shap_val:+.2f} ({impact_dir})",
+            border=1,
+            new_x="LMARGIN",
+            new_y="NEXT",
+            align="C",
+        )
+    pdf.ln(8)
+
+    # 3. Engine B
+    pdf.set_font("Helvetica", style="B", size=14)
+    pdf.cell(
+        0,
+        10,
+        " 3. Visual Engine (Malimg)",
+        border=0,
+        new_x="LMARGIN",
+        new_y="NEXT",
+        fill=True,
+    )
+    pdf.ln(2)
+
+    pdf.set_font("Helvetica", size=11)
+    pdf.cell(
+        0,
+        8,
+        f"Predicted Malware Family: {family_b.upper()} (Confidence: {prob_b:.1f}%)",
+        new_x="LMARGIN",
+        new_y="NEXT",
+    )
+    pdf.ln(2)
+
+    # Embed Image
+    temp_dir = tempfile.gettempdir()
+    img_path = os.path.join(temp_dir, "byte_map.png")
+    image_b.save(img_path)
+
+    pdf.set_font("Helvetica", style="B", size=10)
+    pdf.cell(0, 8, "Grayscale Byte Map Render:", new_x="LMARGIN", new_y="NEXT")
+
+    # Draw image with a border
+    x_pos = pdf.get_x()
+    y_pos = pdf.get_y()
+    pdf.rect(x_pos, y_pos, 70, 70)
+    pdf.image(img_path, x=x_pos, y=y_pos, w=70, h=70)
+
+    report_path = os.path.join(
+        temp_dir, f"Analysis_Report_{os.path.basename(file_path)}.pdf"
+    )
+    pdf.output(report_path)
+    return report_path
+
+
+def get_system_status():
+    """
+    Fetches real-time system metrics for the UI header.
+
+    Returns:
+        str: Formatted markdown string containing time, compute device, and engine status.
+    """
+    current_time = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+    device = "CUDA" if torch.cuda.is_available() else "CPU"
+    return f"**System Time:** {current_time} | **Compute Device:** {device} | **Engines:** Online & Ready"
+
+
+# Build the Gradio UI
+with gr.Blocks() as app:
+    # Use HTML to make the title massively larger than the buttons
+    gr.HTML(
+        "<h1 style='text-align: center; font-size: 3.5rem; margin-bottom: 0.2rem; font-weight: bold;'>Dual-Engine Malware Analysis</h1>"
+    )
+    gr.HTML(
+        "<p style='text-align: center; font-size: 1.2rem; color: gray;'>Upload a Windows Executable (`.exe`, `.dll`) or a Simulated Profile (`.json`)</p>"
+    )
+
+    status_bar = gr.Markdown(get_system_status())
+
+    # Compact inputs
+    with gr.Row():
+        file_input = gr.File(label="Upload File", type="filepath", scale=4)
+        with gr.Column(scale=1):
+            analyze_btn = gr.Button("Analyze File", variant="primary", size="sm")
+            download_btn = gr.File(label="Download Full Report (PDF)", visible=False)
+
+    gr.Markdown("---")
+
+    # Single column layout for maximum horizontal space
+    ensemble_output = gr.Textbox(
+        label="Final Verdict", value="Pending", interactive=False
+    )
+    output_a = gr.Textbox(
+        label="Engine A: Structural (EMBER)",
+        value="Awaiting input...",
+        interactive=False,
+        lines=2,
+    )
+    explain_plot = gr.Label(
+        label="Engine A: Explainability (Feature Impact)", num_top_classes=3
+    )
+    image_output = gr.Image(label="Engine B: Grayscale Byte Map", type="pil")
+    output_b = gr.Label(label="Engine B: Family Probabilities", num_top_classes=25)
+
+    analyze_btn.click(
+        analyze_malware,
+        inputs=[file_input],
+        outputs=[
+            output_a,
+            explain_plot,
+            image_output,
+            output_b,
+            ensemble_output,
+            download_btn,
+        ],
+    )
+
+    # Live update the status bar
+    timer = gr.Timer(1)
+    timer.tick(get_system_status, inputs=None, outputs=status_bar)
+
+    # Load immediately on page open
+    app.load(get_system_status, inputs=None, outputs=status_bar)