shreydan/pascal-multilabel-classifier

1import torch2import torchvision.transforms as T3from timm import create_model4from safetensors.torch import load_model5import numpy as np6from pathlib import Path7import gradio as gr8 9examples = Path('./examples').glob('*')10examples = list(map(str,examples))11 12valid_tfms = T.Compose([13    T.Resize((224,224)),14    T.ToTensor(),15    T.Normalize(16        mean = (0.5,0.5,0.5),17         std = (0.5,0.5,0.5)18    )19])20 21 22model_path = 'model/swin_s3_base_224-pascal/model.safetensors'23model = create_model(24        'swin_s3_base_224',25        pretrained = False,26        num_classes = 2027)28load_model(model,model_path)29model.eval()30 31class_names = [32    "Aeroplane","Bicycle","Bird","Boat","Bottle",33    "Bus","Car","Cat","Chair","Cow","Diningtable",34    "Dog","Horse","Motorbike","Person",35    "Potted plant","Sheep","Sofa","Train","Tv/monitor"36]37 38label2id = {c:idx for idx,c in enumerate(class_names)}39id2label = {idx:c for idx,c in enumerate(class_names)}40 41 42def predict(im):43    im = valid_tfms(im).unsqueeze(0)44    with torch.no_grad():45        logits = model(im)46    47    confidences = logits.sigmoid().flatten()48    predictions = confidences > 0.549    predictions = predictions.float().numpy()50    pred_labels = np.where(predictions==1)[0]51    confidences = confidences[pred_labels].numpy()52    pred_labels = [id2label[label] for label in pred_labels]53    outputs = {l:c for l,c in zip(pred_labels, confidences)}54    return outputs55 56gr.Interface(fn=predict,57             inputs=gr.Image(type="pil"),58             outputs=gr.Label(label='the image contains:'),59             examples=examples).queue().launch()