Update app.py

app.py

@@ -7,12 +7,10 @@ from PIL import Image
 
 from torch.utils.mobile_optimizer import optimize_for_mobile
 
-model = timm.create_model('vit_base_patch16_224', pretrained=True)
-model.head = torch.nn.Linear(in_features=model.head.in_features, out_features=5)
-
-#path = "opt_model.pt"
-
-#model = model.jit.load(path)
+model = timm.create_model('resnet50', pretrained=True)
+model.fc = torch.nn.Linear(in_features=model.fc.in_features, out_features=5)
+path = "epoch_4_Resnet50-0.5contrast.pth"
+model.load_state_dict(torch.load(path))
 
 model.eval()
 
@@ -21,9 +19,9 @@ def transform_image(img_sample):
         transforms.Resize((224, 224)),  # Resize to 224x224
         transforms.ToTensor(),  # Convert PIL image to tensor
         transforms.ColorJitter(contrast=0.5),  # Contrast
-        transforms.RandomAdjustSharpness(sharpness_factor=0.5),
-        transforms.RandomSolarize(threshold=0.75),
-        transforms.RandomAutocontrast(p=1),
+        #transforms.RandomAdjustSharpness(sharpness_factor=0.5),
+        #transforms.RandomSolarize(threshold=0.75),
+        #transforms.RandomAutocontrast(p=1),
     ])
     transformed_img = transform(img_sample)
     return transformed_img
@@ -40,7 +38,7 @@ def predict(Image):
 
     with torch.no_grad():
         grade = torch.softmax(model(img.float()), dim=1)[0]
-    category = ["Normal", "Mild", "Moderate", "Severe", "Proliferative"]
+    category = ["0 - Normal", "1 - Mild", "2 - Moderate", "3 - Severe", "4 - Proliferative"]
     output_dict = {}
     for cat, value in zip(category, grade):
         output_dict[cat] = value.item()