arunm8489

def draw_boxes(x, img):
    c1 = tuple(x[1:3].int())
    c2 = tuple(x[3:5].int())
    cls = int(x[-1])
    label = "{0}".format(classes[cls])
    color = (0,0,255)
    cv2.rectangle(img, c1, c2,color, 2)
    t_size = cv2.getTextSize(label, cv2.FONT_HERSHEY_PLAIN, 1 , 1)[0]
    c2 = c1[0] + t_size[0] + 3, c1[1] + t_size[1] + 4
    cv2.rectangle(img, c1, c2,color, -1)

arunm8489

list(map(lambda x: draw_boxes(x, original_image), output))
cv2.imwrite("out.png", original_image)

arunm8489

im_dim_list = torch.index_select(im_dim, 0, output[:,0].long())

scaling_factor = torch.min(conf_inp_dim/im_dim_list,1)[0].view(-1,1)
output[:,[1,3]] -= (conf_inp_dim - scaling_factor*im_dim_list[:,0].view(-1,1))/2
output[:,[2,4]] -= (conf_inp_dim - scaling_factor*im_dim_list[:,1].view(-1,1))/2
output[:,1:5] /= scaling_factor
    


# adjusting bounding box size between 0 and configuration image size

arunm8489

im_dim = original_img_dim[0], original_img_dim[1]
im_dim = torch.FloatTensor(im_dim).repeat(1,2)

#If there's a GPU availible, put the model on GPU
if CUDA:
    im_dim = im_dim_list.cuda()
    model.cuda()

#Set the model in evaluation mode
model.eval()

arunm8489

# function converting images from opencv format to torch format
def preprocess_image(img, inp_dim):
    """
    Prepare image for inputting to the neural network. 
    
    Returns processed image, original image and original image dimension  
    """

    orig_im = cv2.imread(img)
    dim = orig_im.shape[1], orig_im.shape[0]

arunm8489

CUDA = False
image_name = "dog-cycle-car.png"
nms_thesh = 0.5
#Set up the neural network
print("Loading network.....")
model = Darknet(cfgfile)
model.load_weights(weightsfile)
print("Network successfully loaded")
classes = load_classes(classfile)
print('Classes loaded')

arunm8489

# function to load the classes
def load_classes(class_file):
    fp = open(class_file, "r")
    names = fp.read().split("\n")[:-1]
    return names

arunm8489

def bounding_box_iou(box1, box2):
    """
    Returns the IoU of two bounding boxes 
    
    """
    #Get the coordinates of bounding boxes
    b1_x1, b1_y1, b1_x2, b1_y2 = box1[:,0], box1[:,1], box1[:,2], box1[:,3]
    b2_x1, b2_y1, b2_x2, b2_y2 = box2[:,0], box2[:,1], box2[:,2], box2[:,3]
    
    #get the corrdinates of the intersection rectangle

arunm8489

#Concatenate the batch_id of the image to the detection#this helps us identify which image does the detection correspond to 
#We use a linear straucture to hold ALL the detections from the batch
#the batch_dim is flattened
#batch is identified by extra batch column
#creating a row with index of images

batch_ind = image_pred_class.new(image_pred_class.size(0), 1).fill_(ind)
seq = batch_ind, image_pred_class
if not write:
    output = torch.cat(seq,1)

arunm8489

def bounding_box_iou(box1, box2):
    """
    Returns the IoU of two bounding boxes 
    
    """
    #Get the coordinates of bounding boxes
    b1_x1, b1_y1, b1_x2, b1_y2 = box1[:,0], box1[:,1], box1[:,2], box1[:,3]
    b2_x1, b2_y1, b2_x2, b2_y2 = box2[:,0], box2[:,1], box2[:,2], box2[:,3]
    
    #get the corrdinates of the intersection rectangle