Spaces:
Running
Running
| import gradio as gr | |
| import torch | |
| import onnxruntime as ort | |
| import numpy as np | |
| from PIL import Image | |
| from torchvision import transforms | |
| import torch.nn.functional as F | |
| import matplotlib.pyplot as plt | |
| _metadata_columns = [ | |
| "age", "usePesticide_I", "usePesticide_False", "usePesticide_True", "gender_M", "gender_F", "gender_O", | |
| "familySkinCancerHistory_False", "familySkinCancerHistory_True", "familySkinCancerHistory_I", "familyCancerHistory_True", | |
| "familyCancerHistory_False", "familyCancerHistory_I", "fitzpatrickSkinType_2.0", "fitzpatrickSkinType_1.0", | |
| "fitzpatrickSkinType_4.0", "fitzpatrickSkinType_3.0", "fitzpatrickSkinType_5.0", "macroBodyRegion_CHEST", | |
| "macroBodyRegion_NOSE", "macroBodyRegion_LIP", "macroBodyRegion_BACK", "macroBodyRegion_FOREARM", "macroBodyRegion_ARM", | |
| "macroBodyRegion_LEG", "macroBodyRegion_FACE", "macroBodyRegion_HAND", "macroBodyRegion_SCALP", "macroBodyRegion_NECK", | |
| "macroBodyRegion_FOOT", "macroBodyRegion_EAR", "macroBodyRegion_THIGH", "macroBodyRegion_ABDOMEN", | |
| "hasItched_True", "hasItched_False", "hasItched_I", "hasGrown_I", "hasGrown_False", "hasGrown_True", "hasHurt_True", "hasHurt_False", | |
| "hasHurt_I", "hasChanged_I", "hasChanged_False", "hasChanged_True", "hasBled_False", "hasBled_True", "hasBled_I", "hasElevation_I", | |
| "hasElevation_False", "hasElevation_True" | |
| ] | |
| try: | |
| ort_session = ort.InferenceSession("./pad25_mobilenetv3_folder_1.onnx") | |
| print("ONNX model loaded successfully.") | |
| except Exception as e: | |
| print(f"Error loading ONNX model: {e}") | |
| ort_session = None | |
| LABELS = ['ACK', 'BCC', 'MEL', 'NEV', 'SCC', 'SEK'] | |
| def create_plot(probs_history, steps_labels): | |
| fig, ax = plt.subplots(figsize=(10, 6)) | |
| class_data = {label: [] for label in LABELS} | |
| for step_probs in probs_history: | |
| for label, prob in step_probs.items(): | |
| class_data[label].append(prob * 100) | |
| # Identify top 3 classes based on final probability | |
| final_probs = {label: values[-1] for label, values in class_data.items()} | |
| top_classes = sorted(final_probs, key=final_probs.get, reverse=True)[:3] | |
| annotations = {} | |
| # Plot every class | |
| for name, values in class_data.items(): | |
| x_vals = range(len(values)) | |
| # Style logic | |
| if name in top_classes: # Highlight top classes | |
| line, = ax.plot(x_vals, values, label=name, linewidth=2, marker='o') | |
| color = line.get_color() | |
| # Collect Text Annotations | |
| for x, y in zip(x_vals, values): | |
| if x not in annotations: | |
| annotations[x] = [] | |
| annotations[x].append((y, f"{y:.1f}", color)) | |
| else: | |
| # Other low prob classes (faded) | |
| ax.plot(x_vals, values, label=name, alpha=1, linewidth=1) | |
| # Process annotations to avoid overlap | |
| for x in sorted(annotations.keys()): | |
| points = sorted(annotations[x], key=lambda p: p[0]) | |
| min_dist = 5 | |
| last_text_y = -100 | |
| for i, (y, text, color) in enumerate(points): | |
| text_y = y + 3 | |
| if text_y < last_text_y + min_dist: | |
| text_y = last_text_y + min_dist | |
| ax.text(x, text_y, text, ha='center', fontweight='bold', fontsize=10, color='black') | |
| last_text_y = text_y | |
| ax.set_xticks(range(len(steps_labels))) | |
| ax.set_xticklabels(steps_labels, rotation=30, ha='right') | |
| ax.set_ylabel("Probability (%)") | |
| ax.set_xlabel("Incrementally Added Features") | |
| ax.set_ylim(0, 115) | |
| ax.grid(True, linestyle='--', alpha=0.3) | |
| ax.legend(loc='upper right', bbox_to_anchor=(1.10, 1), borderaxespad=0., framealpha=0.8) | |
| plt.tight_layout() | |
| return fig | |
| def predict(image, age, region, cancer_history, skin_cancer_history, bleed, hurt, itch, grown, changed, elevation): | |
| if ort_session is None: | |
| return "Model not loaded", None | |
| steps = [ | |
| ("Baseline (Image only)", {}), | |
| (f"Age ({age})", {"age": age}), | |
| (f"Region ({region})", {"macroBodyRegion": region}), | |
| ] | |
| symptoms_map = { | |
| "Cancer History": ("familyCancerHistory", cancer_history), | |
| "Skin Cancer History": ("familySkinCancerHistory", skin_cancer_history), | |
| "Bleed": ("hasBled", bleed), | |
| "Hurt": ("hasHurt", hurt), | |
| "Itch": ("hasItched", itch), | |
| "Grew": ("hasGrown", grown), | |
| "Changed": ("hasChanged", changed), | |
| "Elevation": ("hasElevation", elevation) | |
| } | |
| for label, (key, val) in symptoms_map.items(): | |
| steps.append((f"{label} ({val})", {key: val})) | |
| probs_history = [] | |
| steps_labels = [] | |
| if image is not None: | |
| transform = transforms.Compose([ | |
| transforms.Resize((224, 224)), | |
| transforms.ToTensor(), | |
| transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) | |
| ]) | |
| image_pil = Image.open(image).convert('RGB') | |
| image_tensor = transform(image_pil).unsqueeze(0) | |
| else: | |
| image_tensor = torch.zeros(1, 3, 224, 224) | |
| def set_feature(vector, feature_name, value): | |
| col_name = f"{feature_name}_{value}" | |
| if col_name in _metadata_columns: | |
| idx = _metadata_columns.index(col_name) | |
| vector[idx] = 1.0 | |
| accumulated_features = {} | |
| for step_name, new_features in steps: | |
| skip_feature = False | |
| for key, value in new_features.items(): | |
| # I had to add this ugly "None" option in the select ;/ | |
| if value == "None" or value is None or value == []: | |
| skip_feature = True | |
| if skip_feature: | |
| continue | |
| steps_labels.append(step_name) | |
| accumulated_features.update(new_features) | |
| metadata_vector = np.zeros(len(_metadata_columns), dtype=np.float32) | |
| if "age" in accumulated_features and accumulated_features["age"] is not None: | |
| if "age" in _metadata_columns: | |
| val = accumulated_features["age"] | |
| metadata_vector[_metadata_columns.index("age")] = float(val) if val is not None else np.nan | |
| else: | |
| if "age" in _metadata_columns: | |
| metadata_vector[_metadata_columns.index("age")] = np.nan | |
| if "macroBodyRegion" in accumulated_features and accumulated_features["macroBodyRegion"]: | |
| set_feature(metadata_vector, "macroBodyRegion", accumulated_features["macroBodyRegion"]) | |
| symptom_keys = ["familyCancerHistory", "familySkinCancerHistory", "hasBled", "hasHurt", "hasItched", "hasGrown", "hasChanged", "hasElevation"] | |
| for key in symptom_keys: | |
| if key in accumulated_features: | |
| val = accumulated_features[key] | |
| if val != "None": | |
| set_feature(metadata_vector, key, val) | |
| metadata_tensor = torch.tensor(metadata_vector, dtype=torch.float32).unsqueeze(0) | |
| ort_inputs = { | |
| ort_session.get_inputs()[0].name: image_tensor.numpy(), | |
| ort_session.get_inputs()[1].name: metadata_tensor.numpy() | |
| } | |
| ort_outs = ort_session.run(None, ort_inputs) | |
| log_probs = ort_outs[0][0] | |
| probs = F.softmax(torch.tensor(log_probs), dim=0).numpy() | |
| probs_dict = {LABELS[i]: float(probs[i]) for i in range(len(LABELS))} | |
| probs_history.append(probs_dict) | |
| final_result = probs_history[-1] | |
| plot = create_plot(probs_history, steps_labels) | |
| return final_result, plot | |
| def clear_func(): | |
| return None, None, None, "None", "None", "None", "None", "None", "None", "None", "None", None, None | |
| with gr.Blocks() as demo: | |
| with gr.Row(): | |
| gr.Markdown("# PRISM: A Clinically Interpretable Stepwise Framework for Multimodal Skin Cancer Diagnosis") | |
| with gr.Row(): | |
| with gr.Column(): | |
| image = gr.Image(type="filepath", height=534, label="Input Image",) | |
| with gr.Column(): | |
| age = gr.Number(label="Age", value=None) | |
| region = gr.Dropdown(multiselect=False, allow_custom_value=False, label="Region", choices=[None, 'ARM', 'NECK', 'FACE', 'HAND', 'FOREARM', 'CHEST', 'NOSE', 'LEG', | |
| 'THIGH', 'SCALP', 'EAR', 'BACK', 'FOOT', 'ABDOMEN', 'LIP']) | |
| with gr.Row(): | |
| with gr.Column(): | |
| cancer_history = gr.Radio(label="Cancer history", choices=["True", "False", "None"], value="None") | |
| skin_cancer_history = gr.Radio(label="Skin cancer history", choices=["True", "False", "None"], value="None") | |
| bleed = gr.Radio(label="Bled", choices=["True", "False", "None"], value="None") | |
| hurt = gr.Radio(label="Pain", choices=["True", "False", "None"], value="None") | |
| with gr.Column(): | |
| itch = gr.Radio(label="Itch", choices=["True", "False", "None"], value="None") | |
| grown = gr.Radio(label="Grew", choices=["True", "False", "None"], value="None") | |
| changed = gr.Radio(label="Changed", choices=["True", "False", "None"], value="None") | |
| elevation = gr.Radio(label="Elevation", choices=["True", "False", "None"], value="None") | |
| examples = [ | |
| ["assets/examples/98540_74812_0_SCC.png", 91.0, "NECK", "False", "False", "False", "False", "True", "None", "None", "None",], | |
| ["assets/examples/23312_80156_1_BCC.png", 78.0, "NOSE", "True", "False", "True", "True", "True", "False", "False", "True",], | |
| ["assets/examples/33586_53648_1_ACK.png", 43.0, "FOREARM", "True", "False", "True", "True", "True", "False", "False", "True",], | |
| ["assets/examples/61243_97612_0_SEK.png", 73.0, "ARM", "False", "False", "False", "False", "False", "False", "False", "True",], | |
| ["assets/examples/83727_22982_0_NEV.png", 38.0, "THIGH", "False", "True", "False", "False", "True", "True", "True", "False",], | |
| ["assets/examples/86611_83131_0_MEL.png", 69.0, "FOREARM", "False", "True", "False", "True", "False", "False", "False", "False",], | |
| ] | |
| with gr.Row(): | |
| with gr.Column(): | |
| output_plot = gr.Plot(label="Incremental Prediction Change") | |
| with gr.Column(): | |
| output = gr.Label(label="Output", num_top_classes=6) | |
| gr.Examples(examples=examples, | |
| inputs=[image, age, region, cancer_history, skin_cancer_history, | |
| bleed, hurt, itch, grown, changed, elevation]) | |
| with gr.Row(): | |
| with gr.Column(): | |
| submit = gr.Button("Submit") | |
| submit.click(predict, inputs=[image, age, region, cancer_history, skin_cancer_history, bleed, hurt, itch, grown, changed, elevation], outputs=[output, output_plot]) | |
| clear = gr.Button("Clear") | |
| clear.click(clear_func, inputs=[], outputs=[image, age, region, cancer_history, skin_cancer_history, bleed, hurt, itch, grown, changed, elevation, output, output_plot]) | |
| demo.launch() |