Update app.py

app.py

@@ -63,8 +63,10 @@ hyperparameters = {
 results_data = {
     'SSD300': {'Transfer Pruning':{'map': [0.6, 0.65, 0.7, 0.75], 'flops': [1e9, 1.1e9, 1.2e9, 1.3e9], 'params': [1e6, 1.1e6, 1.2e6, 1.3e6]},
               'VIB Pruning':{'map': [0.9, 0.65, 0.7, 0.75], 'flops': [1e9, 1.1e9, 1.2e9, 1.3e9], 'params': [1e6, 1.1e6, 1.2e6, 1.3e6]}},
-    'SSD512': {'map': [0.7, 0.75, 0.8, 0.85], 'flops': [1.2e9, 1.3e9, 1.4e9, 1.5e9], 'params': [1.2e6, 1.3e6, 1.4e6, 1.5e6]},
-    'DETR': {'map': [0.8, 0.85, 0.9, 0.95], 'flops': [1.4e9, 1.5e9, 1.6e9, 1.7e9], 'params': [1.4e6, 1.5e6, 1.6e6, 1.7e6]},
+    'SSD512': {'Transfer Pruning':{'map': [0.6, 0.65, 0.7, 0.75], 'flops': [1e9, 1.1e9, 1.2e9, 1.3e9], 'params': [1e6, 1.1e6, 1.2e6, 1.3e6]},
+              'VIB Pruning':{'map': [0.9, 0.65, 0.7, 0.75], 'flops': [1e9, 1.1e9, 1.2e9, 1.3e9], 'params': [1e6, 1.1e6, 1.2e6, 1.3e6]}},
+    'DETR': {'SPARK':{'map': [0.6, 0.65, 0.7, 0.75], 'flops': [1e9, 1.1e9, 1.2e9, 1.3e9], 'params': [1e6, 1.1e6, 1.2e6, 1.3e6]},
+              'VOC':{'map': [0.9, 0.65, 0.7, 0.75], 'flops': [1e9, 1.1e9, 1.2e9, 1.3e9], 'params': [1e6, 1.1e6, 1.2e6, 1.3e6]}},
 }
 
 # Title of the research
@@ -97,12 +99,21 @@ with col1:
     # Display results table
     
     results = results_data[model]
-    df_results = pd.DataFrame({
-        'Model': [model]*4,
-        'mAP': results[pruning]['map'],
-        'FLOPs': results[pruning]['flops'],
-        'Params': results[pruning]['params']
-    })
+    if model in ['SSD300', 'SSD512']:
+        df_results = pd.DataFrame({
+            'Model': [model]*4,
+            'mAP': results[pruning]['map'],
+            'FLOPs': results[pruning]['flops'],
+            'Params': results[pruning]['params']
+        })
+    else:
+        df_results = pd.DataFrame({
+            'Model': [model]*4,
+            'mAP': results[dataset]['map'],
+            'FLOPs': results[dataset]['flops'],
+            'Params': results[dataset]['params']
+        })
+        
     
     st.markdown(df_results.style.hide(axis="index").to_html(), unsafe_allow_html=True)
 
@@ -110,18 +121,22 @@ with col1:
     st.subheader('Evolution Graphs')
     epochs = [1, 2, 3, 4]
 
-    fig = go.Figure()
-    fig.add_trace(go.Scatter(x=epochs, y=results['map'], mode='lines+markers', name='mAP'))
-    fig.update_layout(title='mAP per Epoch', xaxis_title='Epoch', yaxis_title='mAP')
-    st.plotly_chart(fig, use_container_width=True)
+    # fig = go.Figure()
+    # fig.add_trace(go.Scatter(x=epochs, y=results['map'], mode='lines+markers', name='mAP'))
+    # fig.update_layout(title='mAP per Epoch', xaxis_title='Epoch', yaxis_title='mAP')
+    # st.plotly_chart(fig, use_container_width=True)
 
     fig = go.Figure()
-    fig.add_trace(go.Scatter(x=epochs, y=results['flops'], mode='lines+markers', name='FLOPs', line=dict(color='orange')))
+    if model in ['SSD300', 'SSD512']:
+        ff=pruning
+    else:
+        ff=dataset
+    fig.add_trace(go.Scatter(x=epochs, y=results[ff]['flops'], mode='lines+markers', name='FLOPs', line=dict(color='orange')))
     fig.update_layout(title='FLOPs per Epoch', xaxis_title='Epoch', yaxis_title='FLOPs')
     st.plotly_chart(fig, use_container_width=True)
 
     fig = go.Figure() 
-    fig.add_trace(go.Scatter(x=epochs, y=results['params'], mode='lines+markers', name='Params', line=dict(color='green')))
+    fig.add_trace(go.Scatter(x=epochs, y=results[ff]['params'], mode='lines+markers', name='Params', line=dict(color='green')))
     fig.update_layout(title='Params per Epoch', xaxis_title='Epoch', yaxis_title='Params')
     st.plotly_chart(fig, use_container_width=True)