update app

app.py

@@ -39,34 +39,156 @@ def optimize_model(input_model, sparsity, context, criteria):
     
     return comp_path
 
-def main_interface(model_file, sparsity, action):
-    if action == 'Speed':
-        return optimize_model(model_file, sparsity, 'local', "large_final")
+import matplotlib.pyplot as plt
+import seaborn as sns
+import io
+import numpy as np
 
-    if action == 'Size':
-        return optimize_model(model_file, sparsity, 'global', "large_final")
+def get_model_size(model_path):
+    """Get model size in MB"""
+    size_bytes = os.path.getsize(model_path)
+    size_mb = size_bytes / (1024 * 1024)
+    return round(size_mb, 2)
 
-    if action == 'Consumption':
-        return optimize_model(model_file, sparsity, 'local', "random")
-    else:
-        return "Action not supported"
+def create_size_comparison_plot(original_size, compressed_size):
+    """Create a bar plot comparing model sizes"""
+    # Set seaborn style
+    sns.set_style("darkgrid")
+    
+    # Create figure with higher DPI for better resolution
+    fig = plt.figure(figsize=(10, 6), dpi=150)
+    
+    # Set transparent background
+    fig.patch.set_alpha(0.0)
+    ax = plt.gca()
+    ax.patch.set_alpha(0.0)
+    
+    # Plot bars with custom colors and alpha
+    bars = plt.bar(['Original', 'Compressed'], 
+                  [original_size, compressed_size],
+                  color=['#FF6B00', '#FF9F1C'],
+                  alpha=0.8,
+                  width=0.6)
+    
+    # Add size labels on top of bars with improved styling
+    for bar in bars:
+        height = bar.get_height()
+        plt.text(bar.get_x() + bar.get_width()/2., height + (height * 0.01),
+                f'{height:.2f} MB',
+                ha='center', va='bottom',
+                fontsize=11,
+                fontweight='bold',
+                color='white')
+    
+    # Calculate compression percentage
+    compression_ratio = ((original_size - compressed_size) / original_size) * 100
+    
+    # Customize title and labels with better visibility
+    plt.title(f'Model Size Comparison\nCompression: {compression_ratio:.1f}%', 
+              fontsize=14, 
+              fontweight='bold', 
+              pad=20,
+              color='white')
+    
+    plt.xlabel('Model Version', 
+              fontsize=12, 
+              fontweight='bold',
+              labelpad=10,
+              color='white')
+    
+    plt.ylabel('Size (MB)', 
+              fontsize=12, 
+              fontweight='bold',
+              labelpad=10,
+              color='white')
+    
+    # Customize grid
+    ax.grid(alpha=0.2, color='gray')
+    
+    # Remove top and right spines
+    sns.despine()
+    
+    # Set y-axis limits with some padding
+    max_value = max(original_size, compressed_size)
+    plt.ylim(0, max_value * 1.2)
+    
+    # Add more y-axis ticks
+    plt.yticks(np.linspace(0, max_value * 1.2, 10))
+    
+    # Make tick labels white
+    ax.tick_params(colors='white')
+    for spine in ax.spines.values():
+        spine.set_color('white')
+    
+    # Format axes with white text
+    ax.xaxis.label.set_color('white')
+    ax.yaxis.label.set_color('white')
+    ax.tick_params(axis='x', colors='white')
+    ax.tick_params(axis='y', colors='white')
+    
+    # Format y-axis tick labels
+    ax.yaxis.set_major_formatter(plt.FuncFormatter(lambda x, p: f'{x:.1f}'))
+    
+    # Adjust layout to prevent label cutoff
+    plt.tight_layout()
+    
+    return fig
 
+def main_interface(model_name, sparsity, action):
+    import torchvision.models as models
+    
+    model_mapping = {
+        'ResNet18': models.resnet18(pretrained=True),
+        'ResNet50': models.resnet50(pretrained=True),
+        'MobileNetV2': models.mobilenet_v2(pretrained=True),
+        'EfficientNet-B0': models.efficientnet_b0(pretrained=True),
+        'VGG16': models.vgg16(pretrained=True),
+        'DenseNet121': models.densenet121(pretrained=True)
+    }
+    
+    model = model_mapping[model_name]
+    
+    # Save model temporarily
+    temp_path = "./temp_model.pth"
+    torch.save(model, temp_path)
+    
+    original_size = get_model_size(temp_path)
+    
+    try:
+        if action == 'Speed':
+            compressed_path = optimize_model(temp_path, sparsity, 'local', "large_final")
+        elif action == 'Size':
+            compressed_path = optimize_model(temp_path, sparsity, 'global', "large_final")
+        elif action == 'Consumption':
+            compressed_path = optimize_model(temp_path, sparsity, 'local', "random")
+        else:
+            return None, None
+        
+        compressed_size = get_model_size(compressed_path)
+        size_plot = create_size_comparison_plot(original_size, compressed_size)
+        
+        return compressed_path, size_plot
+    finally:
+        # Clean up temporary file
+        if os.path.exists(temp_path):
+            os.remove(temp_path)
 
-granularity = ['weight', 'filter']
-context = ['local', 'global']
-criteria = ['large_final', 'random']
 
+available_models = ['ResNet18', 'ResNet50', 'MobileNetV2', 'EfficientNet-B0', 'VGG16', 'DenseNet121']
 
 iface = gr.Interface(
     fn=main_interface,
-    inputs= [            
-             gr.File(label="Upload your PyTorch model (.pth file)"),
-             gr.Slider(label="Compression Level", minimum=0, maximum=100),
-             gr.Radio(["Speed", "Size", "Consumption"], label="Select Action")  
-            ],
-    outputs=gr.File(label="Download Compressed Model"),
-    title="FasterAI",
-    description="Upload your neural network model (.pt file) and receive a compressed version.",
+    inputs=[            
+        gr.Dropdown(choices=available_models, label="Select Model", value='ResNet18'),
+        gr.Slider(label="Compression Level", minimum=0, maximum=100, value=50),
+        gr.Radio(["Speed", "Size", "Consumption"], label="Select Action", value="Speed")  
+    ],
+    outputs=[
+        gr.File(label="Download Compressed Model"),
+        gr.Plot(label="Size Comparison")  # Changed from gr.Image to gr.Plot
+    ],
+    title="FasterAI Compressor",
+    description="Select a pre-trained PyTorch model to compress using our optimization techniques.",
 )
 
 iface.launch()