|
import gradio as gr |
|
import gradio as gr |
|
import torch |
|
import torchvision.transforms as transforms |
|
from PIL import Image |
|
from torchvision import models |
|
|
|
|
|
targets = ['Negative','Positive'] |
|
|
|
|
|
transform = 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]), |
|
]) |
|
|
|
|
|
def classify_image(img): |
|
img = Image.fromarray(img.astype('uint8'), 'RGB') |
|
img = transform(img).unsqueeze(0) |
|
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") |
|
img = img.to(device) |
|
|
|
with torch.no_grad(): |
|
prediction = torch.nn.functional.softmax(model(img)[0], dim=0) |
|
confidences = {targets[i]: float(prediction[i]) for i in range(2)} |
|
return confidences |
|
|
|
|
|
demo = gr.Interface(fn=classify_image, |
|
inputs=gr.Image(width=224, height=224), |
|
outputs=gr.Label(num_top_classes=2), |
|
examples=["examples/negative.jpeg", "examples/positive.jpeg"]) |
|
|
|
demo.launch(share=True, debug=True) |
|
|