orcablack/testdep

1import gradio as gr2import cv23import requests4 5import os6from PIL import Image7import timm8import torch9from torchvision.transforms import transforms10import numpy as np11from PIL import ImageFile12import matplotlib.pyplot as plt13import warnings14import glob15 16 17warnings.filterwarnings("ignore")18ImageFile.LOAD_TRUNCATED_IMAGES = True19 20 21def predict(image, model, device, class_name):22 23    prediction_transform = transforms.Compose([transforms.Resize(size=(224, 224)),24                                               transforms.ToTensor(),25                                               transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])])26    try:27      image = prediction_transform(image)[:3,:,:].unsqueeze(0)28    except:29      image = image.convert('RGB')30      image = prediction_transform(image)[:3,:,:].unsqueeze(0)31 32    if device == 'cuda':33        if torch.cuda.is_available():34            image = image.cuda()35        else:36            print("You don't have cuda")37 38    with torch.no_grad():39      model.eval()40      pred = model(image)41 42 43    idx = torch.argmax(pred)44 45    prob = pred[0][idx].item()*10046 47    return prob, class_name[idx]48 49 50model = timm.create_model('resnet50', pretrained=True)51 52model.fc = torch.nn.Sequential(torch.nn.Linear(2048, 256),53                                    torch.nn.Dropout(0.2),54                                    torch.nn.ReLU(),55                                    torch.nn.Linear(256, 64),56                                    torch.nn.Dropout(0.2),57                                    torch.nn.ReLU(),58                                    torch.nn.Linear(64, 32),59                                    torch.nn.Dropout(0.2),60                                    torch.nn.ReLU(),61                                    torch.nn.Linear(32, 4),62                                    torch.nn.Softmax()63                                    )64 65model.load_state_dict(torch.load('model_ResNet50_acc_max.pt',map_location=torch.device('cpu')))66 67display_prob = True68show=True69#path = glob.glob('*.png')70 71def show_preds_image(path):72    #for image in path:73    img = Image.open(path)74        # if show:75        #     plt.imshow(img)76        #     plt.show()77    #img = cv2.imread(path)78    class_name = ['adenocarcinoma',79                'large.cell.carcinoma',80                'normal',81                'squamous.cell.carcinoma']82    prob, result = predict(img, model, 'cpu', class_name)83    if display_prob:84        print('Probability of {} : {:.6f}'.format(result, prob))85   86    return result, prob87 88inputs_image = [89    gr.components.Image(type="filepath", label="Input Image"),90]91 92interface_image = gr.Interface(93    fn=show_preds_image,94    inputs=inputs_image,95    outputs="text",96    title="Cancer Detector App using data from Kaggle",97    cache_examples=False,98).launch()99 100