may be final

app.py

@@ -112,6 +112,9 @@ class crop(object):
         
         return {'LR' : LR_patch, 'GT' : GT_patch}
 
+
+
+
 class augmentation(object):
     
     def __call__(self, sample):
@@ -120,8 +123,6 @@ class augmentation(object):
         hor_flip = random.randrange(0,2)
         ver_flip = random.randrange(0,2)
         rot = random.randrange(0,2)
-    
-#         print(f"Horizontal Flip: {hor_flip}, Vertical Flip: {ver_flip}, Rotation: {rot}")
         if hor_flip:
             temp_LR = np.fliplr(LR_img)
             LR_img = temp_LR.copy()
@@ -253,6 +254,7 @@ class Discriminator(nn.Module):
         self.block_1_2 = D_Block(64, 128, stride=1)
         self.block_1_3 = D_Block(128, 128)
 
+
         # Layer for LR image size
         self.conv_2_1 = nn.Sequential(
             nn.Conv2d(in_channels, 64, (3, 3), stride=1, padding=1), nn.LeakyReLU()
@@ -294,48 +296,12 @@ class Discriminator(nn.Module):
         )
 
         x = self.flatten(x)
-
-#         print(x.shape)
-
         x = self.fc1(x)
         x = self.fc2(self.relu(x))
-        # outputs the range of 0 to 1
-#         print(x.shape)
         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 MeanShift(nn.Conv2d):
@@ -353,6 +319,7 @@ class MeanShift(nn.Conv2d):
             p.requires_grad = False
 
 
+
 class perceptual_loss(nn.Module):
     def __init__(self, vgg):
         super(perceptual_loss, self).__init__()
@@ -398,134 +365,12 @@ class TVLoss(nn.Module):
     @staticmethod
     def tensor_size(t):
         return t.size()[1] * t.size()[2] * t.size()[3]
