Update app.py

app.py

@@ -130,6 +130,30 @@ class VBDDataset(Dataset):
         return self.image_ids.shape[0]
     
 
+def get_train_transform():
+    return A.Compose([
+        A.Flip(0.5),  
+        ToTensorV2(p=1.0)  
+    ], bbox_params={'format': 'pascal_voc', 'label_fields': ['labels']})
+
+def get_valid_transform():
+    return A.Compose([
+        ToTensorV2(p=1.0)  
+    ])
+
+
+def collate_fn(batch):
+    return tuple(zip(*batch))
+    
+    
+
+def format_prediction_string(labels, boxes, scores):
+    pred_strings = []
+    for j in zip(labels, scores, boxes):
+        pred_strings.append("{0} {1:.4f} {2} {3} {4} {5}".format(
+            j[0], j[1], j[2][0], j[2][1], j[2][2], j[2][3]))
+    return " ".join(pred_strings)
+    
 def generate_diagnostic_report(predictions, labels, threshold=0.5):
     # Initialize an empty report string
     report = "Diagnostic Report:\n\n"
@@ -251,9 +275,22 @@ uploaded_file = st.sidebar.file_uploader("Upload Chest X-ray image", type=["png"
 
 # Load the model (use your model loading function)
 # Ensure the model path is correct and accessible
-model = create_model(num_classes=14)
-model.load_state_dict(torch.load('Doctoria CXR Thoraric Full Model.pth', map_location=torch.device('cpu')))
-model.eval()
+
+#model = create_model(num_classes=15)
+#model.load_state_dict(torch.load('Doctoria CXR Thoraric Full Model.pth', map_location=torch.device('cpu')))
+#model.eval()
+
+def load_model(model_path):
+    # Create an instance of the VinDetector model
+    model = VinDetector(num_classes=14)  # Adjust num_classes as needed
+    
+    # Load the saved state_dict
+    model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')))
+    
+    model.eval()  # Set the model to evaluation mode
+    return model
+
+model = load_model('Doctoria CXR Thoraric Full Model.pth')
 
 def process_image(image_path):
     # Load and transform the image