Update app.py

app.py

@@ -1,50 +1,60 @@
-import gradio as gr
-import torch
-import torch.nn as nn
-from torchvision import models, transforms
-from PIL import Image
-
-# Load the trained model
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
-# Define model
-model = models.resnet50(weights=None)
-in_features = model.fc.in_features
-model.fc = nn.Sequential(
-    nn.Linear(in_features, 512),
-    nn.ReLU(),
-    nn.Dropout(0.4),
-    nn.Linear(512, 47)  # 47 classes
-)
-model.load_state_dict(torch.load("cancer_model.pth", map_location=device))
-model.to(device)
-model.eval()
-
-# Label mapping (update this with actual class names if available)
-class_names = [f"Class {i}" for i in range(47)]
-
-# Transforms (same as validation)
-transform = transforms.Compose([
-    transforms.Lambda(lambda img: img.convert('RGB') if img.mode != 'RGB' else img),
-    transforms.Resize((224, 224)),
-    transforms.ToTensor(),
-    transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
-])
-
-# Prediction function
-def predict(img):
-    img = transform(img).unsqueeze(0).to(device)
-    with torch.no_grad():
-        outputs = model(img)
-        probs = torch.softmax(outputs, dim=1)
-        confidences, predicted = torch.max(probs, 1)
-    return {class_names[i]: float(probs[0][i]) for i in range(len(class_names))}
-
-# Gradio UI
-gr.Interface(
-    fn=predict,
-    inputs=gr.Image(type="pil"),
-    outputs=gr.Label(num_top_classes=5),
-    title="Multi-Cancer Classifier",
-    description="Upload a histopathology or cancer-related image. The model will predict its cancer type (47 classes)."
-).launch()
+import gradio as gr
+import torch
+import torch.nn as nn
+from torchvision import models, transforms
+from PIL import Image
+
+# Load the trained model
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+# Define model
+model = models.resnet50(weights=None)
+in_features = model.fc.in_features
+model.fc = nn.Sequential(
+    nn.Linear(in_features, 512),
+    nn.ReLU(),
+    nn.Dropout(0.4),
+    nn.Linear(512, 47)  # 47 classes
+)
+model.load_state_dict(torch.load("cancer_model.pth", map_location=device))
+model.to(device)
+model.eval()
+
+# Actual class names (replace generic labels)
+class_names = [
+    "lymph_fl", "colon_bnt", "cervix_koc", "colon_aca", "lymph_cll", "Skin_Melanoma", "Skin_Chickenpox",
+    "brain_glioma_tumor", "kidney_tumor", "all_pro", "Skin_Dyshidrotic Eczema", "Skin_Ringworm",
+    "cervix_dyk", "pancreatic_normal", "lung_colon_aca", "lung_lung_scc", "oral_scc", "kidney_stone",
+    "all_benign", "kidney_normal", "Skin_Basal Cell Carcinoma", "Skin_Nail Fungus", "all_early",
+    "pancreatic_tumor", "brain_pituitary_tumor", "lymph_mcl", "lung_colon_n", "Skin_Dermato Fibroma",
+    "brain_meningioma_tumor", "brain_normal", "breast_benign", "Skin_Actinic Keratosis",
+    "lung_lung_aca", "Skin_Vascular Lesion", "oral_normal", "Skin_Seborrheic Keratosis",
+    "Skin_Normal Skin", "all_pre", "Skin_Pigmented Benign Keratosis", "cervix_pab", "cervix_sfi",
+    "breast_malignant", "Skin_Acne", "cervix_mep", "Skin_Squamous Cell Carcinoma", "Skin_Nevus", "kidney_cyst"
+]
+
+# Transforms (must match training/validation)
+transform = transforms.Compose([
+    transforms.Lambda(lambda img: img.convert('RGB') if img.mode != 'RGB' else img),
+    transforms.Resize((224, 224)),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
+])
+
+# Prediction function
+def predict(img):
+    img = transform(img).unsqueeze(0).to(device)
+    with torch.no_grad():
+        outputs = model(img)
+        probs = torch.softmax(outputs, dim=1)
+        confidences, predicted = torch.max(probs, 1)
+    return {class_names[i]: float(probs[0][i]) for i in range(len(class_names))}
+
+# Gradio UI
+gr.Interface(
+    fn=predict,
+    inputs=gr.Image(type="pil"),
+    outputs=gr.Label(num_top_classes=5),
+    title="Multi-Cancer Classifier",
+    description="Upload a histopathology or cancer-related image. The model will predict its cancer type from 47 classes."
+).launch()