Update web_app.py

web_app.py

@@ -4,9 +4,6 @@ from fastai.vision.learner import create_body
 import streamlit as st
 from PIL import Image
 import cv2 as cv
-
-# ---------Backend--------------------------------------------------------------
-
 import os
 import glob
 import time
@@ -58,8 +55,6 @@ class ColorizationDataset(Dataset):
     def __len__(self):
         return len(self.paths)
 
-
-# A handy function to make our dataloaders
 def make_dataloaders(batch_size=16, n_workers=4, pin_memory=True, **kwargs):
     dataset = ColorizationDataset(**kwargs)
     dataloader = DataLoader(dataset, batch_size=batch_size, num_workers=n_workers,
@@ -313,10 +308,6 @@ def update_losses(model, loss_meter_dict, count):
 
 
 def lab_to_rgb(L, ab):
-    """
-    Takes a batch of images
-    """
-
     L = (L + 1.) * 50.
     ab = ab * 110.
     Lab = torch.cat([L, ab], dim=1).permute(0, 2, 3, 1).cpu().numpy()
@@ -352,10 +343,6 @@ def log_results(loss_meter_dict):
     for loss_name, loss_meter in loss_meter_dict.items():
         print(f"{loss_name}: {loss_meter.avg:.5f}")
 
-
-# pip install fastai==2.4
-
-
 def build_res_unet(n_input=1, n_output=2, size=256):
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     body = create_body(resnet18(), pretrained=True, n_in=n_input, cut=-2)
@@ -388,8 +375,6 @@ class MyDataset(torch.utils.data.Dataset):
         ab = img_lab[[1, 2], ...] / 110.
         return {'L': L, 'ab': ab}
 
-
-# A handy function to make our dataloaders
 def make_dataloaders2(batch_size=16, n_workers=4, pin_memory=True, **kwargs):
     dataset = MyDataset(**kwargs)
     dataloader = DataLoader(dataset, batch_size=batch_size, num_workers=n_workers,