app.py

import torch
import torchvision
import gradio as gr
import pandas as pd
import json

from utils.load_image import load_image
from utils.load_model import load_model
from utils.get_layers import get_layers
from utils.get_image_path import get_image_path
from optimization import ActivationExtractor, optimize_activation_batch
from visualization import comparative_visualization

print("All modules Imported successfully")

sample_images = {
    "Pizza": "sample_images/Pizza.jpg",
    "Elephant": "sample_images/Elephant.jpg",
    "the-leaning-tower": "sample_images/the-leaning-tower.jpg"
}

def update_image(selected_sample):
  """
  Update the image input based on the selected sample image.

  Args:
    selected_sample (str): The name of the selected sample image.

  Returns:
    dict: A dictionary to update the image input component.
  """
  return gr.update(value=sample_images[selected_sample])


def optimize_neural_activation(
    image_path:str,
    use_real_image=True,
    architectures=["resnet50"],
    resnet50_layers=None,
    vgg16_layers=None,
    mobilenet_v2_layers=None,
    efficientnet_b0_layers=None,
    regularizations = {
      'reduction':'mean',
      'l2': 0.01,
      'tv': 0.02,
      'sparsity': 0.01,
      'entropy':0.01, 
      'l1': 0.01,
      'linf':0.01,
      'feature_map_sparsity':0.01,
      'clip': (0.0, 1.0)
      },
    steps:int=50,
    neuron_ids = [2,20,250],
    sample_image_select = None,
    log_freq: int=10
    ):

  """
  Main optimization loop for gradio interface

  Args:
    image_path (str): Path to user-uploaded or sample image.
    use_real_image (bool): Whether to use a real world image or a random noise.
    architectures (str): List of architectures to optimize for from. Can be selected from the list ["resnet50", "vgg16", "mobilenet_v2", "efficientnet_b0".
    selected_layers (list): List of specific layers to optimize.
    regularizations (dict): regularization weights. Use from the following: 'l2': 0.01,'tv': 0.02,'sparsity': 0.01,'entropy':0.01,'l1': 0.01,'linf':0.01,'feature_map_sparsity':0.01,'clip': (0.0, 1.0)
    steps (int): No. of steps of iteration
    log_freq (int): Frequency of logging optimization progress.

  Returns:
    fig: the matplotlib figure object with comparative visualizations.
    df: Pandas DataFrame summarizing the results.
  """

  selected_layers = {
        "resnet50": resnet50_layers or [],
        "vgg16": vgg16_layers or [],
        "mobilenet_v2": mobilenet_v2_layers or [],
        "efficientnet_b0": efficientnet_b0_layers or []
    }

  layer_dropdowns = {
        "resnet50": resnet50_layers or [],
        "vgg16": vgg16_layers or [],
        "mobilenet_v2": mobilenet_v2_layers or [],
        "efficientnet_b0": efficientnet_b0_layers or []
  }

  if layer_dropdowns:
    for arch, layers in layer_dropdowns.items():
      print(f"Processing {arch} with selected layers: {layers}")

  print(f"Dropdown values for resnet50: {resnet50_layers}")
  print(f"Dropdown values for vgg16: {vgg16_layers}")
  print(f"Dropdown values for mobilenet_v2: {mobilenet_v2_layers}")
  print(f"Dropdown values for efficientnet_b0: {efficientnet_b0_layers}")


  results_list = []
  summary = []

  steps = max(1, int(steps)) # ensuring that steps is a positive integer

  try:
    if isinstance(neuron_ids,str):
      neuron_ids=list(map(int, neuron_ids.split(",")))
  except ValueError:
    raise ValueError("Invalid ids provided. Please enter a comma-separated list of integers.")

  try:
    if isinstance(regularizations, str):
      regularizations=json.loads(regularizations)
  except json.JSONDecodeError:
    raise ValueError("Invalid JSON provided for regularizations")

  for arch in architectures:

    print(f"\n--- Optimizing for Architecture: {arch} ---" )
    model = load_model(architecture=arch)
    #print(f"Model {arch} loaded successfully: {model}") # for debugging
    layers = get_layers(model)

    layers_dict = dict(layers)
    # Validate user-provided layer names or select the first 3 layers by default
    # Filtering ensures that only valid layer names (those present in the model) are processed.
    default_layers = {
    "resnet50": ["layer1.0.conv1", "layer1.0.conv2", "layer1.0.downsample.1"],
    "vgg16": ["features.0", "features.5", "features.10"],
    "mobilenet_v2": ["features.0.0", "features.2.0", "features.5.0"],
    "efficientnet_b0": ["features.0.0", "features.2.0.block.0.0", "features.4.0.block.0.0"]}
    filtered_layers = []
    selected_layers = selected_layers or default_layers.get(arch, list(layers_dict.keys())[:3])
    for name in selected_layers[arch]:
      print(f"name:{name}, selectedLayers:{selected_layers}") #debugging
      if name in layers_dict:
        filtered_layers.append((name, layers_dict[name]))
      else:
        raise ValueError(f"Layer {name} not found in the selected architecture. Available layers are {layers_dict.keys()}")

    if not filtered_layers:
      print(f"No valid layers selected for architecture{arch}. Please select or use default values.")

    for layer_name, layer in filtered_layers:

      # Initialize a fresh input tensor for each layer by initializing the input inside this for loop
      # Problem: If the input tensor is reused across iterations, the optimized input from
      #          one layer's run affects subsequent runs. This results in:
      #          - Loss curves starting from a pre-optimized state rather than a random initialization.
      #          - Incorrect visualizations and inconsistencies across layer-specific optimizations.
      # Solution: By reinitializing the input tensor inside the loop, we ensure:
      #          - Each optimization starts from a fresh, randomly initialized input.
      #          - Independent and unbiased optimization for each layer.

      if use_real_image and image_path:
        image_path = get_image_path(image_path, sample_image_select)
        input_tensor = load_image(image_path)
        #print(f"Image loaded successfully with shape: {input_tensor.shape}") # for debugging

        print(f"Using real-world as starting input for {layer_name}")
      else:
        input_tensor = torch.randn(size=(1,3,224,224), requires_grad=True)
        print(f"Using random_noise as starting input for {layer_name}")


      print(f"Optimizing for {layer_name} | Checking activation size...")

      with ActivationExtractor(model=model, target_layer=layer) as extractor:
        _ = model(input_tensor)
        activation_shape = extractor.activation.shape
        num_channels = extractor.activation.size(1) # or extractor.activation.shape[1] or activation_shape[1]

      #num_channels = layer.out_channels if hasattr(layer, "out_channels") else 64
      target_neurons = [3, 20, 50]
      adjusted_neurons = [min(neuron, num_channels-1) for neuron in neuron_ids]

      print(f"Target neuron: {adjusted_neurons} out of {num_channels} neurons (activation_shape: {activation_shape})")


      result = optimize_activation_batch(model=model,
                                    target_neurons=adjusted_neurons,
                                    input_data=input_tensor,
                                    target_layer=layer,
                                    lr=0.1,
                                    steps=steps,
                                    regularizations_dict = regularizations,
                                    log_freq=log_freq)

      results_list.append({
          "architecture":arch,
          "layer_name":layer_name,
          "adjusted_neurons":adjusted_neurons,
          "optimized_input":result["optimized_input"],
          "loss_history": result["loss_history"]
      })

      summary.append({
          "Architecture": arch,
          "Layer Name": layer_name,
          "Input Shape": tuple(result["optimized_input"].shape),
          "Neuron Ids": adjusted_neurons,
          "Final Loss": result["loss_history"][-1]
      })

      #visualize_results(
          #input_tensor=results["optimized_input"],
          #loss_history=results["loss_history"],
          #neuron_id=adjusted_neurons,
          #layer_name=layer_name)

  # comparative plot
  fig = comparative_visualization(results_list)

  # summary table
  df = pd.DataFrame(summary)

  plot_file="output_plot.png"
  fig.savefig(plot_file)

  table_file="summary_table.csv"
  df.to_csv(table_file,index=False)

  return fig, df, plot_file, table_file


