File size: 2,354 Bytes
e214012 |
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 53 54 55 56 57 58 59 60 61 62 63 64 65 66 |
import torch
import gradio as gr
from PIL import Image
from torchvision import transforms
from statistics import mean
Image.MAX_IMAGE_PIXELS = None
def read_img_file(f):
img = Image.open(f)
if img.mode != 'RGB':
img = img.convert('RGB')
return img
_transform_test_random=transforms.Compose([
transforms.RandomCrop((256,256)),
transforms.ToTensor(),
transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))])
_transform_test_random_vit = transforms.Compose([
transforms.RandomCrop((252,252)),
transforms.ToTensor(),
transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))])
def detect(img, model_choices):
if model_choices == "EVA-02 ViT L/14":
model = torch.load("./model_eva.pth",map_location="cpu").cpu().eval()
_transform = _transform_test_random_vit
elif model_choices == "ConvNext Large":
model = torch.load("./model_convnext.pth",map_location="cpu").cpu().eval()
_transform = _transform_test_random
elif model_choices == "EfficientNet-V2 B0":
model = torch.load("./model_effnet.pth",map_location="cpu").cpu().eval()
_transform = _transform_test_random
output = ""
with torch.inference_mode():
tmp=[]
for _ in range(10):
img_random_crop = _transform(img)
outputs = model.forward(img_random_crop.unsqueeze(0))
outputs = torch.sigmoid(outputs).cpu().numpy()
tmp.append(outputs[0][0])
output+=f"{str(tmp)}\n"
output+=f"10 try method: {mean(tmp)}\n"
# print(tmp)
# print("10 try method: ", mean(tmp))
with torch.inference_mode():
img_crop = _transform(img)
outputs = model.forward(img_crop.unsqueeze(0))
outputs = torch.sigmoid(outputs).cpu().numpy()
output+=f"1 try method: {outputs}\n"
# print("1 try method: ",outputs)
return output
model_choices = ["ConvNext Large", "EVA-02 ViT L/14", "EfficientNet-V2 B0"]
demo = gr.Interface(fn=detect,
inputs=[gr.Image(type="pil", label="Input Image"),
gr.Radio(
model_choices,
type="value",
value="EVA-02 ViT L/14",
label="Choose Detector Model",
)], outputs="text")
demo.launch() |