-    
-
-# Create the directory if it doesn't exist
-#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'])
-    
-#     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)
-        
-#     generator.train()
-    
-#     l2_loss = nn.MSELoss()
-#     g_optim = optim.Adam(generator.parameters(), lr=1e-4)
-        
-#     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()
-    
-#     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)
-                        
-#             ## 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
-
-#             g_optim.zero_grad()
-#             d_optim.zero_grad()
-#             d_loss.backward()
-#             d_optim.step()
-            
-#             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'])
-
-#             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_optim.zero_grad()
-#             d_optim.zero_grad()
-#             g_loss.backward()
-#             g_optim.step()
-
-#         fine_epoch += 1
-
-#         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)
 
 
 
 config = {
     'LR_path': '/kaggle/input/lumber-spine-dataset/LRimages',
     'GT_path': '/kaggle/input/lumber-spine-dataset/HRimages',
-#     'LR_path': '/kaggle/input/custom-dataset/custom_dataset/train_LR',
-#     'GT_path': '/kaggle/input/custom-dataset/custom_dataset/train_HR',
     'res_num': 16,
     'num_workers': 0,
     'batch_size': 16,
@@ -543,101 +388,36 @@ config = {
     'generator_path': None,  # Set the path if available
     'mode': 'train'
 }
-    
 
-# List all folders in /kaggle/input/
-#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))]
-
-# Print the list of folders
-# print("All folders in /kaggle/input/:")
-# for folder in all_folders:
-#     print(folder)
-
-#train(config)
-
-# def preprocess_input(image_path):
-#     # Load the input image
-#     input_image = Image.open(image_path).convert("RGB")
-#     input_image = np.array(input_image) / 127.5 - 1.0
-#     input_image = input_image.transpose(2, 0, 1).astype(np.float32)
-#     return torch.tensor(input_image).unsqueeze(0)
-
-
-def preprocess_input(image):
-    # Convert Gradio image to PIL format
-    image_pil = to_pil_image(image)
-    image_pil=image_pil.convert("RGB")
-    # Convert image to numpy array and normalize
-    input_image = np.array(image_pil) / 127.5 - 1.0
-    # Transpose dimensions
-    input_image = input_image.transpose(2, 0, 1).astype(np.float32)
-    # Convert to PyTorch tensor and add batch dimension
-    input_tensor = torch.tensor(input_image).unsqueeze(0)
-    return input_tensor
 
+def preprocess_input(input_image):
+    input_image = np.array(input_image) / 127.5 - 1.0
+    input_image = input_image.transpose(2, 0, 1).astype(np.float32)
+    return torch.tensor(input_image).unsqueeze(0)
 
 
-    
-# def test_single_image(generator_path, input_image_path):
-#     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-#     num_gpus = torch.cuda.device_count()
-#     # Load the generator model
-#     generator = Generator()
-    
-#     if num_gpus == 2:
-#         state_dict = torch.load(generator_path)
-#         if 'module.' in list(state_dict.keys())[0]:
-#             state_dict = {k[7:]: v for k, v in state_dict.items()}
-#         # Load the state dictionary to the model
-#         generator.load_state_dict(state_dict)
-#         generator = DataParallel(generator, device_ids= [0,1]).to(device)
-        
-#     else:
-#         generator = generator.to(device)
-#         generator.load_state_dict(torch.load(generator_path))
-#     generator.eval()
-
-#     # Preprocess the input image
-#     input_image = preprocess_input(input_image_path).to(device)
-
-#     # Perform inference
-#     with torch.no_grad():
-#         output = generator(input_image)
-#         output = output[0].cpu().numpy()
-#         output = (output + 1.0) / 2.0
-#         output = output.transpose(1, 2, 0)
-#         plt.imshow(output)
-#         plt.show()
-
-def test_single_image(input_image):
-    return input_image
+def test_single_image( input_image):
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    num_gpus = torch.cuda.device_count()
-    generator_path = 'pre_trained_model_064.pt'
+    # generator_path = './pre_trained_model_064.pt'
+    generator_path = './MedSRGAN_gene_016.pt'
+
+
 
     # Load the generator model
-    generator = Generator()
-
-    if num_gpus == 2:
-        state_dict = torch.load(generator_path)
-        if 'module.' in list(state_dict.keys())[0]:
-            state_dict = {k[7:]: v for k, v in state_dict.items()}
-        # Load the state dictionary to the model
-        generator.load_state_dict(state_dict)
-        generator = torch.nn.DataParallel(generator, device_ids=[0, 1]).to(device)
-    else:
-        generator = generator.to(device)
-        generator.load_state_dict(torch.load(generator_path))
+    generator = Generator()          #for ours
+    state_dict = torch.load(generator_path, map_location='cpu')
+    if 'module.' in list(state_dict.keys())[0]:
+        state_dict = {k[7:]: v for k, v in state_dict.items()}
+    
+    generator.load_state_dict(state_dict)
     generator.eval()
 
     # Preprocess the input image
-    input_tensor = preprocess_input(input_image).to(device)
+    input_image = preprocess_input(input_image).to(device)
 
     # Perform inference
     with torch.no_grad():
-        output = generator(input_tensor)
+        output = generator(input_image)
         output = output[0].cpu().numpy()
         output = (output + 1.0) / 2.0
         output = output.transpose(1, 2, 0)
@@ -646,23 +426,13 @@ def test_single_image(input_image):
         return output_image
 
 
-
-
-
-# Load the generator model path
-#generator_path = 'pre_trained_model_064.pt'
-
-
-
-#test_single_image('pre_trained_model_064.pt',uploaded_image_data)
-# Define Gradio interface
 uploaded_image_data = gr.components.Image(type="pil", label="Upload Image")
-with gr.Column(scale=2, min_width=400):
-            output = gr.components.Image(type="pil", label="Super-Resolated Image")
+# with gr.Column(scale=2, min_width=400):
+output = gr.components.Image(type="pil", label="Super-Resolated Image")
 
 
 # Deploy the interface
-gr.Interface(test_single_image, inputs=uploaded_image_data, outputs=output, 
+gr.Interface(test_single_image, inputs=uploaded_image_data , outputs=output, 
               title="Image Super-Resolution", 
               description="Upload an image to see it super-resolated."
               ).launch()