Spaces:
Sleeping
Sleeping
### 1. Imports and class names setup ### | |
import gradio as gr | |
import os | |
import torch | |
from model import create_ViT | |
from timeit import default_timer as timer | |
# Setup class names | |
class_names = ['Normal', 'Malignant'] | |
### 2. Model and transforms preparation ### | |
ViT, manual_transforms = create_ViT() | |
# Load saved weights | |
ViT.load_state_dict( | |
torch.load(f='Final_ViT_Model_50_Epochs.pth', | |
map_location=torch.device('cpu') # Load model to cpu | |
) | |
) | |
### 3. Predict function ### | |
def predict(img): | |
start_time = timer() | |
img = manual_transforms(img).unsqueeze(0) | |
ViT.eval() | |
with torch.inference_mode(): | |
pred_prob = torch.sigmoid(ViT(img).squeeze()) | |
pred_label_idx = int(torch.round(pred_prob)) | |
if pred_label_idx: | |
pred_labels_and_preds = {class_names[pred_label_idx-1]: float(pred_prob.item())} | |
else: | |
if 1-float(pred_prob.item()) < 75: | |
pred_labels_and_preds = {class_names[pred_label_idx]: 1-float(pred_prob.item()) } | |
else: | |
pred_labels_and_preds = {class_names[pred_label_idx+1]: 1-float(pred_prob.item())} | |
end_time = timer() | |
pred_time = round(end_time-start_time, 4) | |
return pred_labels_and_preds, pred_time | |
### 4. Gradio App ### | |
# Create interface for gradio | |
title = 'SkinGuard' | |
description = 'An AI model that can predict whether a CLOSE-UP picture of skin is normal or shows a sign of Skin Cancer.\n\nINSTRUCTIONS:\n\n 1. Take a picture of the target area\n\n 2. IMPORTANT: Crop the image to ONLY show skin lesion / target area.' | |
article = 'This is our demo for the DEDA Entrepreneurship Competition 2024. Created by Imran, Brian, Lukas, and Rohit, all in Baez CSE Period 5. \n\n IMPORTANT NOTE: If you have followed the instruction in the description of this demo and STILL get malignant results constantly, do NOT take it as an official diagnosis. Contact a healthcare professional for more information if you are concerned.' | |
# Create examples list | |
example_list = [['examples/' + example] for example in os.listdir('examples')] | |
# Create the gradio demo | |
demo = gr.Interface(fn=predict,# maps intputs to outputs | |
inputs=gr.Image(type='pil'), | |
outputs=[gr.Label(num_top_classes=1, label='Predictions'), | |
gr.Label(label='Prediction Time (s)')], | |
examples=example_list, | |
title=title, | |
description=description, | |
article=article) | |
# Launch the model!! | |
demo.launch(debug=False, | |
share=True) | |