| import torch | |
| from PIL import Image | |
| from torchvision import transforms | |
| from architecture import ResNetLungCancer | |
| import gradio as gr | |
| device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") | |
| model = ResNetLungCancer(num_classes=4) | |
| model.load_state_dict(torch.load('Model/lung_cancer_detection_model.pth', map_location=device)) | |
| model = model.to(device) | |
| model.eval() | |
| preprocess = transforms.Compose([ | |
| transforms.Resize(256), | |
| transforms.CenterCrop(224), | |
| transforms.ToTensor(), | |
| transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) | |
| ]) | |
| class_names = ['Adenocarcinoma', 'Large Cell Carcinoma', 'Normal', 'Squamous Cell Carcinoma'] | |
| def predict(image): | |
| image = Image.fromarray(image.astype('uint8'), 'RGB') | |
| input_tensor = preprocess(image).unsqueeze(0).to(device) | |
| with torch.no_grad(): | |
| output = model(input_tensor) | |
| predicted_class = torch.argmax(output, dim=1).item() | |
| return class_names[predicted_class] | |
| iface = gr.Interface( | |
| fn=predict, | |
| inputs=gr.Image(), | |
| outputs=gr.Label(num_top_classes=1), | |
| examples=[ | |
| ["Data/test/large.cell.carcinoma/000108.png"], | |
| ["Data/test/normal/7 - Copy (3).png"] | |
| ] | |
| ) | |
| iface.launch(share=True) |