init

app.py

@@ -40,33 +40,27 @@ X_demo = (Y_demo - B) / K
 
 df_oracle = make_dataframe(X_demo, Y_demo, 'Oracle')
 
-def prepare_data():
+def prepare_data(sel_num):
     interval = (Y_UB - Y_LB) / NUM_SEG
     all_x, all_y = [], []
+    prob = []
     for i in range(NUM_SEG):
         uniform_y_distribution = torch.distributions.Uniform(Y_UB - (i+1)*interval, Y_UB-i*interval)
-        y_uniform = uniform_y_distribution.sample((NUM_TRAIN_SAMPLES, 1))
+        y_uniform = uniform_y_distribution.sample((sel_num[i], 1))
 
         noise_distribution = torch.distributions.Normal(loc=0, scale=NOISE_SIGMA)
-        noise = noise_distribution.sample((NUM_TRAIN_SAMPLES, 1))
+        noise = noise_distribution.sample((sel_num[i], 1))
         y_uniform_oracle = y_uniform - noise
 
         x_uniform = (y_uniform_oracle - B) / K
-        all_x.append(x_uniform)
-        all_y.append(y_uniform)
-    return all_x, all_y
-
-def select_data(all_x, all_y, sel_num):
-    sel_x, sel_y = [], []
-    prob = []
-    for i in range(NUM_SEG):
-        sel_x += all_x[i][:sel_num[i]]
-        sel_y += all_y[i][:sel_num[i]]
+        all_x += x_uniform
+        all_y += y_uniform
         prob += [torch.tensor(sel_num[i]).float() for _ in range(sel_num[i])]
-    sel_x = torch.stack(sel_x)
-    sel_y = torch.stack(sel_y)
+
+    all_x = torch.stack(all_x)
+    all_y = torch.stack(all_y)
     prob = torch.stack(prob)
-    return sel_x, sel_y, prob
+    return all_x, all_y, prob
 
 
 def unzip_dataloader(training_loader):
@@ -79,7 +73,6 @@ def unzip_dataloader(training_loader):
     all_y = torch.cat(all_y)
     return all_x, all_y
 
-# Train the model
 def train(train_loader, training_bundle, num_epochs):
     training_df = make_dataframe(*unzip_dataloader(train_loader))
     for epoch in range(num_epochs):
@@ -116,8 +109,6 @@ def visualize(training_df, training_bundle, epoch):
 
 
 def make_video():
