Saatvik/Multilabel

1import os2import torch3from torch import nn4from model import create_resnext5from torchvision import transforms6 7transform = transforms.Compose([8    transforms.Resize((224, 224)),9    transforms.ToTensor(),10    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])11])12 13model = create_resnext()14model = nn.DataParallel(model)15model.load_state_dict(torch.load('resnext.pth', map_location='cpu'))16 17classes = ['person',18  'bicycle',19  'car',20  'motorcycle',21  'airplane',22  'bus',23  'train',24  'truck',25  'boat',26  'traffic light',27  'fire hydrant',28  'stop sign',29  'parking meter',30  'bench',31  'bird',32  'cat',33  'dog',34  'horse',35  'sheep',36  'cow',37  'elephant',38  'bear',39  'zebra',40  'giraffe',41  'backpack',42  'umbrella',43  'handbag',44  'tie',45  'suitcase',46  'frisbee',47  'skis',48  'snowboard',49  'sports ball',50  'kite',51  'baseball bat',52  'baseball glove',53  'skateboard',54  'surfboard',55  'tennis racket',56  'bottle',57  'wine glass',58  'cup',59  'fork',60  'knife',61  'spoon',62  'bowl',63  'banana',64  'apple',65  'sandwich',66  'orange',67  'broccoli',68  'carrot',69  'hot dog',70  'pizza',71  'donut',72  'cake',73  'chair',74  'couch',75  'potted plant',76  'bed',77  'dining table',78  'toilet',79  'tv',80  'laptop',81  'mouse',82  'remote',83  'keyboard',84  'cell phone',85  'microwave',86  'oven',87  'toaster',88  'sink',89  'refrigerator',90  'book',91  'clock',92  'vase',93  'scissors',94  'teddy bear',95  'hair drier',96  'toothbrush']97 98from time import time99 100def predict(img):101    start = time()102    103    newimg = transform(img).unsqueeze(dim=0)104    model.eval()105    with torch.inference_mode():106        ylogit = model(newimg).detach()107        yprob = torch.sigmoid(ylogit)108    109    names_with_probs = {classes[i]: float(yprob[0][i]) for i in range(len(classes))}110    111    predtime = time() - start112    113    return names_with_probs, predtime114 115import gradio as gr116 117example_list = [["examples/" + example] for example in os.listdir("examples/")]118 119title = 'Multilabel Classifier'120description = "A ResNeXt-50 feature extractor computer vision model to identify objects present in an image out of 80 classes"121article = "Github repo-> https://github.com/Saatvik-Sinha/DSEG660-Multilabel-Classification-Challenge"122 123demo = gr.Interface(fn=predict,124                    inputs=gr.Image(type="pil"),125                    outputs=[gr.Label(num_top_classes=80, label="Predictions"),126                             gr.Number(label="Prediction time (s)")],127                    examples=example_list, 128                    title=title,129                    description=description,130                    article=article)131 132demo.launch()133