Update model.py

the code structure modified

model.py

@@ -48,53 +48,49 @@ class StyleLoss(nn.Module):
         self.loss = F.mse_loss(G, self.target)
         return input  
 
-
-#Normalization
-
-cnn_normalization_mean = torch.tensor([0.485, 0.456, 0.406])
-cnn_normalization_std = torch.tensor([0.229, 0.224, 0.225])
-#image transformation
-# def image_transform(image):
-#     if isinstance(image, str):
-#         # If image is a path to a file, open it using PIL
-#         image = Image.open(image).convert('RGB')
-#     else:
-#         # If image is a NumPy array, convert it to a PIL image
-#         image = Image.fromarray(image.astype('uint8'), 'RGB')
-#     # Apply the same transformations as before
-#     image = transform(image).unsqueeze(0)
-#     return image
+#Image Transform
+        
+transform = transforms.Compose([
+    transforms.Resize((128,128)),  # scale imported image
+    transforms.ToTensor()])  # transform it into a torch tensor
 
 def image_transform(image):
-    if image is None:
-        return None
-
     if isinstance(image, str):
         # If image is a path to a file, open it using PIL
-        with open(image, "rb") as f:
-            image = Image.open(f).convert('RGB')
+        image = Image.open(image).convert('RGB')
     else:
-        # If image is already a PIL image, just convert it to RGB mode
-        image = image.convert('RGB')
-    
+        # If image is a NumPy array, convert it to a PIL image
+        image = Image.fromarray(image.astype('uint8'), 'RGB')
     # Apply the same transformations as before
-    image =image_transform(image).unsqueeze(0)
+    image = transform(image).unsqueeze(0)
     return image
 
 
-#Defining a model
-# weights=weights=torchvision.models.VGG19_Weights.IMAGENET1K_V1
-# cnn = models.vgg19(weights=weights).features.eval()
-weights = models.vgg19(pretrained='imagenet')
-cnn = weights.features.eval()
+        
+#Defining the Model
+cnn = models.vgg19(pretrained=True).features.eval()
 
-    
-#getting the input optimizer
-def get_input_optimizer(input_img):
-    # this line to show that input is a parameter that requires a gradient
-    optimizer = optim.LBFGS([input_img])
-    return optimizer
+#Normalization
+
+cnn_normalization_mean = torch.tensor([0.485, 0.456, 0.406]).to(device)
+cnn_normalization_std = torch.tensor([0.229, 0.224, 0.225]).to(device)
+
+# create a module to normalize input image so we can easily put it in a
+# ``nn.Sequential``
+class Normalization(nn.Module):
+    def __init__(self, mean, std):
+        super(Normalization, self).__init__()
+        # .view the mean and std to make them [C x 1 x 1] so that they can
+        # directly work with image Tensor of shape [B x C x H x W].
+        # B is batch size. C is number of channels. H is height and W is width.
+        self.mean = torch.tensor(mean).view(-1, 1, 1)
+        self.std = torch.tensor(std).view(-1, 1, 1)
+
+    def forward(self, img):
+        # normalize ``img``
+        return (img - self.mean) / self.std
 
+    
 # desired depth layers to compute style/content losses :
 
 
@@ -159,3 +155,8 @@ def get_style_model_and_losses(cnn, normalization_mean, normalization_std,
     model = model[:(i + 1)]
 
     return model, style_losses, content_losses 
+                                   
+def get_input_optimizer(input_img):
+    # this line to show that input is a parameter that requires a gradient
+    optimizer = optim.LBFGS([input_img])
+    return optimizer