diff --git a/Application/Classifiers/Classifier.py b/Application/Classifiers/Classifier.py index 589c27d..d5dc6e8 100644 --- a/Application/Classifiers/Classifier.py +++ b/Application/Classifiers/Classifier.py @@ -5,57 +5,102 @@ import numpy as np import tensorflow as tf import cv2 -import os +import os import json - from Application.Classifiers.ClassifierInterface import ClassifierInterface class Classifier(ClassifierInterface): def __init__(self): - self.threshold = .5 + print("1") + self.model_path = os.path.join(os.path.dirname(__file__), "./class1.pb") + self.odapi = self.DetectorAPI(path_to_ckpt=self.model_path) + self.threshold = 0.9 with open(os.path.join(os.path.dirname(__file__), "coco_map.json")) as file: mapping = json.load(file) self.classes = dict() for element in mapping: - self.classes[element["id"]-1] = element["display_name"] + self.classes[element["id"]] = element["display_name"] - self.net = cv2.dnn.readNet(os.path.join(os.path.dirname(__file__),"yolov4.weights"),os.path.join(os.path.dirname(__file__),"yolov4.cfg")) - #self.net.setPreferableBackend(cv2.dnn.DNN_BACKEND_CUDA) - #self.net.setPreferableTarget(cv2.dnn.DNN_TARGET_CUDA) - self.layer_names = self.net.getLayerNames() - self.outputlayers = [self.layer_names[i[0] - 1] for i in self.net.getUnconnectedOutLayers()] - - print("Classifier Initiated") - - def tagLayer(self, imgs): + def detect(self, img): # get the results from the net + boxes, scores, classes, num = self.odapi.process_frame(img) + res = False + for i in range(len(boxes)): + if scores[i] > self.threshold: + if classes[i] in self.classes: + #print(self.classes[classes[i]]) + return self.classes[classes[i]] - results = [] - for i, contours in enumerate(imgs[19:20]): - #print(i) - for contour in contours: - height,width,channels = contour.shape - - dim = max(height, width) - if dim > 320: - img2 = np.zeros(shape=[dim, dim, 3], dtype=np.uint8) - else: - img2 = np.zeros(shape=[320,320, 3], dtype=np.uint8) - img2[:height,:width] = contour - blob = cv2.dnn.blobFromImage(img2,1/256,(320,320),(0,0,0),True,crop=False) #reduce 416 to 320 - self.net.setInput(blob) - outs = self.net.forward(self.outputlayers) - for out in outs: - for detection in out: - scores = detection - class_id = np.argmax(scores) - confidence = scores[class_id] - if confidence > self.threshold: - if self.classes[class_id] not in results: - cv2.imshow("changes x", img2) - cv2.waitKey(10) & 0XFF - results.append(self.classes[class_id]) - #print(self.classes[x], score) - - return results + + def tagLayer(self, data): + res = [] + for cnts in data: + for cnt in cnts: + if cnt.any(): + x = self.detect(cnt) + if x not in res: + res.append(x) + if x is not None: + print(x) + cv2.imshow("changes x", cnt) + cv2.waitKey(10) & 0XFF + return res + + # Detector API can be changed out given the I/O remains the same + # this way you can use a different N-Net if you like to + class DetectorAPI: + def __init__(self, path_to_ckpt): + self.path_to_ckpt = path_to_ckpt + gpus = tf.config.experimental.list_physical_devices('GPU') + if gpus: + try: + for gpu in gpus: + tf.config.experimental.set_memory_growth(gpu, True) + except RuntimeError as e: + print(e) + self.detection_graph = tf.Graph() + with self.detection_graph.as_default(): + od_graph_def = tf.GraphDef() + with tf.gfile.GFile(self.path_to_ckpt, 'rb') as fid: + serialized_graph = fid.read() + od_graph_def.ParseFromString(serialized_graph) + tf.import_graph_def(od_graph_def, name='') + + self.default_graph = self.detection_graph.as_default() + self.sess = tf.Session(graph=self.detection_graph) + + # Definite input and output Tensors for detection_graph + self.image_tensor = self.detection_graph.get_tensor_by_name('image_tensor:0') + # Each box represents a part of the image where a particular object was detected. + self.detection_boxes = self.detection_graph.get_tensor_by_name('detection_boxes:0') + # Each score represent how level of confidence for each of the objects. + # Score is shown on the result image, together with the class label. + self.detection_scores = self.detection_graph.get_tensor_by_name('detection_scores:0') + self.detection_classes = self.detection_graph.get_tensor_by_name('detection_classes:0') + self.num_detections = self.detection_graph.get_tensor_by_name('num_detections:0') + + def process_frame(self, image): + # Expand dimensions since the trained_model expects images to have shape: [1, None, None, 3] + image_np_expanded = np.expand_dims(image, axis=0) + # Actual detection. + + (boxes, scores, classes, num) = self.sess.run( + [self.detection_boxes, self.detection_scores, self.detection_classes, self.num_detections], + feed_dict={self.image_tensor: image_np_expanded}) + + im_height, im_width,_ = image.shape + boxes_list = [None for i in range(boxes.shape[1])] + for i in range(boxes.shape[1]): + boxes_list[i] = ( + int(boxes[0, i, 0] * im_height), + int(boxes[0, i, 1] * im_width), + int(boxes[0, i, 2] * im_height), + int(boxes[0, i, 3] * im_width) + ) + + return boxes_list, scores[0].tolist(), [int(x) for x in classes[0].tolist()], int(num[0]) + + def close(self): + self.sess.close() + self.default_graph.close() diff --git a/Application/Classifiers/ClassifierInterface.py b/Application/Classifiers/ClassifierInterface.py index 1aa11bc..227171b 100644 --- a/Application/Classifiers/ClassifierInterface.py +++ b/Application/Classifiers/ClassifierInterface.py @@ -1,6 +1,6 @@ class ClassifierInterface: - def tagLayers(self, layers): + def tagLayer(self, layers): """takes layers, returns list (len(), same as input) of lists with tags for corresponfing layers""" pass \ No newline at end of file diff --git a/Application/Config.py b/Application/Config.py index e762d0c..d766ef0 100644 --- a/Application/Config.py +++ b/Application/Config.py @@ -13,7 +13,7 @@ class Config: "maxLength": None, "ttolerance": 60, "videoBufferLength": 16, - "noiseThreashold": 0.1, + "noiseThreashold": 0.3, "noiseSensitivity": 3/4, "LayersPerContour": 5, "averageFrames": 10 diff --git a/Application/Layer.py b/Application/Layer.py index ad7b7b3..9d909b0 100644 --- a/Application/Layer.py +++ b/Application/Layer.py @@ -100,7 +100,7 @@ class Layer: # the loop isn't nessecary (?) if the number of clusters is known, since it isn't the loop tries to optimize while True: - kmeans = KMeans(init="random", n_clusters=clusterCount, n_init=5, max_iter=300, random_state=42) + kmeans = KMeans(init="random", n_clusters=clusterCount, n_init=10, max_iter=300, random_state=42) kmeans.fit(mapped) labels = list(kmeans.labels_) diff --git a/Application/LayerManager.py b/Application/LayerManager.py index 7b0b82a..a191724 100644 --- a/Application/LayerManager.py +++ b/Application/LayerManager.py @@ -27,7 +27,7 @@ class LayerManager: def cleanLayers(self): self.freeMin() self.sortLayers() - #self.cleanLayers2() + self.cleanLayers2() self.freeMax() def removeStaticLayers(self): diff --git a/Application/VideoReader.py b/Application/VideoReader.py index 8a766bf..7781bfb 100644 --- a/Application/VideoReader.py +++ b/Application/VideoReader.py @@ -24,6 +24,7 @@ class VideoReader: self.buffer = Queue(config["videoBufferLength"]) self.vc = cv2.VideoCapture(videoPath) self.stopped = False + self.getWH() if setOfFrames is not None: self.listOfFrames = sorted(setOfFrames) diff --git a/main.py b/main.py index 445fb70..61cd80d 100644 --- a/main.py +++ b/main.py @@ -14,8 +14,8 @@ def main(): start = time.time() config = Config() - config["inputPath"] = os.path.join(os.path.dirname(__file__), "generate test footage/out.mp4") - #config["importPath"] = os.path.join(os.path.dirname(__file__), "output/short.txt") + config["inputPath"] = os.path.join(os.path.dirname(__file__), "generate test footage/3.mp4") + config["importPath"] = os.path.join(os.path.dirname(__file__), "output/short.txt") config["outputPath"] = os.path.join(os.path.dirname(__file__), "output/short.mp4") vr = VideoReader(config) @@ -34,7 +34,7 @@ def main(): layerManager = LayerManager(config, layers) layerManager.cleanLayers() - #layerManager.tagLayers() + layerManager.tagLayers() layers = layerManager.layers exporter = Exporter(config) exporter.export(layers, raw=False)