Guest User

Untitled

a guest
May 21st, 2021
128
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
  1. import cv2
  2. import numpy as np
  3. path_to_frozen_inference_graph = 'frozen_inference_graph_coco.pb'
  4. path_coco_model= 'mask_rcnn_inception_v2_coco_2018_01_28.pbtxt'
  5. VIDEO = 'race.mp4'
  6. net = cv2.dnn.readNetFromTensorflow(path_to_frozen_inference_graph,path_coco_model)
  7. colors = np.random.randint(0, 255, (80, 3))
  8.  
  9. video = cv2.VideoCapture(VIDEO)
  10. while True:
  11.     grabbed,frame=video.read()
  12.     if not grabbed:
  13.         break
  14.  
  15.     img=cv2.resize(frame,(650,550))
  16.     height, width, _ = img.shape
  17.     black_image = np.zeros((height, width, 3), np.uint8)
  18.     black_image[:] = (150, 150, 0)
  19.     blob = cv2.dnn.blobFromImage(img, swapRB=True)
  20.     net.setInput(blob)
  21.     boxes, masks = net.forward(["detection_out_final", "detection_masks"])
  22.     detection_count = boxes.shape[2]
  23.     for i in range(detection_count):
  24.         box = boxes[0, 0, i]
  25.         class_id = box[1]
  26.         score = box[2]
  27.         if score < 0.5:
  28.             continue
  29.  
  30.         x = int(box[3] * width)
  31.         y = int(box[4] * height)
  32.         x2 = int(box[5] * width)
  33.         y2 = int(box[6] * height)
  34.  
  35.  
  36.         roi = black_image[y: y2, x: x2]
  37.         roi_height, roi_width, _ = roi.shape
  38.  
  39.         mask = masks[i, int(class_id)]
  40.  
  41.         mask = cv2.resize(mask, (roi_width, roi_height))
  42.  
  43.         _, mask = cv2.threshold(mask, 0.5, 255, cv2.THRESH_BINARY)
  44.  
  45.         cv2.rectangle(img, (x, y), (x2, y2), (255, 0, 0), 3)
  46.         contours, _ = cv2.findContours(np.array(mask, np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
  47.         color = colors[int(class_id)]
  48.         for cnt in contours:
  49.             cv2.fillPoly(roi, [cnt], (int(color[0]), int(color[1]), int(color[2])))
  50.  
  51.  
  52.  
  53.     cv2.imshow("Black image", black_image)
  54.  
  55.     key = cv2.waitKey(1) & 0xFF
  56.     if key == ord("q"):
  57.              break
  58.  
  59.  
  60. video.release()
  61. cv2.destroyAllWindows()
RAW Paste Data