Spaces:
Sleeping
Sleeping
| import os | |
| import streamlit as st | |
| import cv2 | |
| import sys | |
| import argparse | |
| import numpy as np | |
| import json | |
| import torch | |
| import torch.nn.functional as F | |
| import detectron2.data.transforms as T | |
| import torchvision | |
| from collections import OrderedDict | |
| from scipy import spatial | |
| import matplotlib.pyplot as plt | |
| from packaging import version | |
| from detectron2.engine import DefaultPredictor | |
| from detectron2.utils.visualizer import Visualizer | |
| from detectron2.config import get_cfg | |
| from detectron2 import model_zoo | |
| from detectron2.data import Metadata | |
| from detectron2.structures.boxes import Boxes | |
| from detectron2.structures import Instances | |
| from plots.plot_pca_point import plot_pca_point | |
| from plots.plot_histogram_dist import plot_histogram_dist | |
| from plots.plot_gradcam import plot_gradcam | |
| def extract_features(model, img, box): | |
| height, width = img.shape[1:3] | |
| inputs = [{"image": img, "height": height, "width": width}] | |
| with torch.no_grad(): | |
| img = model.preprocess_image(inputs) | |
| features = model.backbone(img.tensor) | |
| features_ = [features[f] for f in model.roi_heads.box_in_features] | |
| box_features = model.roi_heads.box_pooler(features_, [box]) | |
| output_features = F.avg_pool2d(box_features, [7, 7]) | |
| output_features = output_features.view(-1, 256) | |
| return output_features | |
| def forward_model_full(model, cfg, cv_img): | |
| height, width = cv_img.shape[:2] | |
| transform_gen = T.ResizeShortestEdge( | |
| [cfg.INPUT.MIN_SIZE_TEST, cfg.INPUT.MIN_SIZE_TEST], cfg.INPUT.MAX_SIZE_TEST | |
| ) | |
| image = transform_gen.get_transform(cv_img).apply_image(cv_img) | |
| image = torch.as_tensor(image.astype("float32").transpose(2, 0, 1)) | |
| inputs = [{"image": image, "height": height, "width": width}] | |
| with torch.no_grad(): | |
| images = model.preprocess_image(inputs) | |
| features = model.backbone(images.tensor) | |
| proposals, _ = model.proposal_generator(images, features, None) | |
| features_ = [features[f] for f in model.roi_heads.box_in_features] | |
| box_features = model.roi_heads.box_pooler(features_, [x.proposal_boxes for x in proposals]) | |
| box_head = model.roi_heads.box_head(box_features) | |
| predictions = model.roi_heads.box_predictor(box_head) | |
| output_features = F.avg_pool2d(box_features, [7, 7]) | |
| output_features = output_features.view(-1, 256) | |
| probs = model.roi_heads.box_predictor.predict_probs(predictions, proposals) | |
| pred_instances, pred_inds = model.roi_heads.box_predictor.inference(predictions, proposals) | |
| pred_instances = model.roi_heads.forward_with_given_boxes(features, pred_instances) | |
| pred_instances = model._postprocess(pred_instances, inputs, images.image_sizes) | |
| instances = pred_instances[0]["instances"] | |
| instances.set("probs", probs[0][pred_inds]) | |
| instances.set("features", output_features[pred_inds]) | |
| return instances, cv_img | |
| def load_model(): | |
| cfg = get_cfg() | |
| cfg.merge_from_file(model_zoo.get_config_file("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml")) | |
| cfg.MODEL.ROI_HEADS.NUM_CLASSES = 3 | |
| cfg.MODEL.WEIGHTS = MODEL | |
| cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = TH | |
| cfg.MODEL.DEVICE = "cpu" | |
| metadata = Metadata() | |
| metadata.set( | |
| evaluator_type="coco", | |
| thing_classes=["neoplastic", "aphthous", "traumatic"], | |
| thing_dataset_id_to_contiguous_id={"1": 0, "2": 1, "3": 2} | |
| ) | |
| predictor = DefaultPredictor(cfg) | |
| model = predictor.model | |
| return dict( | |
| predictor=predictor, | |
| model=model, | |
| metadata=metadata, | |
| cfg=cfg | |
| ) | |
| def draw_box(file_name, box, type, model, resize_input=False): | |
| height, width, channels = img.shape | |
| pred_v = Visualizer(img[:, :, ::-1], model["metadata"], scale=1) | |
| instances = Instances((height, width), pred_boxes=Boxes(torch.tensor(box).unsqueeze(0)), pred_classes=torch.tensor([type])) | |
| pred_v = pred_v.draw_instance_predictions(instances) | |
| pred = pred_v.get_image()[:, :, ::-1] | |
| pred = cv2.resize(pred, (800, 800)) | |
| return pred | |
| def explain(img, model): | |
| state.write("Loading features...") | |
| database = json.load(open(FEATURES_DATABASE)) | |
| state.write("Computing logits...") | |
| instances, input = forward_model_full(model["model"], model["cfg"], img) | |
| instances.remove("pred_masks") | |
| pred_v = Visualizer(cv2.cvtColor(input, cv2.COLOR_BGR2RGB), model["metadata"], scale=1) | |
| pred_v = pred_v.draw_instance_predictions(instances.to("cpu")) | |
| pred = pred_v.get_image()[:, :, ::-1] | |
| pred = cv2.resize(pred, (800, 800)) | |
| pred = cv2.cvtColor(pred, cv2.COLOR_BGR2RGB) | |
| if version.parse(st.__version__) >= version.parse("1.11.0"): | |
| tabs = st.tabs(["Result", "Detection"] + [f"Lesion #{i}" for i in range(0, len(instances))]) | |
| lesion_tabs = tabs[2:] | |
| detection_tab = tabs[1] | |
| with tabs[0]: | |
| st.header("Image processed") | |
| st.success("Use the tabs on the right to see the detected lesions and detailed explanations for each lesion") | |
| else: | |
| tabs = [st.container() for i in range(0, len(instances)+1)] | |
| lesion_tabs = tabs[1:] | |
| detection_tab = tabs[0] | |
| state.write("Populating first tab...") | |
| with detection_tab: | |
| st.header("Detected lesions") | |
| st.image(pred) | |
| for i, (tab, box, type, scores, features) in enumerate(zip(lesion_tabs, instances.pred_boxes, instances.pred_classes, instances.probs, instances.features)): | |
| state.write(f"Populating tab for lesion #{i}...") | |
| healthy_prob = scores[-1].item() | |
| scores = scores[:-1] | |
| features = features.tolist() | |
| with tab: | |
| st.header(f"Lesion #{i}") | |
| state.write(f"Populating classes for lesion #{i}...") | |
| lesion_img = draw_box(img, box.cpu(), type, model) | |
| lesion_img = cv2.cvtColor(lesion_img, cv2.COLOR_BGR2RGB) | |
| classes = ["healty", "neoplastic", "aphthous", "traumatic"] | |
| y_pos = np.arange(len(classes)) | |
| probs = [healthy_prob] + scores.cpu().numpy().tolist() | |
| probs_fig = plt.figure() | |
| plt.bar(y_pos, probs, align="center") | |
| plt.xticks(y_pos, classes) | |
| plt.ylabel("Probability") | |
| plt.title("Class") | |
| st.subheader("Classification") | |
| col1, col2 = st.columns(2) | |
| col1.image(lesion_img) | |
| col2.pyplot(probs_fig) | |
| st.subheader("Feature space") | |
| col1, col2 = st.columns(2) | |
| state.write(f"Populating PCA for lesion #{i}...") | |
| fig = plot_pca_point(point=features, features_database=FEATURES_DATABASE, pca_model=PCA_MODEL, fig_h=800, fig_w=600, fig_dpi=100) | |
| col1.pyplot(fig) | |
| state.write(f"Populating histogram for lesion #{i}...") | |
| fig = plot_histogram_dist(point=features, features_database=FEATURES_DATABASE, fig_h=800, fig_w=600, fig_dpi=100) | |
| col2.pyplot(fig) | |
| state.write(f"Populating Gradcam++ for lesion #{i}...") | |
| st.subheader("Gradcam++") | |
| fig = plot_gradcam(model=MODEL, img=img, instance=i, fig_h=1600, fig_w=1200, fig_dpi=200, th=TH, layer="backbone.bottom_up.res5.2.conv3") | |
| st.pyplot(fig) | |
| state.write("All done...") | |
| FILE = "./test.jpg" | |
| MODEL = "./models/model.pth" | |
| PCA_MODEL = "./models/pca.pkl" | |
| FEATURES_DATABASE = "./assets/features/features.json" | |
| st.header("Explainable Oral Lesion Detection") | |
| st.markdown("""Demo for the paper [Explainable diagnosis of oral cancer via deep learning and case-based reasoning](https://mlpi.ing.unipi.it/doctoralai/) | |
| Upload an image using the form below and click on "Process" | |
| """) | |
| FILE = st.file_uploader("Image", type=["jpg", "jpeg", "png"]) | |
| TH = st.slider("Threshold", min_value=0.0, max_value=1.0, value=0.5) | |
| process = st.button("Process") | |
| state = st.empty() | |
| if process: | |
| state.write("Loading model...") | |
| model = load_model() | |
| nparr = np.fromstring(FILE.getvalue(), np.uint8) | |
| img = cv2.imdecode(nparr, cv2.IMREAD_COLOR) | |
| #img = cv2.imread(FILE) | |
| img = cv2.resize(img, (800, 800)) | |
| explain(img, model) | |