import time import torch import numpy as np import matplotlib.pyplot as plt import matplotlib.patches as patches def boxes_iou(box1, box2): """ Returns the IOU between `box1` and `box2` (i.e intersection area divided by union area) """ # Get the Width and Height of each bounding box width_box1 = box1[2] height_box1 = box1[3] width_box2 = box2[2] height_box2 = box2[3] # Calculate the area of the each bounding box area_box1 = width_box1 * height_box1 area_box2 = width_box2 * height_box2 # Find the vertical edges of the union of the two bounding boxes mx = min(box1[0] - width_box1/2.0, box2[0] - width_box2/2.0) Mx = max(box1[0] + width_box1/2.0, box2[0] + width_box2/2.0) # Calculate the width of the union of the two bounding boxes union_width = Mx - mx # Find the horizontal edges of the union of the two bounding boxes my = min(box1[1] - height_box1/2.0, box2[1] - height_box2/2.0) My = max(box1[1] + height_box1/2.0, box2[1] + height_box2/2.0) # Calculate the height of the union of the two bounding boxes union_height = My - my # Calculate the width and height of the area of intersection of the two bounding boxes intersection_width = width_box1 + width_box2 - union_width intersection_height = height_box1 + height_box2 - union_height # If the the boxes don't overlap then their IOU is zero if intersection_width <= 0 or intersection_height <= 0: return 0.0 # Calculate the area of intersection of the two bounding boxes intersection_area = intersection_width * intersection_height # Calculate the area of the union of the two bounding boxes union_area = area_box1 + area_box2 - intersection_area # Calculate the IOU iou = intersection_area/union_area return iou def nms(boxes, iou_thresh): """ Performs Non maximal suppression technique to `boxes` using `iou_thresh` threshold """ # print(boxes.shape) # If there are no bounding boxes do nothing if len(boxes) == 0: return boxes # Create a PyTorch Tensor to keep track of the detection confidence # of each predicted bounding box det_confs = torch.zeros(len(boxes)) # Get the detection confidence of each predicted bounding box for i in range(len(boxes)): det_confs[i] = boxes[i][4] # Sort the indices of the bounding boxes by detection confidence value in descending order. # We ignore the first returned element since we are only interested in the sorted indices _,sortIds = torch.sort(det_confs, descending = True) # Create an empty list to hold the best bounding boxes after # Non-Maximal Suppression (NMS) is performed best_boxes = [] # Perform Non-Maximal Suppression for i in range(len(boxes)): # Get the bounding box with the highest detection confidence first box_i = boxes[sortIds[i]] # Check that the detection confidence is not zero if box_i[4] > 0: # Save the bounding box best_boxes.append(box_i) # Go through the rest of the bounding boxes in the list and calculate their IOU with # respect to the previous selected box_i. for j in range(i + 1, len(boxes)): box_j = boxes[sortIds[j]] # If the IOU of box_i and box_j is higher than the given IOU threshold set # box_j's detection confidence to zero. if boxes_iou(box_i, box_j) > iou_thresh: box_j[4] = 0 return best_boxes def detect_objects(model, img, iou_thresh, nms_thresh): # Start the time. This is done to calculate how long the detection takes. start = time.time() # Set the model to evaluation mode. model.eval() # Convert the image from a NumPy ndarray to a PyTorch Tensor of the correct shape. # The image is transposed, then converted to a FloatTensor of dtype float32, then # Normalized to values between 0 and 1, and finally unsqueezed to have the correct # shape of 1 x 3 x 416 x 416 img = torch.from_numpy(img.transpose(2,0,1)).float().div(255.0).unsqueeze(0) # Feed the image to the neural network with the corresponding NMS threshold. # The first step in NMS is to remove all bounding boxes that have a very low # probability of detection. All predicted bounding boxes with a value less than # the given NMS threshold will be removed. list_boxes = model(img, nms_thresh) # Make a new list with all the bounding boxes returned by the neural network boxes = list_boxes[0][0] + list_boxes[1][0] + list_boxes[2][0] # Perform the second step of NMS on the bounding boxes returned by the neural network. # In this step, we only keep the best bounding boxes by eliminating all the bounding boxes # whose IOU value is higher than the given IOU threshold boxes = nms(boxes, iou_thresh) # Stop the time. finish = time.time() # Print the time it took to detect objects print('\n\nIt took {:.3f}'.format(finish - start), 'seconds to detect the objects in the image.\n') # Print the number of objects detected print('Number of Objects Detected:', len(boxes), '\n') return boxes def load_class_names(namesfile): # Create an empty list to hold the object classes class_names = [] # Open the file containing the COCO object classes in read-only mode with open(namesfile, 'r') as fp: # The coco.names file contains only one object class per line. # Read the file line by line and save all the lines in a list. lines = fp.readlines() # Get the object class names for line in lines: # Make a copy of each line with any trailing whitespace removed line = line.rstrip() # Save the object class name into class_names class_names.append(line) return class_names def print_objects(boxes, class_names): print('Objects Found and Confidence Level:\n') for i in range(len(boxes)): box = boxes[i] if len(box) >= 7 and class_names: cls_conf = box[5] cls_id = box[6] print('%i. %s: %f' % (i + 1, class_names[cls_id], cls_conf)) def plot_boxes(img, boxes, class_names, plot_labels, color = None): # Define a tensor used to set the colors of the bounding boxes colors = torch.FloatTensor([[1,0,1],[0,0,1],[0,1,1],[0,1,0],[1,1,0],[1,0,0]]) # Define a function to set the colors of the bounding boxes def get_color(c, x, max_val): ratio = float(x) / max_val * 5 i = int(np.floor(ratio)) j = int(np.ceil(ratio)) ratio = ratio - i r = (1 - ratio) * colors[i][c] + ratio * colors[j][c] return int(r * 255) # Get the width and height of the image width = img.shape[1] height = img.shape[0] # Create a figure and plot the image fig, a = plt.subplots(1,1) a.imshow(img) # Plot the bounding boxes and corresponding labels on top of the image for i in range(len(boxes)): # Get the ith bounding box box = boxes[i] # Get the (x,y) pixel coordinates of the lower-left and lower-right corners # of the bounding box relative to the size of the image. x1 = int(np.around((box[0] - box[2]/2.0) * width)) y1 = int(np.around((box[1] - box[3]/2.0) * height)) x2 = int(np.around((box[0] + box[2]/2.0) * width)) y2 = int(np.around((box[1] + box[3]/2.0) * height)) # Set the default rgb value to red rgb = (1, 0, 0) # Use the same color to plot the bounding boxes of the same object class if len(box) >= 7 and class_names: cls_conf = box[5] cls_id = box[6] classes = len(class_names) offset = cls_id * 123457 % classes red = get_color(2, offset, classes) / 255 green = get_color(1, offset, classes) / 255 blue = get_color(0, offset, classes) / 255 # If a color is given then set rgb to the given color instead if color is None: rgb = (red, green, blue) else: rgb = color # Calculate the width and height of the bounding box relative to the size of the image. width_x = x2 - x1 width_y = y1 - y2 # Set the postion and size of the bounding box. (x1, y2) is the pixel coordinate of the # lower-left corner of the bounding box relative to the size of the image. rect = patches.Rectangle((x1, y2), width_x, width_y, linewidth = 2, edgecolor = rgb, facecolor = 'none') # Draw the bounding box on top of the image a.add_patch(rect) # If plot_labels = True then plot the corresponding label if plot_labels: # Create a string with the object class name and the corresponding object class probability conf_tx = class_names[cls_id] + ': {:.1f}'.format(cls_conf) # Define x and y offsets for the labels lxc = (img.shape[1] * 0.266) / 100 lyc = (img.shape[0] * 1.180) / 100 # Draw the labels on top of the image a.text(x1 + lxc, y1 - lyc, conf_tx, fontsize = 12, color = 'k', bbox = dict(facecolor = rgb, edgecolor = rgb, alpha = 0.6)) plt.axis("off") plt.savefig("output.jpg") plt.show()