Update model.py

model.py

@@ -5,14 +5,12 @@ import torch.nn as nn
 import torchvision.models as models
 from PIL import Image
 from torchvision import transforms
-import time
 import torch.nn.functional as F
 import torch.optim as optim
 import matplotlib.pyplot as plt
 import torchvision.transforms as transforms
 import copy
 import torchvision.models as models
-import torchvision.transforms.functional as TF
 from PIL import Image
 import numpy as np
 
@@ -68,25 +66,19 @@ transform = transforms.Compose([
     transforms.ToTensor()])  # transform it into a torch tensor
 
 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)
+  
+    if image is not None:
+        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
 
 
-        
-
-
-#Normalization
-
-cnn_normalization_mean = torch.tensor([0.485, 0.456, 0.406])
-cnn_normalization_std = torch.tensor([0.229, 0.224, 0.225])
-
 # create a module to normalize input image so we can easily put it in a
 # ``nn.Sequential``
 class Normalization(nn.Module):