Update app.py

app.py

@@ -303,37 +303,37 @@ class Discriminator(nn.Module):
         return self.sigmoid(x)
     
 
-class vgg19(nn.Module):
-    def __init__(self, pre_trained=True, require_grad=False):
-        super(vgg19, self).__init__()
-
-        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        num_gpus = torch.cuda.device_count()
-        vgg_features = models.vgg19(pretrained=pre_trained).features
-        self.seq_list = nn.ModuleList([nn.Sequential(ele) for ele in vgg_features])
-
-        self.vgg_layer = ['conv1_1', 'relu1_1', 'conv1_2', 'relu1_2', 'pool1',
-                          'conv2_1', 'relu2_1', 'conv2_2', 'relu2_2', 'pool2',
-                          'conv3_1', 'relu3_1', 'conv3_2', 'relu3_2', 'conv3_3', 'relu3_3', 'conv3_4', 'relu3_4', 'pool3',
-                          'conv4_1', 'relu4_1', 'conv4_2', 'relu4_2', 'conv4_3', 'relu4_3', 'conv4_4', 'relu4_4', 'pool4',
-                          'conv5_1', 'relu5_1', 'conv5_2', 'relu5_2', 'conv5_3', 'relu5_3', 'conv5_4', 'relu5_4', 'pool5']
-
-        if not require_grad:
-            for parameter in self.parameters():
-                parameter.requires_grad = False
-        if num_gpus == 2:
-            self = DataParallel(self, device_ids=[0,1]).to(device)
-        else:
-            self.to(device)  # Move the entire model to the selected device
-
-    def forward(self, x):
-        vgg_outputs = []
-
-        for layer in self.seq_list:
-            x = layer(x)
-            vgg_outputs.append(x)
-
-        return vgg_outputs
+# class vgg19(nn.Module):
+#     def __init__(self, pre_trained=True, require_grad=False):
+#         super(vgg19, self).__init__()
+
+#         device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+#         num_gpus = torch.cuda.device_count()
+#         vgg_features = models.vgg19(pretrained=pre_trained).features
+#         self.seq_list = nn.ModuleList([nn.Sequential(ele) for ele in vgg_features])
+
+#         self.vgg_layer = ['conv1_1', 'relu1_1', 'conv1_2', 'relu1_2', 'pool1',
+#                           'conv2_1', 'relu2_1', 'conv2_2', 'relu2_2', 'pool2',
+#                           'conv3_1', 'relu3_1', 'conv3_2', 'relu3_2', 'conv3_3', 'relu3_3', 'conv3_4', 'relu3_4', 'pool3',
+#                           'conv4_1', 'relu4_1', 'conv4_2', 'relu4_2', 'conv4_3', 'relu4_3', 'conv4_4', 'relu4_4', 'pool4',
+#                           'conv5_1', 'relu5_1', 'conv5_2', 'relu5_2', 'conv5_3', 'relu5_3', 'conv5_4', 'relu5_4', 'pool5']
+
+#         if not require_grad:
+#             for parameter in self.parameters():
+#                 parameter.requires_grad = False
+#         if num_gpus == 2:
+#             self = DataParallel(self, device_ids=[0,1]).to(device)
+#         else:
+#             self.to(device)  # Move the entire model to the selected device
+
+#     def forward(self, x):
+#         vgg_outputs = []
+
+#         for layer in self.seq_list:
+#             x = layer(x)
+#             vgg_outputs.append(x)
+
+#         return vgg_outputs
     
 
 
@@ -400,123 +400,123 @@ class TVLoss(nn.Module):
     
 
 # Create the directory if it doesn't exist
-output_directory = '/kaggle/working/model/'
-os.makedirs(output_directory, exist_ok=True)
+#output_directory = '/kaggle/working/model/'
+#os.makedirs(output_directory, exist_ok=True)
 
 
 
-def train(config):
-#     print("config : ", config)
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    num_gpus = torch.cuda.device_count()
-    transform = transforms.Compose([crop(config['scale'], config['patch_size']), augmentation()])
-    dataset = mydata(GT_path=config['GT_path'], LR_path=config['LR_path'], in_memory=config['in_memory'], transform=transform)
-    loader = DataLoader(dataset, batch_size=config['batch_size'], shuffle=True, num_workers=config['num_workers'])
+# def train(config):
+# #     print("config : ", config)
+#     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+#     num_gpus = torch.cuda.device_count()
+#     transform = transforms.Compose([crop(config['scale'], config['patch_size']), augmentation()])
+#     dataset = mydata(GT_path=config['GT_path'], LR_path=config['LR_path'], in_memory=config['in_memory'], transform=transform)
+#     loader = DataLoader(dataset, batch_size=config['batch_size'], shuffle=True, num_workers=config['num_workers'])
     
