File size: 4,370 Bytes
0325572 9128cc2 0325572 9128cc2 0325572 9128cc2 0325572 74fe289 0325572 74fe289 9128cc2 0325572 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 |
# import pytorch related dependencies
import torch
from PIL import Image
from torch import nn
import numpy as np
import torchvision as torchvision
import torchvision.transforms as transforms
from pytorch_grad_cam import GradCAM, ScoreCAM, GradCAMPlusPlus, AblationCAM, XGradCAM, EigenCAM, FullGrad
from pytorch_grad_cam.utils.image import show_cam_on_image
import gradio as gr
import timm
# model setup
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
classes = ["Preparation",
"CalotTriangleDissection",
"ClippingCutting",
"GallbladderDissection",
"GallbladderPackaging",
"CleaningCoagulation",
"GallbladderRetraction"]
model = timm.create_model('efficientnet_b3')# This is a very well known network but it is designed for 1000 classes and not just cats and dogs this is why we need the next line
model.classifier = nn.Linear(1536, 7)
#state_dict_trained = torch.hub.load_state_dict_from_url("https://github.com/tobiascz/demotime/raw/main/checkpoints/ham10k_checkpoint_mobile_0.82_epoch24.pt", model_dir=".", map_location = device)
import os
print(os.getcwd())
state_dict_trained = torch.load('checkpoints/state_dict_timm_effnet_b3_e6_val_f1=0.75.pt', map_location=torch.device(device))
sd = model.state_dict()
print(state_dict_trained.keys())
for k,v in sd.items():
if not "classifier" in k:
sd[k] = state_dict_trained[f'model.model.{k}']
sd['classifier.weight'] = state_dict_trained['model.fc_phase.weight']
sd['classifier.bias'] = state_dict_trained['model.fc_phase.bias']
model.load_state_dict(sd) ## Here we load the trained weights (state_dict) in our model
model.eval() # This
# image pre-processing
norm_mean = (0.4914, 0.4822, 0.4465)
norm_std = (0.2023, 0.1994, 0.2010)
transform = transforms.Compose([ # resize image to the network input size
transforms.CenterCrop((400,400)),
transforms.ToTensor(),
transforms.Normalize(norm_mean, norm_std)
])
# convert tensot to numpy array
def tensor2npimg(tensor, mean, std):
# inverse of normalization
tensor = tensor.clone()
mean_tensor = torch.as_tensor(list(mean), dtype=tensor.dtype, device=tensor.device).view(-1,1,1)
std_tensor = torch.as_tensor(list(std), dtype=tensor.dtype, device=tensor.device).view(-1,1,1)
tensor.mul_(std_tensor).add_(mean_tensor)
# convert tensor to numpy format for plt presentation
npimg = tensor.numpy()
npimg = np.transpose(npimg,(1,2,0)) # C*H*W => H*W*C
return npimg
# draw Grad-CAM on image
# target layer could be any layer before the final attention block
# Some common choices are:
# FasterRCNN: model.backbone
# Resnet18 and 50: model.layer4[-1]
# VGG and densenet161: model.features[-1]
# mnasnet1_0: model.layers[-1]
# ViT: model.blocks[-1].norm1
# SwinT: model.layers[-1].blocks[-1].norm1
def image_grad_cam(model, input_tensor, input_float_np, target_layers):
cam = GradCAM(model=model, target_layers=target_layers, use_cuda=False)
grayscale_cam = cam(input_tensor=input_tensor, aug_smooth=True, eigen_smooth=True)
grayscale_cam = grayscale_cam[0, :]
return show_cam_on_image(input_float_np, grayscale_cam, use_rgb=True)
# config the predict function for Gradio, input type of image is numpy.nparray
def predict(input_img):
# numpy.nparray -> PIL.Image
leasionExample = Image.fromarray(input_img.astype('uint8'), 'RGB')
# normalize the image to fit the input size of our model
leasion_tensor = transform(leasionExample)
input_float_np = tensor2npimg(leasion_tensor, norm_mean, norm_std)
leasion_tensor = leasion_tensor.unsqueeze(dim=0)
# predict
with torch.no_grad():
outputs = model(leasion_tensor)
outputs = torch.exp(outputs)
# probabilities of all classes
pred_softmax = torch.softmax(outputs, dim=1).cpu().numpy()[0]
# class with hightest probability
# diagnostic suggestions
# return label dict and suggestion
return {classes[i]: float(pred_softmax[i]) for i in range(len(classes))}
# start gradio application
gr.Interface(
fn=predict,
inputs=gr.inputs.Image(),
outputs=[gr.outputs.Label(label="Predict Result")],
examples=[['images/video01_000014_prep.png'],['images/video01_001403.png'],['images/video01_001528_pack.png']],
title="Surgical Workflow Classifier"
).launch() |