Update app.py

app.py

@@ -41,7 +41,15 @@ CRC_LABELS = {
             7: "STR",
             8: "TUM",
         }
-
+BRACS_LABELS = {
+        0: "Normal",
+        1: "Pathological Benign",
+        2: "Usual Ductal Hyperplasia",
+        3: "Flat Epithelial Atypia",
+        4: "Atypical Ductal Hyperplasia",
+        5: "Ductal Carcinoma In Situ"
+        6: "Invasive Carcinoma"
+}
 print("Loading DinoV2 base model...")
 dinov2 = torch.hub.load('facebookresearch/dinov2', 'dinov2_vitg14_reg')
 
@@ -59,7 +67,7 @@ dinov2.eval()
 
 #torch.save(dinov2.state_dict(), "teacher_checkpoint_load.pt")
 
-def setup_linear(path):
+def setup_linear(path, classes = 4):
     print(f"Loading {path} linear classifier...")
     # Load the best checkpoint from the latest run
     linear_checkpoint = torch.load(path)
@@ -67,25 +75,12 @@ def setup_linear(path):
     linear_bias = linear_checkpoint["state_dict"]["head.bias"]
 
     # Create linear layer
-    linear = torch.nn.Linear(1536, 4)
+    linear = torch.nn.Linear(1536, classes)
     linear.weight.data = linear_weights
     linear.bias.data = linear_bias
     linear.eval()
     return linear
 
-def setup_linear_crc(path):
-    print(f"Loading {path} linear classifier...")
-    # Load the best checkpoint from the latest run
-    linear_checkpoint = torch.load(path)
-    linear_weights = linear_checkpoint["state_dict"]["head.weight"]
-    linear_bias = linear_checkpoint["state_dict"]["head.bias"]
-
-    # Create linear layer
-    linear = torch.nn.Linear(1536, 9)
-    linear.weight.data = linear_weights
-    linear.bias.data = linear_bias
-    linear.eval()
-    return linear
 
 # Move models to GPU if available
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -95,7 +90,7 @@ breakhis_path = "./breakhis_best.ckpt"
 breakhis_linear = setup_linear(breakhis_path).to(device)
 
 gleason_path = "./gleason_best.ckpt"
-gleason_linear = setup_linear(gleason_path).to(device)
+gleason_linear = setup_linear(gleason_path, 9).to(device)
 
 bach_path = "./bach_best.ckpt"
 bach_linear = setup_linear(bach_path).to(device)
@@ -103,6 +98,8 @@ bach_linear = setup_linear(bach_path).to(device)
 crc_path = "./crc_best.ckpt"
 crc_linear = setup_linear_crc(crc_path).to(device)
 
+#bracs_path = "./bracs_best.ckpt"
+#bracs_linear = setup_linear(bracs_path, 7).to(device)
 
 print(f"Models loaded on {device}")
 
@@ -142,6 +139,10 @@ def predict_bach(image):
 def predict_crc(image):
 
     return predict_class(image, crc_linear, "crc")
+    
+def predict_bracs(image):
+
+    return predict_class(image, bracs_linear, "bracs")
 
 
 def predict_class(image, linear, dataset):
@@ -181,7 +182,8 @@ def predict_class(image, linear, dataset):
             probs_dict[BACH_LABELS[idx]] = float(prob)
         elif dataset == "crc":
             probs_dict[CRC_LABELS[idx]] = float(prob)
-
+        elif dataset == "bracs":
+            probs_dict[BRACS_LABELS[idx]] = float(prob)
 
     return probs_dict
 
@@ -277,11 +279,30 @@ bach = gr.Interface(
         ],  # You can add example image paths here
     theme=gr.themes.Soft()
 )
+bracs = gr.Interface(
+    fn=predict_bracs,
+    inputs=gr.Image(type="filepath", label="Upload Cancer Image"),
+    outputs=gr.Label(num_top_classes=7, label="Bracs Tumor Type Prediction"),
+    title="Tumor Classification",
+    description="""
+    Upload a prostate cancer image to predict the tumor type. Your image must be at 40X magnification. Do not otherwise modify your image.
+
+    This model uses a custom-trained DinoV2 foundation model for pathology images with a linear classifier for tumor classification.
 
+    Images are classified as Normal, Pathological Benign, Usual Ductal Hyperplasia, Flat Epithelial Atypia, 
+        Atypical Ductal Hyperplasia, Ductal Carcinoma In Situ, Invasive Carcinoma
+
+    For this particular demo, images *must* be one of the sample classes - unsupported classes will yield confusing and/or useless results.
+    """,
+    examples=[
+        ],  # You can add example image paths here
+    theme=gr.themes.Soft()
+)
 
 
+        
 
-demo = gr.TabbedInterface([breakhis, gleason, crc, bach],["BreakHis", "Gleason", "CRC", "Bach"])
+demo = gr.TabbedInterface([breakhis, gleason, crc, bach, bracs],["BreakHis", "Gleason", "CRC", "Bach", "BRACS"])
 
 
 if __name__ == "__main__":