-    generator = Generator()
+#     generator = Generator()
     
-    if config['fine_tuning']:
-        generator.load_state_dict(torch.load(config['generator_path']))
-    generator = nn.DataParallel(generator, device_ids=[0, 1]) if num_gpus == 2 else generator.to(device)
+#     if config['fine_tuning']:
+#         generator.load_state_dict(torch.load(config['generator_path']))
+#     generator = nn.DataParallel(generator, device_ids=[0, 1]) if num_gpus == 2 else generator.to(device)
         
-    generator.train()
+#     generator.train()
     
-    l2_loss = nn.MSELoss()
-    g_optim = optim.Adam(generator.parameters(), lr=1e-4)
+#     l2_loss = nn.MSELoss()
+#     g_optim = optim.Adam(generator.parameters(), lr=1e-4)
         
-    pre_epoch = 0
-    fine_epoch = 0
+#     pre_epoch = 0
+#     fine_epoch = 0
     
-    #### Train using L2_loss
-    while pre_epoch < config['pre_train_epoch']:
-        for i, tr_data in enumerate(loader):
-            gt = tr_data['GT'].to(device)
-            lr = tr_data['LR'].to(device)
-            output = generator(lr)
-            loss = l2_loss(gt, output)
-
-            g_optim.zero_grad()
-            loss.backward()
-            g_optim.step()
-
-        pre_epoch += 1
-
-        if pre_epoch % 2 == 0:
-            print(pre_epoch)
-            print(loss.item())
-            print('=========')
-        
-        if pre_epoch % 4 == 0:
-            torch.save(generator.state_dict(), '/kaggle/working/model/pre_trained_model_%03d.pt' % pre_epoch)
-
-    #### Train using perceptual & adversarial loss
-    vgg_net = vgg19().to(device)
-    vgg_net = vgg_net.eval()
-    discriminator = Discriminator()
-    if num_gpus == 2:
-        discriminator = DataParallel(discriminator, device_ids = [0,1]).to(device)
-    else:
-        discriminator = discriminator.to(device)
-    discriminator.train()
+#     #### Train using L2_loss
+#     while pre_epoch < config['pre_train_epoch']:
+#         for i, tr_data in enumerate(loader):
+#             gt = tr_data['GT'].to(device)
+#             lr = tr_data['LR'].to(device)
+#             output = generator(lr)
+#             loss = l2_loss(gt, output)
+
+#             g_optim.zero_grad()
+#             loss.backward()
+#             g_optim.step()
+
+#         pre_epoch += 1
+
+#         if pre_epoch % 2 == 0:
+#             print(pre_epoch)
+#             print(loss.item())
+#             print('=========')
+        
+#         if pre_epoch % 4 == 0:
+#             torch.save(generator.state_dict(), '/kaggle/working/model/pre_trained_model_%03d.pt' % pre_epoch)
+
+#     #### Train using perceptual & adversarial loss
+#     vgg_net = vgg19().to(device)
+#     vgg_net = vgg_net.eval()
+#     discriminator = Discriminator()
+#     if num_gpus == 2:
+#         discriminator = DataParallel(discriminator, device_ids = [0,1]).to(device)
+#     else:
+#         discriminator = discriminator.to(device)
+#     discriminator.train()
     
-    d_optim = optim.Adam(discriminator.parameters(), lr=1e-4)
-    scheduler = optim.lr_scheduler.StepLR(g_optim, step_size=2000, gamma=0.1)
+#     d_optim = optim.Adam(discriminator.parameters(), lr=1e-4)
+#     scheduler = optim.lr_scheduler.StepLR(g_optim, step_size=2000, gamma=0.1)
     
-    VGG_loss = perceptual_loss(vgg_net)
-    cross_ent = nn.BCELoss()
-    tv_loss = TVLoss()
-    real_label = torch.ones((gt.size(0), 2)).to(device)
-    fake_label = torch.zeros((gt.size(0), 2)).to(device)
-
-    while fine_epoch < config['fine_train_epoch']:
-        scheduler.step()
-        
-        for i, tr_data in enumerate(loader):
-            gt = tr_data['GT'].to(device)
-            lr = tr_data['LR'].to(device)
+#     VGG_loss = perceptual_loss(vgg_net)
+#     cross_ent = nn.BCELoss()
+#     tv_loss = TVLoss()
+#     real_label = torch.ones((gt.size(0), 2)).to(device)
+#     fake_label = torch.zeros((gt.size(0), 2)).to(device)
+
+#     while fine_epoch < config['fine_train_epoch']:
+#         scheduler.step()
+        
+#         for i, tr_data in enumerate(loader):
+#             gt = tr_data['GT'].to(device)
+#             lr = tr_data['LR'].to(device)
                         
