Hey there, I was wondering if anyone can help me with my inference code.

  
      github.com
  
  
    amrit-das/Custom-Model-Training-PyTorch/blob/master/predict.py 22
import torch
import torch.nn as nn
from torchvision.models import resnet18
from torchvision.transforms import transforms
import matplotlib.pyplot as plt
import numpy as np
from torch.autograd import Variable
import torch.functional as F
from PIL import Image
import os
import sys
import argparse

parser = argparse.ArgumentParser(description = 'To Predict from a trained model')
parser.add_argument('-i','--image', dest = 'image_name', required = True, help='Path to the image file')
parser.add_argument('-m','--model', dest = 'model_name', required = True, help='Path to the model')
parser.add_argument('-n','--num_class',dest = 'num_classes', required = True, help='Number of training classes')
args = parser.parse_args()

path_to_model = "./models/"+args.model_name


  This file has been truncated. show original

  
  
    
    
  
  


So, my script is predicting the class of the image. Say I have classes as a ball, a bat and an apple. My code predicts its a ball ( suppose I passed an image of the ball). My question is how can I draw a bounding rectangle around the ball.