def run_interface():
  print("Inside run_interface()")
  def get_available_layers(architecture):
    """ Fetch layer names for selected architecture. """
    model = load_model(architecture=architecture)
    layers = get_layers(model)
    return [name for name,_ in layers]

  with gr.Blocks() as interface:
    gr.Markdown("# Neural Activation Optimizer")
    gr.Markdown("Optimize neural activations across different architectures and layers using real-world images or random noise")

    # Add Markdown instructions
    gr.Markdown("""
    ## Steps to Use the Tool:
    1. **Upload an Image or Select a Sample:**  
        - Upload a custom image or select one of the preloaded sample images (e.g., Pizza, Elephant, or The Leaning Tower).
    2. **Toggle Real-World Image Usage:**  
        - Choose whether to use the real-world image or initialize with random noise by selecting nothing.
    3. **Select Architectures and Layers:**  
        - Choose from supported architectures like ResNet50, VGG16, MobileNetV2, or EfficientNetB0.  
        - Specify the layers to optimize for each selected architecture.
    4. **Customize Neuron IDs:**  
        - Provide neuron IDs to optimize, separated by commas (e.g., `3,20,50`).
    5. **Adjust Regularizations and Parameters:**  
        - Edit the regularization settings using the JSON field provided.  
        - Specify the number of optimization steps and logging frequency.
    6. **Run the Optimization:**  
        - Click the "Run Optimization" button to start the process.  
        - Visualize the optimized input and loss curves in real-time.
    7. **Export Results:**  
        - Download the generated plots and summary table for further analysis.

    ## Supported Features:
    - Pre-trained architectures: ResNet, VGG, MobileNet, EfficientNet.
    - Dynamic customization of layers and regularizations.
    - Detailed visualizations and quick feedback.
    - Downloadable results for reproducibility.
    """)

    with gr.Row():
      image_input=gr.Image(type="filepath", label="Upload image or use sample")
      sample_image_select=gr.Dropdown(
          choices=list(sample_images.keys()),
          label="Select Sample Image",
          value="Sample Images in Dropdown"
      )
      real_image_toggle=gr.Checkbox(value=True, label="Use Real-World Image")

      neuron_input = gr.Textbox(
          value = "3,20,50",
          label="Neuron IDs (comma-separated, e.g., 3, 20, 50)",
          lines=1
      )

      log_freq_input = gr.Number(
        value=10,
        label = "Log Frequency (Steps)"
      )

    sample_image_select.change(
      fn=update_image,
      inputs=[sample_image_select],
      outputs=[image_input]
    )

    architecture_select = gr.CheckboxGroup(
        choices=["resnet50", "vgg16", "mobilenet_v2", "efficientnet_b0"],
        label="Select Architectures"

    )
    #layer_dropdowns = {
        #"resnet50": gr.Dropdown(choices=[], multiselect=True, label="Select layers for ResNet50"),
        #"vgg16": gr.Dropdown(choices=[], multiselect=True, label="Select layers for VGG16"),
        #"mobilenet_v2": gr.Dropdown(choices=[], multiselect=True, label="Select layers for MobileNetV2"),
        #"efficientnet_b0": gr.Dropdown(choices=[], multiselect=True, label="Select layers for EfficientNetB0")}

    resnet50_dropdown = gr.Dropdown(choices=[], multiselect=True, label="Select layers for ResNet50", visible=False)
    vgg16_dropdown = gr.Dropdown(choices=[], multiselect=True, label="Select layers for VGG16", visible=False)
    mobilenet_v2_dropdown = gr.Dropdown(choices=[], multiselect=True, label="Select layers for MobileNetV2", visible=False)
    efficientnet_b0_dropdown = gr.Dropdown(choices=[], multiselect=True, label="Select layers for EfficientNetB0", visible=False)
 

    regularizations_input=gr.Textbox(
        label="Regularizations (Editable)",
        value=json.dumps({'reduction':'mean','l2': 0.01,'tv': 0.02,'sparsity': 0.01,'entropy':0.01,'l1': 0.01,'linf':0.01,'feature_map_sparsity':0.01,'clip': [0.0, 1.0]}, indent=4),
        lines=10
    )

    step_input=gr.Number(value=50, label="Number of Steps (Default: 50)")

    output_plot=gr.Plot(label="Comparative Visualizations")
    output_dataframe=gr.DataFrame(label="Summary of Results")
    download_plot=gr.File(label='Download Plot')
    download_table=gr.File(label="Download Summary Table")
      
    def update_layer_dropdowns(selected_architectures):
      dropdowns = {
      "resnet50": get_available_layers("resnet50") if "resnet50" in selected_architectures else [],
      "vgg16": get_available_layers("vgg16") if "vgg16" in selected_architectures else [],
      "mobilenet_v2": get_available_layers("mobilenet_v2") if "mobilenet_v2" in selected_architectures else [],
      "efficientnet_b0": get_available_layers("efficientnet_b0") if "efficientnet_b0" in selected_architectures else []
      }
      return (
        gr.update(choices=dropdowns["resnet50"], visible="resnet50" in selected_architectures),
        gr.update(choices=dropdowns["vgg16"], visible="vgg16" in selected_architectures),
        gr.update(choices=dropdowns["mobilenet_v2"], visible="mobilenet_v2" in selected_architectures),
        gr.update(choices=dropdowns["efficientnet_b0"], visible="efficientnet_b0" in selected_architectures)
        )

     


    architecture_select.change(
        fn=update_layer_dropdowns, inputs=[architecture_select], 
        outputs=[
          resnet50_dropdown,
          vgg16_dropdown,
          mobilenet_v2_dropdown,
          efficientnet_b0_dropdown
          ]
    )

    run_button=gr.Button("Run Optimization")
    run_button.click(
        fn=optimize_neural_activation,
        inputs=[image_input,
        real_image_toggle,
        architecture_select,
        resnet50_dropdown,
        vgg16_dropdown,
        mobilenet_v2_dropdown,
        efficientnet_b0_dropdown,
        regularizations_input,
        step_input,
        neuron_input,
        sample_image_select,
        log_freq_input],
        outputs=[output_plot, output_dataframe, download_plot, download_table]
    )

  interface.launch(share=True)

if __name__=="__main__":
  print("Starting the Gradio interface...")
  run_interface()