-            ## Training Discriminator
-            output = generator(lr)
-            fake_prob = discriminator(output, lr)
-            real_prob = discriminator(gt, lr)
-            d_loss_real = cross_ent(real_prob, real_label)
-            d_loss_fake = cross_ent(fake_prob, fake_label)
+#             ## Training Discriminator
+#             output = generator(lr)
+#             fake_prob = discriminator(output, lr)
+#             real_prob = discriminator(gt, lr)
+#             d_loss_real = cross_ent(real_prob, real_label)
+#             d_loss_fake = cross_ent(fake_prob, fake_label)
             
-            d_loss = d_loss_real + d_loss_fake
+#             d_loss = d_loss_real + d_loss_fake
 
-            g_optim.zero_grad()
-            d_optim.zero_grad()
-            d_loss.backward()
-            d_optim.step()
+#             g_optim.zero_grad()
+#             d_optim.zero_grad()
+#             d_loss.backward()
+#             d_optim.step()
             
-            output = generator(lr)
-            fake_prob = discriminator(output, lr)
+#             output = generator(lr)
+#             fake_prob = discriminator(output, lr)
             
-            _percep_loss, hr_feat, sr_feat = VGG_loss((gt + 1.0) / 2.0, (output + 1.0) / 2.0, layer=config['feat_layer'])
+#             _percep_loss, hr_feat, sr_feat = VGG_loss((gt + 1.0) / 2.0, (output + 1.0) / 2.0, layer=config['feat_layer'])
 
-            L2_loss = l2_loss(output, gt)
-            percep_loss = config['vgg_rescale_coeff'] * _percep_loss
-            adversarial_loss = config['adv_coeff'] * cross_ent(fake_prob, real_label)
-            total_variance_loss = config['tv_loss_coeff'] * tv_loss(config['vgg_rescale_coeff'] * (hr_feat - sr_feat)**2)
+#             L2_loss = l2_loss(output, gt)
+#             percep_loss = config['vgg_rescale_coeff'] * _percep_loss
+#             adversarial_loss = config['adv_coeff'] * cross_ent(fake_prob, real_label)
+#             total_variance_loss = config['tv_loss_coeff'] * tv_loss(config['vgg_rescale_coeff'] * (hr_feat - sr_feat)**2)
             
-            g_loss = percep_loss + adversarial_loss + total_variance_loss + L2_loss
+#             g_loss = percep_loss + adversarial_loss + total_variance_loss + L2_loss
             
-            g_optim.zero_grad()
-            d_optim.zero_grad()
-            g_loss.backward()
-            g_optim.step()
+#             g_optim.zero_grad()
+#             d_optim.zero_grad()
+#             g_loss.backward()
+#             g_optim.step()
 
-        fine_epoch += 1
+#         fine_epoch += 1
 
-        if fine_epoch % 2 == 0:
-            print(fine_epoch)
-            print(g_loss.item())
-            print(d_loss.item())
-            print('=========')
+#         if fine_epoch % 2 == 0:
+#             print(fine_epoch)
+#             print(g_loss.item())
+#             print(d_loss.item())
+#             print('=========')
 
-        if fine_epoch % 4 == 0:
-            torch.save(generator.state_dict(), '/kaggle/working/model/MedSRGAN_gene_%03d.pt' % fine_epoch)
-            torch.save(discriminator.state_dict(), '/kaggle/working/model/MedSRGAN_discrim_%03d.pt' % fine_epoch)
+#         if fine_epoch % 4 == 0:
+#             torch.save(generator.state_dict(), '/kaggle/working/model/MedSRGAN_gene_%03d.pt' % fine_epoch)
+#             torch.save(discriminator.state_dict(), '/kaggle/working/model/MedSRGAN_discrim_%03d.pt' % fine_epoch)
 
 
 
@@ -545,14 +545,14 @@ config = {
     
 
 # List all folders in /kaggle/input/
-input_directory = '/kaggle/input/lumber-spine-dataset/'
+#input_directory = '/kaggle/input/lumber-spine-dataset/'
 # input_directory = '/kaggle/input/custom-dataset/custom_dataset/'
-all_folders = [f for f in os.listdir(input_directory) if os.path.isdir(os.path.join(input_directory, f))]
+#all_folders = [f for f in os.listdir(input_directory) if os.path.isdir(os.path.join(input_directory, f))]
 
 # Print the list of folders
-print("All folders in /kaggle/input/:")
-for folder in all_folders:
-    print(folder)
+# print("All folders in /kaggle/input/:")
+# for folder in all_folders:
+#     print(folder)
 
 #train(config)