class MyDarkflow: def __init__(self): self.confidence = 0.4 self.option = { 'model': 'cfg/yolo.cfg', 'load': 'bin/yolo.weights', 'threshold': self.confidence, 'gpu': 0.7 } self.tfnet = TFNet(self.option) self.colors = { 'car': (238, 23, 23), 'truck': (0, 255, 21), 'bus': (3, 0, 255), 'person': (0, 255, 243) } def highlight_vehicles(self, img): results = self.tfnet.return_predict(img) for result in results: #Pega posição e tipo do veículo tl = (result['topleft']['x'], result['topleft']['y']) br = (result['bottomright']['x'], result['bottomright']['y']) label = result['label'] #Dá cor à label if label not in self.colors: self.colors[label] = 200 * np.random.rand(3) #Desenha quadrado em volta do veículo img = cv2.rectangle(img, tl, br, self.colors[label], 3) img = cv2.putText(img, label, tl, cv2.FONT_HERSHEY_COMPLEX, 1, (0, 0, 0), 2) return (img, len(results))
#options = {"model": "cfg/yolo-face.cfg", "load": "weight/yolo-face_final.weights", "threshold": 0.1, "gpu": 1.0} options = { "model": "cfg/yolo-face.cfg", "load": "weight/yolo-face_4000.weights", "threshold": 0.1, "gpu": 1.0 } tfnet = TFNet(options) count = 0 tracker = cv2.TrackerMIL_create() cam = cv2.VideoCapture(0) while True: _, camcv = cam.read() camcv = cv2.resize(camcv, (448, 448)) result = tfnet.return_predict(camcv) num_people = len(result) for i in range(num_people): #print(result[i]['topleft']) top_left_x = result[i]['topleft']['x'] top_left_y = result[i]['topleft']['y'] bottom_right_x = result[i]['bottomright']['x'] bottom_right_y = result[i]['bottomright']['y'] if count == 0 and top_left_x != 0: count += 1 bbox = (top_left_x, top_left_y, bottom_right_x, bottom_right_y) ok = tracker.init(camcv, bbox) elif count == 0: continue else:
def predict(videoPath): if os.path.exists("yolo.csv"): os.remove("yolo.csv") with open('yolo.csv', 'a') as writeFile: writer = csv.writer(writeFile) writer.writerow(['FrameNumber', 'PredictionString']) writeFile.close() options = {"model": "cfg/455.cfg", "load": 35250, "threshold": 0.4} frameNum = 0 tfnet = TFNet(options) skip = False fileDir = os.listdir(os.getcwd() + "/frames") for filename in fileDir: img = cv2.imread(os.getcwd() + "/frames/" + filename, cv2.IMREAD_COLOR) #print(os.getcwd() + "/frames/" + filename) while img is None: print("image is none") img = cv2.imread(os.getcwd() + "/frames/" + filename, cv2.IMREAD_COLOR) time.sleep(2) if img is None: skip = True os.remove(os.getcwd() + "/frames/" + filename) break saved_img = img if not skip: img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) print("not skipping") # use YOLO to predict the image result = tfnet.return_predict(img) #print("result equals " + str(result)) img.shape id = 0 csvString = "" for i in range(0, len(result)): tl = (result[i]['topleft']['x'], result[i]['topleft']['y']) br = (result[i]['bottomright']['x'], result[i]['bottomright']['y']) label = result[i]['label'] #print(result[i]) id = id + 1 csvString += result[i]['label'] + " " + str(result[i]['confidence']) + " " + str(result[i]['bottomright']['x']) + " " \ + str(result[i]['bottomright']['x']) + " " + str(result[i]['topleft']['y']) + " " + str(result[i]['bottomright']['y']) + " " saved_img = cv2.rectangle(saved_img, tl, br, (0, 255, 0), 7) saved_img = cv2.putText(saved_img, str(id) + " " + label, tl, cv2.FONT_HERSHEY_COMPLEX, 1, (0, 0, 0), 2) cv2.imwrite(os.getcwd() + "/output-images/" + filename, saved_img) with open('yolo.csv', 'a') as writeFile: writer = csv.writer(writeFile) writer.writerow([frameNum, csvString]) frameNum = frameNum + 1 csvString = "" writeFile.close() print("Finished processing image " + str(frameNum) + "/" + str(len(fileDir))) skip = False print("Making video") img_array = [] for index in range(0, len(os.listdir(os.getcwd() + "/output-images/"))): print("sorted file name : " + str(img)) img = cv2.imread(os.getcwd() + "/output-images/" + str(index) + ".jpg", cv2.IMREAD_COLOR) height, width, layers = img.shape size = (width, height) img_array.append(img) print(len(img_array)) out = cv2.VideoWriter( os.getcwd() + "/" + videoPath.split("/")[-1].split(".")[0] + "-yolo.avi", cv2.VideoWriter_fourcc(*'DIVX'), 15, size) for i in range(len(img_array)): out.write(img_array[i]) out.release() return json.dumps({ 'csvpath': os.getcwd() + "/yolo.csv", 'videopath': os.getcwd() + "/" + videoPath.split("/")[-1].split(".")[0] + "-yolo.avi" })
while (True): # Capture frame-by-frame ret, frame = cap.read() frame = preprocess(frame) start_time = datetime.datetime.now() if is_target_exist: # ok, bbox = tracker.update(frame) bbox = siamFC.track(frame) # Draw bounding box p1 = (int(bbox[0]), int(bbox[1])) p2 = (int(bbox[0] + bbox[2]), int(bbox[1] + bbox[3])) cv2.rectangle(frame, p1, p2, (255, 0, 0), 2, 1) else: results = tfnet.return_predict(frame) if not is_target_exist else [] is_contain_person = False maximum_index = 0 maximum_confidence = -1 for i, res in enumerate(results): if res['label'] == 'person': is_contain_person = True if res['confidence'] > maximum_confidence: maximum_confidence = res['confidence'] maximum_index = i if is_contain_person: res = results[maximum_index] bbox = (res['topleft']['x'], res['topleft']['y'], res['bottomright']['x'] - res['topleft']['x'], res['bottomright']['y'] - res['topleft']['y'])