def find_houghp_lines(image, rho, theta, threshold, minLineLength, maxLineGap): #img = util.create_canny_edge_image(image) houghp_lines = shape_detector.find_lines_in_image_houghp( image, rho, theta, threshold, minLineLength, maxLineGap) log(f"Found {len(houghp_lines)} Hough lines") return houghp_lines
def _join_lines_with_association_symbols(self): """ Joins the found lines with the advanced association shapes, such as inheritance, aggregation and so on. :return: """ lines_entities = self.get_generic_entities(types=[ClassDiagramTypes.ASSOCIATION_ENTITY]) symbol_entities = self.get_generic_entities(types=[ClassDiagramTypes.ASSOCIATION_SYMBOL]) for l in lines_entities: line = l.shapes[0] # GenericEntity of type ASSOCIATION_ENTITY always has just one shape, which is a Line line_start = line.start_xy() line_end = line.end_xy() for s in symbol_entities: symbol_bb = s.bounding_box() if util.is_point_in_area(line_start, symbol_bb) or util.is_point_in_area(line_end, symbol_bb): s = s.shapes[0] # TODO: Don't assume we have only one shape! l.add_shape(s) if s.shape is ShapeType.TRIANGLE: l.set(constants.STR_GENERIC_ENTITY_LABEL_NAME, "Inheritance") elif s.shape is ShapeType.RECTANGLE: l.set(constants.STR_GENERIC_ENTITY_LABEL_NAME, "Aggregation") l.type = ClassDiagramTypes.ASSOCIATION_ENTITY_ADVANCED # Change type from simple to advanced association log("Association line was joined with symbol")
def find_hough_lines(image, rho, theta, threshold): #img = util.create_canny_edge_image(image) hough_lines = shape_detector.find_lines_in_image_hough( image, rho, theta, threshold) log(f"Found {len(hough_lines)} HoughP lines") return hough_lines
def convert(self): log("transform to class primitives") self.generic_entities = self.generic_entities + self._extract_classes() self.generic_entities = self.generic_entities + self._extract_associations() self._join_lines_with_association_symbols() self._link_associations_with_classes() return self.generic_entities
def draw_contours_on_image(contours, image, color=(0, 255, 0)): """ Draws the given contours onto the given image, :param contours: Contours that are drawn. :param image: The image the contours are being drawn onto. :return: """ image = image.copy() util.log(f"Draw {len(contours)} contours") cv2.drawContours(image, contours, -1, color, 2) return image
def get_next_line_with_corresponding_point(point, lines): if point is None: return None else: lines = list(filter(lambda x: x.contains_point(point) is False, lines)) log(f"continue with {len(lines)} lines") for l in lines: log(f"check point {point}") closest_point = get_closest_corresponding_point(point, l) if closest_point is None: return get_opposite_line_end(point, l) else: return get_next_line_with_corresponding_point(closest_point, lines) return point
def export(self): log("Exporting image with labeled basic shapes") entities = [] for s in self.shape_detector.get_shapes(): if util.has_no_contour_children(s.contour_index, self.shape_detector.hierarchy): ge = GenericEntity() ge.add_shape(s) ge.set(constants.STR_GENERIC_ENTITY_LABEL_NAME, s.shape_name()) entities.append(ge) self.image = draw_util.draw(self.image, entities) return self.image
def _extract_advanced_associations(self, image): """ Tries to extract the associations such as inheritance, aggregation, composition between classes. :param image: The image the associations are extracted from :return: An array of GenericEntities, were each GenericEntity contains an extracted association """ found_associations = [] # Extract inheritance shapes shapes, _, _ = self.shape_detector.find_shapes_in_image(image) for shape in shapes: if shape.shape is ShapeType.TRIANGLE or shape.shape is ShapeType.RECTANGLE: log(f"Advanced association found: {shape}") assoc = GenericEntity(ClassDiagramTypes.ASSOCIATION_SYMBOL) assoc.add_shape(shape) found_associations.append(assoc) log(f"{len(found_associations)} advanced associations found") return found_associations
def _extract_simple_associations(self, image): """ Tries to extract the associations that are simple lines between classes. :param image: The image the associations are extracted from :return: An array of GenericEntities, were each GenericEntity contains the extracted association """ line_detector = LineDetector() line_detector.init(image) line_detector.find_lines() lines = line_detector.merge_lines() log(f"{len(lines)} lines found") found_associations = [] for l in lines: new_assoc = GenericEntity(ClassDiagramTypes.ASSOCIATION_ENTITY) new_assoc.add_shape(l) found_associations.append(new_assoc) log(f"{len(found_associations)} simple associations found") return found_associations
def nothing(arg): print(arg) rho = cv2.getTrackbarPos(tbRho, hough_winname) + 1 theta = cv2.getTrackbarPos(tbTheta, hough_winname) + 1 threshold = cv2.getTrackbarPos(tbThreshold, hough_winname) + 1 theta = (np.pi / 180) * (theta) # HoughP Lines houghp_img = shape_detector.image.copy() houghp_img = util.create_canny_edge_image(houghp_img) houghp_lines = shape_detector.find_lines_in_image_houghp(houghp_img, rho, theta, threshold) log(f"Found {len(houghp_lines)} Hough lines") for x in range(0, len(houghp_lines)): for x1, y1, x2, y2 in houghp_lines[x]: cv2.line(houghp_img, (x1, y1), (x2, y2), (0, 0, 255), 2) # Hough Lines hough_img = shape_detector.image.copy() hough_img = util.create_canny_edge_image(hough_img) hough_lines = shape_detector.find_lines_in_image_hough(hough_img, rho, theta, threshold) log(f"Found {len(hough_lines)} HoughP lines") for x in range(0, len(hough_lines)): for rho, theta in hough_lines[x]: a = np.cos(theta) b = np.sin(theta) x0 = a * rho y0 = b * rho x1 = int(x0 + 1000 * (-b)) y1 = int(y0 + 1000 * (a)) x2 = int(x0 - 1000 * (-b)) y2 = int(y0 - 1000 * (a)) cv2.line(hough_img, (x1, y1), (x2, y2), (0, 0, 255), 2) cv2.imshow("Hough Lines", hough_img) cv2.imshow("HoughP Lines", houghp_img) pass
def image_interaction(ch, image, img_path): # C - Canny Edge if ch == 99: image = util.create_canny_edge_image(image) # D - Dilate if ch == 100: log("Dilate") image = util.dilate(image) # E - Erode if ch == 101: log("Erode") image = util.erode(image) # C - Find Classes if ch == 97: shape_detector = ShapeDetector(img_path) shapes = shape_detector.find_shapes() diagram_converter = DiagramTypeDetector.find_converter(shape_detector) entities = diagram_converter._extract_classes() image = draw_util.draw_entities_on_image(entities, image) log(f"{len(shapes)} shapes in image found") # S - Extract Shapes if ch == 115: shape_detector = ShapeDetector(img_path) shapes = shape_detector.find_shapes() shape_detector.save_found_shapes() log(f"{len(shapes)} shapes in image found") for s in shapes: util.print_contour_details(s.contour) # B - Binary if ch == 98: image = util.create_binary_img(image) # I - Inverted if ch == 105: image = util.create_inverted_image(image) # R - Reset if ch == 114: shape_detector = ShapeDetector(img_path) image = shape_detector.image original = image.copy() return image
def _extract_classes(self): """ Extracts the class entities from a class diagram sketch. :return: An array of GenericEntities, were each GenericEntity contains a found class """ """ Extracts the class entities from a class diagram sketch. :return: An array of GenericEntities, were each GenericEntity contains a found class """ log(f"Extract classes from {len(self.shape_detector.shapes)} shapes") found_classes = [] sorted_contours = self.shape_detector.sort_contours_by_parent() class_counter = 0 for k, v in sorted_contours.items(): v = [v for k,v in enumerate(v) if util.area_contour(v) > ClassDiagramConverter.MIN_AREA_CLASS_RECTANGLES] contour_groups = util.group_contours_by_x_pos(v) for group_key, group_value in contour_groups.items(): # Create class entities if len(group_value) == 3: new_class = GenericEntity(ClassDiagramTypes.CLASS_ENTITY) new_class.set(constants.STR_GENERIC_ENTITY_LABEL_NAME, f"Class {class_counter}") # Add shapes to class new_class.add_shape(self.shape_detector.create_shape(group_value[0])) new_class.add_shape(self.shape_detector.create_shape(group_value[1])) new_class.add_shape(self.shape_detector.create_shape(group_value[2])) # new_class.set("name_contour", group_value[0]) # new_class.set("attribute_contour", group_value[1]) # new_class.set("method_contour", group_value[2]) found_classes.append(new_class) class_counter += 1 log(f"{len(found_classes)} class entities found") return found_classes
def __init__(self, image=None, options=None): self.orig_image = None """ Reference to the original image. A working copy will be created from this image. """ self.image = None """ Working copy of the original image. All image processing happens will be applied on this image. """ self.options = options """ Options that define the behaviour of the ShapeDetector. """ self.preprocessed_image = None """ Preprocessed working copy image. """ self.shapes = [] """ Holds all found shapes. """ self.contours = None self.hierarchy = None if image is not None: self._load(image) util.log("ShapeDetector initialized")
def export(self): log("Exporting class diagram to image") # Label classes class_entities = self.converter.get_generic_entities(types=[ClassDiagramTypes.CLASS_ENTITY]) log(f"\t... with {len(class_entities)} classes") self.image = draw_util.draw_bounding_boxes(self.image, class_entities, labels=True) # Extract text from class entities if 'ocr' in self.opts and self.opts['ocr']: for c in class_entities: for s in c.shapes: s.ocr() # Draw bounding boxes of advanced associations advanced_association_entities = self.converter.get_generic_entities( types=[ClassDiagramTypes.ASSOCIATION_ENTITY_ADVANCED]) log(f"\t... with {len(advanced_association_entities)} advanced associations") self.image = draw_util.draw_bounding_boxes(self.image, advanced_association_entities, color=constants.COLOR_RED, labels=True) # Draw normal entities association_entities = self.converter.get_generic_entities(types=[ClassDiagramTypes.ASSOCIATION_ENTITY]) log(f"\t... with {len(association_entities)} normal associations") self.image = draw_util.draw_entities_on_image(self.image, association_entities, color=constants.COLOR_YELLOW) # Print relations between classes association_entities = association_entities + advanced_association_entities for i, assoc in enumerate(association_entities): from_class = assoc.get(ClassDiagramConverter.STR_ASSOC_FROM) to_class = assoc.get(ClassDiagramConverter.STR_ASSOC_TO) if from_class is not None and to_class is not None: log(f"Association {i} of type {assoc.get(constants.STR_GENERIC_ENTITY_LABEL_NAME)} points " f"from {from_class.get(constants.STR_GENERIC_ENTITY_LABEL_NAME)} " f"to {to_class.get(constants.STR_GENERIC_ENTITY_LABEL_NAME)}") return self.image
def _link_associations_with_classes(self): """ Links the found classes and associations with each other. """ log(f"Try linking classes with associations") class_entities = self.get_generic_entities(types=[ClassDiagramTypes.CLASS_ENTITY]) assoc_entities = self.get_generic_entities(types=[ClassDiagramTypes.ASSOCIATION_ENTITY]) advanced_entities = self.get_generic_entities(types=[ClassDiagramTypes.ASSOCIATION_ENTITY_ADVANCED]) # Link class entities with remaining associations for c in class_entities: class_bounding_box = c.bounding_box(adjustment=constants.BOUNDING_BOX_ADJUSTMENT) # ... with advanced associations for a in advanced_entities: for advanced_shape in a.shapes: if type(advanced_shape) is Shape: advanced_bounding_box = advanced_shape.bounding_box() if util.do_bounding_boxes_intersect(advanced_bounding_box, class_bounding_box) or util.do_bounding_boxes_intersect(class_bounding_box, advanced_bounding_box): a.set(ClassDiagramConverter.STR_ASSOC_FROM, c) elif type(advanced_shape) is Line: line_start = advanced_shape.start_xy() line_end = advanced_shape.end_xy() if util.is_point_in_area(line_start, class_bounding_box) or util.is_point_in_area(line_end, class_bounding_box): a.set(ClassDiagramConverter.STR_ASSOC_TO, c) # ... with simple associations for a in assoc_entities: line = a.shapes[0] # GenericEntity of type ASSOCIATION_ENTITY always has just one shape, which is a Line line_start = line.start_xy() line_end = line.end_xy() if util.is_point_in_area(line_start, class_bounding_box): a.set(ClassDiagramConverter.STR_ASSOC_FROM, c) log("FROM association found") elif util.is_point_in_area(line_end, class_bounding_box): a.set(ClassDiagramConverter.STR_ASSOC_TO, c) log("TO association found")
# Line Segment Detector lsd_image = cv2.imread(img_path) lsd_image = imutils.resize(lsd_image, width=700) lsd_gray = cv2.cvtColor(lsd_image, cv2.COLOR_BGR2GRAY) lsd_gauss = cv2.GaussianBlur(lsd_gray,(9,9),0) lsd_edges = cv2.Canny(lsd_gauss, 100, 150) #cv2.imshow(wincanny, lsd_edges) LSD = cv2.createLineSegmentDetector() lines, width, prec, nfa = LSD.detect(lsd_edges) lsd_color = cv2.cvtColor(lsd_gray, cv2.COLOR_GRAY2BGR) log(f"{len(lines)} lines through LSD found") all_lines = [] for i in range(len(lines)): for x1, y1, x2, y2 in lines[i]: line = Line(Point(x1,y1), Point(x2, y2)) all_lines.append(line) # Merged lines merged_lines = purge_lines(all_lines) log(f"Lines were merged. Still {len(merged_lines)} lines") # path_lines = [] # for l in merged_lines: # #path_line = find_path(l, merged_lines) # short_circuit_end_point = short_circuit_line(l, merged_lines)
# Press D for dilation # Press C for Canny Edge # Press H for Hough Lines # Press P for HoughP Lines # Loop for get trackbar pos and process it while True: ch = cv2.waitKey(5) # C - Canny Edge if ch == 99: image = util.create_canny_edge_image(image, min=50, max=150) # D - Dilate if ch == 100: log("Dilate") image = util.dilate(image) # E - Erode if ch == 101: log("Erode") image = util.erode(image) # 1 - Draw contours if ch == 49: shape_detector.load(image) shapes = shape_detector.find_shapes() contours = [s.contour for s in shape_detector.shapes] image = draw_util.draw_labeled_contours(contours, shape_detector.hierarchy, image)
from detector.detector import * from detector.export.basic_shape_image_exporter import BasicShapeImageExporter ap = argparse.ArgumentParser() if __name__ == '__main__': ap.add_argument("-i", "--image", required=True, help="Path to the image you want to detect") ap.add_argument("-s", "--save", required=False, help="Path the result will be saved at") args = vars(ap.parse_args()) img_path = args["image"] output_path = args["save"] print(img_path) if os.path.isfile(img_path): shape_detector = ShapeDetector(img_path) shapes = shape_detector.find_shapes() # Print shapes for s in shapes: util.log(f"\t{s}") exporter = BasicShapeImageExporter(shape_detector.image, shape_detector) img = exporter.export() util.save_image(img, output_path) else: raise ValueError("Image doesn't exist")
#img_path = "img/class.jpeg" #img_path = "img/class_many.jpeg" #img_path = "img/class2.jpeg" #img_path = "img/usecase.jpeg" #img_path = "img/circles.jpeg" #img_path = "img/ocr_test.jpeg" #img_path = "img/ocr_test2.jpeg" # ap.add_argument("-i", "--image", required=True, help="Path to the image you want to detect") # args = vars(ap.parse_args()) # print(args["image"]) # Detect all shapes shape_detector = ShapeDetector(img_path) shapes = shape_detector.find_shapes() log(f"{len(shapes)} shapes in image found") # Detect type of diagram diagram_converter = DiagramTypeDetector.find_converter(shape_detector) # Convert shapes into diagram class_entities = diagram_converter._extract_classes() assoc_entities = diagram_converter._extract_associations() # contours = [c.get("name_contour") for c in entities] +\ # [c.get("attribute_contour") for c in entities] +\ # [c.get("method_contour") for c in entities] #img = shape_detector.get_image_remove_shape_type(ShapeType.RECTANGLE) # Draw class contours
def convert(self): log("transform to use case primitives")
required=False, help="Set this parameter in order to toggle the shape detection.", action="store_true") if __name__ == '__main__': init_args() args = vars(ap.parse_args()) img_path = args["image"] output_path = args["save"] print(img_path) if os.path.isfile(img_path): opts = {'ocr': args['ocr'], 'contour_epsilon': args['epsilon']} util.log(f"Passed options: {str(opts)}") shape_detector = ShapeDetector(img_path, opts) img = shape_detector.image if not args['custom']: # Start default diagram detection shapes = shape_detector.find_shapes() # Print shapes for s in shapes: util.log(f"\t{s}") diagram_converter = DiagramTypeDetector.find_converter( shape_detector) diagram_converter.convert()
def __init__(self): self.edge_image = None self.lines = [] self.LSD = None util.log("LineDetector initialized")