How to draw bounding rectangles on the detected parts of the image

Amrit_Das · December 8, 2018, 10:49am

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

amrit-das/Custom-Model-Training-PyTorch/blob/master/predict.py

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.

albanD · December 8, 2018, 1:08pm

From a very quick read it looks like your model is only doing classification no? btw the link is broken as you changed the name of the file.

You will need a model that does detection so that it will provide you with bounding boxes for objects it detects.

Amrit_Das · December 9, 2018, 7:55am

I got your point could you please point me at a example or a project, I can get reference from ??

kowshik_thopalli · December 9, 2018, 8:49am

Hi. If you didnt know already the class of problems you are referring to are called “object detection”.

There a multiple deep learning approaches to this problem such as R-CNN, Fast R-CNN, YOLO etc.
Here is an example of one such implementation https://github.com/eriklindernoren/PyTorch-YOLOv3.git
I strongly suggest you this lecture for an overview https://www.youtube.com/watch?v=nDPWywWRIRo