Spaces:
Sleeping
Sleeping
File size: 1,367 Bytes
2794d4e |
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 |
from model import AQC_NET
import torch
import torch.nn as nn
import torchvision.transforms as T
from PIL import Image
import numpy as np
import gradio as gr
model = AQC_NET(pretrain=True,num_label=5)
def predict(image_name):
model.eval()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)
inputs = preprocess(image_name)
inputs = inputs.to(device)
with torch.no_grad():
outputs = model(inputs.unsqueeze(0))
values, indices = torch.topk(outputs, k=5)
print(values,indices)
return {i.item(): v.item() for i, v in zip(indices[0], values.detach()[0])}
def preprocess(image_name):
transforms = T.Compose([
T.Resize((256,256)),
T.ToTensor(),
T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
image = transforms(image_name)
return image
def run_gradio():
title = "AQC_NET PH"
description = "trial AQC_NET"
examples = ["test_image.jpg","test_img.jpg"]
inputs = [
gr.inputs.Image(type="pil", label="Input Image")
]
gr.Interface(
predict,
inputs,
outputs = 'label',
title=title,
description=description,
examples=examples,
theme="huggingface",
).launch(debug=True, enable_queue=True)
#print(predict("test_image.jpg"))
run_gradio() |