Advertisement
Mochinov

Untitled

Jan 13th, 2022
1,333
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
Python 5.32 KB | None | 0 0
  1. import cv2
  2. import numpy as np
  3. import multiprocessing
  4. from yolo import Yolo
  5. import time
  6. def multiproces(frame,end_frame):
  7.     video = cv2.VideoCapture('k.m4v')
  8.     video.set(cv2.CAP_PROP_POS_FRAMES, frame)
  9.     success, kadr = video.read()
  10.     while frame <= end_frame:
  11.         bboxes, idxs, indeks = detect(kadr, target_class_id)
  12.         for box in bboxes:
  13.             box[0] = round(box[0] * fxy)
  14.             box[1] = round(box[1] * fxy)
  15.             box[2] = round(box[2] * fxy)
  16.             box[3] = round(box[3] * fxy)
  17.         multiTracker = cv2.legacy.MultiTracker_create()
  18.         kadr_res = cv2.resize(kadr, None, fx=fxy, fy=fxy)
  19.         for bbox in bboxes:
  20.             multiTracker.add(createTracker(), kadr_res, bbox)
  21.         j = 0
  22.         while j <= 30:
  23.             _, bboxes = multiTracker.update(kadr_res)
  24.             j += 5
  25.             frame += 5
  26.             video.set(cv2.CAP_PROP_POS_FRAMES, frame)
  27.             success, kadr = video.read()
  28.             kadr_res = cv2.resize(kadr, None, fx=fxy, fy=fxy)
  29.             boxes = []
  30.             for i in idxs:
  31.                 boxes.append(bboxes[i[0]])
  32.             for i in range(6):
  33.                 mas[frame-i] = boxes
  34.             draw(kadr,boxes)
  35. #индекс и сохранение
  36. #рисовка и многопоточность
  37. def draw(kadr, boxes):
  38.     for i in boxes:
  39.         x = round(i[0]*1/fxy)
  40.         y = round(i[1]*1/fxy)
  41.         w = round(i[2]*1/fxy)
  42.         h = round(i[3]*1/fxy)
  43.         cv2.rectangle(kadr, (x, y), (x + w, y + h), colors, 2)
  44.     cv2.imshow("output", kadr)
  45.     cv2.waitKey(0)
  46. video = cv2.VideoCapture('k.m4v')
  47. length = int(video.get(cv2.CAP_PROP_FRAME_COUNT))
  48. mas = []
  49. for i in range(length):
  50.     mas.append(i)
  51. fxy = 1/8
  52. colors = [255, 255, 255]#цвет рамки
  53. weights = "yolov3.weights"#Веса и классы
  54. config = "yolov3.cfg"
  55. labels = "yolov3.txt"
  56. target_class_id = 0
  57. conf_thresh = 0.5
  58. nms_thresh = 0.5
  59. net = Yolo(config, weights, labels, conf_thresh, nms_thresh)
  60. def createTracker():#Создание отдельного трекера
  61.     tracker = cv2.legacy.TrackerCSRT_create()
  62.     return tracker
  63. def detect(kadr, target_class_id):#Выхов функцции детекции
  64.     data, idxs, indeks = net.detect(kadr, target_class_id)
  65.     return data, idxs, indeks
  66. if __name__ == '__main__':
  67.     video = cv2.VideoCapture('k.m4v')
  68.     length = int(video.get(cv2.CAP_PROP_FRAME_COUNT))
  69.     delenie = round(length/4)
  70.     p1 = multiprocessing.Process(target=multiproces, args=(0,delenie*1,))
  71.     p2 = multiprocessing.Process(target=multiproces, args=(delenie*1,delenie*2,))
  72.     p3 = multiprocessing.Process(target=multiproces, args=(delenie*2,delenie*3,))
  73.     p4 = multiprocessing.Process(target=multiproces, args=(delenie*3,length,))
  74.     p1.start()
  75.     p2.start()
  76.     p3.start()
  77.     p4.start()
  78.  
  79.  
  80.     p1.join()
  81.     p2.join()
  82.     p3.join()
  83.     p4.join()
  84.  
  85.  
  86.  
  87.  
  88.  
  89.  
  90.  
  91.  
  92.  
  93.  
  94.  
  95.  
  96.  
  97.  
  98.  
  99.  
  100. import cv2
  101. import numpy as np
  102.  
  103. class Yolo:
  104.     def __init__(self, cfg, weights, names, conf_thresh, nms_thresh): #Создание и загрузка весов детекции объетов
  105.         self.ct = conf_thresh
  106.         self.nmst = nms_thresh
  107.         self.net = cv2.dnn.readNet(weights, cfg)
  108.         print("Finished: " + str(weights))
  109.         self.classes = [0]
  110.         layer_names = self.net.getLayerNames()
  111.         self.output_layers = [layer_names[i[0]-1] for i in self.net.getUnconnectedOutLayers()]
  112.     def detect(self, img, target_id):#Финальная функция возвращает индекс параметры рисования и координаты объетов
  113.         b, c, ids, idxs = self.get_detection_data(img, target_id)
  114.         return b, idxs, ids
  115.     def get_detection_data(self, img, target_id):#Подгон разрешения и обработка по уже найденым объетам
  116.         layer_outputs = self.get_inf(img)
  117.         height, width = img.shape[:2]
  118.         b, c, ids, idxs = self.thresh(layer_outputs, width, height, target_id)
  119.         return b, c, ids, idxs
  120.     def get_inf(self, img):#Подгон разрешения и запуск нейросети
  121.         blob = cv2.dnn.blobFromImage(img, 1 / 255.0, (416, 416), swapRB=True, crop=False)
  122.         self.net.setInput(blob)
  123.         layer_outputs = self.net.forward(self.output_layers)
  124.         return layer_outputs
  125.     def thresh(self, layer_outputs, width, height, target_id):#Сам алгоритм классов индексов и координат
  126.         boxes = []
  127.         confidences = []
  128.         class_ids = []
  129.         for output in layer_outputs:
  130.             for detection in output:
  131.                 scores = detection[5:]
  132.                 class_id = np.argmax(scores)
  133.                 confidence = scores[class_id]
  134.                 if confidence > self.ct and class_id == target_id:
  135.                     box = detection[0:4] * np.array([width, height, width, height])
  136.                     (cx, cy, w, h) = box.astype('int')
  137.                     tx = int(cx - (w / 2))
  138.                     ty = int(cy - (h / 2))
  139.                     boxes.append([tx, ty, int(w), int(h)])
  140.                     confidences.append(float(confidence))
  141.                     class_ids.append(class_id)
  142.         idxs = cv2.dnn.NMSBoxes(boxes, confidences, self.ct, self.nmst)
  143.         return boxes, confidences, class_ids, idxs
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement