| |
| import argparse |
| import torch |
| from PIL import Image |
| from torchvision import datasets, transforms |
|
|
|
|
| IMAGENET_MEAN = [0.485, 0.456, 0.406] |
| IMAGENET_STD = [0.229, 0.224, 0.225] |
| HARDCODED_WNID_TO_IDX_101 = {'n01440764': 0, 'n01443537': 1, 'n01484850': 2, 'n01491361': 3, 'n01494475': 4, 'n01496331': 5, 'n01498041': 6, 'n01514668': 7, 'n01514859': 8, 'n01531178': 9, 'n01537544': 10, 'n01560419': 11, 'n01582220': 12, 'n01592084': 13, 'n01601694': 14, 'n01608432': 15, 'n01614925': 16, 'n01622779': 17, 'n01630670': 18, 'n01632458': 19, 'n01632777': 20, 'n01644900': 21, 'n01664065': 22, 'n01665541': 23, 'n01667114': 24, 'n01667778': 25, 'n01675722': 26, 'n01677366': 27, 'n01685808': 28, 'n01687978': 29, 'n01693334': 30, 'n01695060': 31, 'n01698640': 32, 'n01728572': 33, 'n01729322': 34, 'n01729977': 35, 'n01734418': 36, 'n01735189': 37, 'n01739381': 38, 'n01740131': 39, 'n01742172': 40, 'n01749939': 41, 'n01751748': 42, 'n01753488': 43, 'n01755581': 44, 'n01756291': 45, 'n01770081': 46, 'n01770393': 47, 'n01773157': 48, 'n01773549': 49, 'n01773797': 50, 'n01774384': 51, 'n01774750': 52, 'n01775062': 53, 'n01776313': 54, 'n01795545': 55, 'n01796340': 56, 'n01798484': 57, 'n01806143': 58, 'n01818515': 59, 'n01819313': 60, 'n01820546': 61, 'n01824575': 62, 'n01828970': 63, 'n01829413': 64, 'n01833805': 65, 'n01843383': 66, 'n01847000': 67, 'n01855672': 68, 'n01860187': 69, 'n01877812': 70, 'n01883070': 71, 'n01910747': 72, 'n01914609': 73, 'n01924916': 74, 'n01930112': 75, 'n01943899': 76, 'n01944390': 77, 'n01950731': 78, 'n01955084': 79, 'n01968897': 80, 'n01978287': 81, 'n01978455': 82, 'n01984695': 83, 'n01985128': 84, 'n01986214': 85, 'n02002556': 86, 'n02006656': 87, 'n02007558': 88, 'n02011460': 89, 'n02012849': 90, 'n02013706': 91, 'n02018207': 92, 'n02018795': 93, 'n02027492': 94, 'n02028035': 95, 'n02037110': 96, 'n02051845': 97, 'n02058221': 98, 'n02077923': 99, 'n02391049': 100} |
|
|
|
|
| def preprocess_image(image_path, input_image_size): |
| image = Image.open(image_path).convert("RGB") |
| transform_list = [] |
| if image.size[0] != input_image_size or image.size[1] != input_image_size: |
| transform_list.extend([ |
| transforms.Resize(input_image_size, interpolation=Image.BICUBIC), |
| transforms.CenterCrop(input_image_size), |
| ]) |
| transform_list.extend([ |
| transforms.ToTensor(), |
| transforms.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD), |
| ]) |
| transform = transforms.Compose(transform_list) |
| tensor = transform(image).unsqueeze(0) |
| return tensor |
|
|
|
|
| def load_wnid_to_name(cls_map_path): |
| if not cls_map_path: |
| return None |
| wnid_to_name = {} |
| with open(cls_map_path, "r", encoding="utf-8") as handle: |
| for line in handle: |
| line = line.strip() |
| if not line: |
| continue |
| parts = line.split() |
| if len(parts) < 3: |
| continue |
| wnid = parts[0] |
| class_name = " ".join(parts[2:]) |
| wnid_to_name[wnid] = class_name |
| return wnid_to_name if wnid_to_name else None |
|
|
|
|
| def load_idx_to_name_from_val_dir(val_dir, wnid_to_name): |
| if not val_dir: |
| idx_to_wnid = {idx: wnid for wnid, idx in HARDCODED_WNID_TO_IDX_101.items()} |
| else: |
| dataset = datasets.ImageFolder(val_dir) |
| idx_to_wnid = {v: k for k, v in dataset.class_to_idx.items()} |
| idx_to_name = {} |
| for idx, wnid in idx_to_wnid.items(): |
| idx_to_name[idx] = wnid_to_name.get(wnid, wnid) if wnid_to_name else wnid |
| return idx_to_name if idx_to_name else None |
|
|
|
|
| def topk_from_logits(logits, idx_to_class, k=5): |
| probs = torch.softmax(logits, dim=1) |
| values, indices = torch.topk(probs, k=k, dim=1) |
| values = values.squeeze(0).tolist() |
| indices = indices.squeeze(0).tolist() |
| results = [] |
| for score, idx in zip(values, indices): |
| cls_name = idx_to_class.get(idx, str(idx)) if idx_to_class else str(idx) |
| results.append((idx, cls_name, score)) |
| return results |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser(description="TorchScript single-image inference.") |
| parser.add_argument("--torchscript", type=str, required=True, |
| help="Path to TorchScript .pt file") |
| parser.add_argument("--image_path", type=str, required=True, |
| help="Path to input image") |
| parser.add_argument("--input_image_size", type=int, default=224) |
| parser.add_argument("--val_dir", type=str, |
| default=None, |
| help="Val dir to derive class index -> wnid mapping") |
| parser.add_argument("--cls_map_path", type=str, |
| default="/scratch/general/vast/j.yan/nas_tvm/cls_map.txt", |
| help="Path to cls_map.txt for wnid -> class name mapping") |
| parser.add_argument("--topk", type=int, default=1) |
| args = parser.parse_args() |
|
|
| model = torch.jit.load(args.torchscript, map_location="cpu") |
| model.eval() |
|
|
| input_tensor = preprocess_image(args.image_path, args.input_image_size) |
| wnid_to_name = load_wnid_to_name(args.cls_map_path) |
| idx_to_class = load_idx_to_name_from_val_dir(args.val_dir, wnid_to_name) |
|
|
| with torch.no_grad(): |
| logits = model(input_tensor) |
|
|
| print("[torchscript] top-{}:".format(args.topk)) |
| for _, cls_name, score in topk_from_logits(logits, idx_to_class, k=args.topk): |
| print(f" class={cls_name} prob={score:.6f}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
|
|