-    if osp.isfile('movie.mp4'):
-        os.remove('movie.mp4')
     (
         ffmpeg
             .input('train_log/*.png', pattern_type='glob', framerate=3)
@@ -143,7 +134,6 @@ class ReweightL2(_Loss):
         loss = loss.sum()
         return loss
 
-# we use a linear layer to regress the weight from height
 class LinearModel(nn.Module):
     def __init__(self, input_dim, output_dim):
         super(LinearModel, self).__init__()
@@ -206,15 +196,9 @@ class DummyDataset(Dataset):
     def __len__(self):
         return len(self.inputs)
 
-def run(num1, num2, num3, num4, num5, random_seed, submit):
-    sel_num = [num1, num2, num3, num4, num5]
-    sel_num = [int(num/100*NUM_PER_BUCKET) for num in sel_num]
-    torch.manual_seed(int(random_seed))
-    all_x, all_y = prepare_data()
-    sel_x, sel_y, prob = select_data(all_x, all_y, sel_num)
-    train_loader = DataLoader(DummyDataset(sel_x, sel_y, prob), BATCH_SIZE, shuffle=True)
 
-    training_df = make_dataframe(sel_x, sel_y)
+def vis_training_data(all_x, all_y):
+    training_df = make_dataframe(all_x, all_y)
     g = sns.jointplot(data=training_df, x='x', y='y', color='#003ea1', alpha=0.1, linewidths=0, s=100,
                       marginal_kws=dict(bins=torch.linspace(Y_LB, Y_UB, steps=NUM_SEG+1), rug=True),
                       xlim=((Y_LB - B) / K, (Y_UB - B) / K),
@@ -230,15 +214,28 @@ def run(num1, num2, num3, num4, num5, random_seed, submit):
     plt.savefig('training_data.png',bbox_inches='tight')
     plt.close()
 
+def clean_up_logs():
+    if not osp.exists('train_log'):
+        os.mkdir('train_log')
+    for f in os.listdir('train_log'):
+        os.remove(osp.join('train_log', f))
+    if osp.isfile('movie.mp4'):
+        os.remove('movie.mp4')
+
+def run(num1, num2, num3, num4, num5, random_seed, submit):
+    sel_num = [num1, num2, num3, num4, num5]
+    sel_num = [int(num/100*NUM_PER_BUCKET) for num in sel_num]
+    torch.manual_seed(int(random_seed))
+    all_x, all_y, prob = prepare_data(sel_num)
+    train_loader = DataLoader(DummyDataset(all_x, all_y, prob), BATCH_SIZE, shuffle=True)
+    vis_training_data(all_x, all_y)
+
     if submit == 0:
         text = "Press \"Start Regressing\" if your are happy with the training data. Regression takes ~10s."
     else:
         text = "Press \"Prepare Training Data\" before changing the training data. You may also change the random seed."
     if submit == 1:
-        if not osp.exists('train_log'):
-            os.mkdir('train_log')
-        for f in os.listdir('train_log'):
-            os.remove(osp.join('train_log', f))
+        clean_up_logs()
         regress(train_loader)
         make_video()
     output = 'train_log/{:05d}.png'.format(NUM_EPOCHS) if submit==1 else None
@@ -246,34 +243,36 @@ def run(num1, num2, num3, num4, num5, random_seed, submit):
     return 'training_data.png', text, output, video
 
 
-iface = gr.Interface(
-    fn=run,
-    inputs=[
-            gr.inputs.Slider(0, 100, default=50, step=0.1, label='Label percentage in [0, 2)'),
-            gr.inputs.Slider(0, 100, default=50, step=0.1, label='Label percentage in [2, 4)'),
-            gr.inputs.Slider(0, 100, default=50, step=0.1, label='Label percentage in [4, 6)'),
-            gr.inputs.Slider(0, 100, default=50, step=0.1, label='Label percentage in [6, 8)'),
-            gr.inputs.Slider(0, 100, default=50, step=0.1, label='Label percentage in [8, 10)'),
-            gr.inputs.Number(default=0, label='Random Seed', optional=False),
-            gr.inputs.Radio(['Prepare Training Data', 'Start Regressing!'],
-                            type="index", default=None, label='Mode', optional=False),
-            ],
-    outputs=[
-        gr.outputs.Image(type="file", label="Training data"),
-        gr.outputs.Textbox(type="auto", label='What\' s next?'),
-        gr.outputs.Image(type="file", label="Regression result"),
-        gr.outputs.Video(type='mp4', label='Training process')
-    ],
-    live=True,
-    allow_flagging='never',
-    title="Balanced MSE for Imbalanced Visual Regression [CVPR 2022]",
-    description="Welcome to the demo of Balanced MSE &#9878;. In this demo, we will work on a simple task: imbalanced <i>linear</i> regression. <br>"
-                "To get started, drag the sliders &#128071;&#128071; and create your label distribution! "
-                "<small>(Examples are at the bottom of the page.)</small>",
-    examples=[
-        [0.1, 0.8, 6.4, 51.2, 100, 0, 'Prepare Training Data'],
-        [1, 10, 100, 10, 1, 0, 'Prepare Training Data'],
-    ],
+if __name__ == '__main__':
     
-)
-iface.launch()
+    iface = gr.Interface(
+        fn=run,
+        inputs=[
+                gr.inputs.Slider(0, 100, default=50, step=0.1, label='Label percentage in [0, 2)'),
+                gr.inputs.Slider(0, 100, default=50, step=0.1, label='Label percentage in [2, 4)'),
+                gr.inputs.Slider(0, 100, default=50, step=0.1, label='Label percentage in [4, 6)'),
+                gr.inputs.Slider(0, 100, default=50, step=0.1, label='Label percentage in [6, 8)'),
+                gr.inputs.Slider(0, 100, default=50, step=0.1, label='Label percentage in [8, 10)'),
+                gr.inputs.Number(default=0, label='Random Seed', optional=False),
+                gr.inputs.Radio(['Prepare Training Data', 'Start Regressing!'],
+                                type="index", default=None, label='Mode', optional=False),
+                ],
+        outputs=[
+            gr.outputs.Image(type="file", label="Training data"),
+            gr.outputs.Textbox(type="auto", label='What\' s next?'),
+            gr.outputs.Image(type="file", label="Regression result"),
+            gr.outputs.Video(type='mp4', label='Training process')
+        ],
+        live=True,
+        allow_flagging='never',
+        title="Balanced MSE for Imbalanced Visual Regression [CVPR 2022]",
+        description="Welcome to the demo of Balanced MSE &#9878;. In this demo, we will work on a simple task: imbalanced <i>linear</i> regression. <br>"
+                    "To get started, drag the sliders &#128071;&#128071; and create your label distribution! "
+                    "<small>(Examples are at the bottom of the page.)</small>",
+        examples=[
+            [0.1, 0.8, 6.4, 51.2, 100, 0, 'Prepare Training Data'],
+            [1, 10, 100, 10, 1, 0, 'Prepare Training Data'],
+        ],
+        
+    )
+    iface.launch()