diff --git a/generate_gt_for_training.py b/generate_gt_for_training.py
index 9b7f02b..8ca5cd3 100644
--- a/generate_gt_for_training.py
+++ b/generate_gt_for_training.py
@@ -418,7 +418,7 @@ def visualize_textline_segmentation(dir_xml, dir_out, dir_imgs):
     help="directory of images where textline segmentation will be overlayed", )
 
 def visualize_layout_segmentation(xml_file, dir_xml, dir_out, dir_imgs):
-    assert xml_file and dir_xml, "A single xml file -xml or a dir of xml files -dx is required not both of them"
+    assert xml_file or dir_xml, "A single xml file -xml or a dir of xml files -dx is required not both of them"
     if dir_xml:
         xml_files_ind = os.listdir(dir_xml)
     else:
@@ -442,7 +442,7 @@ def visualize_layout_segmentation(xml_file, dir_xml, dir_out, dir_imgs):
         co_text, co_graphic, co_sep, co_img, co_table, co_noise, y_len, x_len = get_layout_contours_for_visualization(xml_file)
         
         
-        added_image = visualize_image_from_contours_layout(co_text['paragraph'], co_text['header'], co_text['drop-capital'], co_sep, co_img, co_text['marginalia'],  img)
+        added_image = visualize_image_from_contours_layout(co_text['paragraph'], co_text['header']+co_text['heading'], co_text['drop-capital'], co_sep, co_img, co_text['marginalia'],  img)
 
         cv2.imwrite(os.path.join(dir_out, f_name+'.png'), added_image)
 
diff --git a/gt_gen_utils.py b/gt_gen_utils.py
index a734020..0ac15a2 100644
--- a/gt_gen_utils.py
+++ b/gt_gen_utils.py
@@ -306,6 +306,7 @@ def get_layout_contours_for_visualization(xml_file):
     co_noise=[]
     
     types_text = []
+    types_graphic = []
     
     for tag in region_tags:
         if tag.endswith('}TextRegion') or tag.endswith('}Textregion'):
@@ -325,6 +326,9 @@ def get_layout_contours_for_visualization(xml_file):
                                 if len(types_text_without_paragraph) == 0:
                                     if "type" in nn.attrib:
                                         c_t_in['paragraph'].append( np.array( [ [ int(x.split(',')[0]) , int(x.split(',')[1]) ]  for x in p_h] ) )
+                                    else:
+                                        c_t_in['paragraph'].append( np.array( [ [ int(x.split(',')[0]) , int(x.split(',')[1]) ]  for x in p_h] ) )
+                                        
                                 elif len(types_text_without_paragraph) >= 1:
                                     if "type" in nn.attrib:
                                         if nn.attrib['type'] in types_text_without_paragraph:
@@ -332,10 +336,15 @@ def get_layout_contours_for_visualization(xml_file):
                                         else:
                                             c_t_in['paragraph'].append( np.array( [ [ int(x.split(',')[0]) , int(x.split(',')[1]) ]  for x in p_h] ) )
                                             
+                                    else:
+                                        c_t_in['paragraph'].append( np.array( [ [ int(x.split(',')[0]) , int(x.split(',')[1]) ]  for x in p_h] ) )
+                                            
                             else:
                                 if "type" in nn.attrib:
                                     if nn.attrib['type'] in all_defined_textregion_types:
                                         c_t_in[nn.attrib['type']].append( np.array( [ [ int(x.split(',')[0]) , int(x.split(',')[1]) ]  for x in p_h] ) )
+                                else:
+                                    c_t_in['paragraph'].append( np.array( [ [ int(x.split(',')[0]) , int(x.split(',')[1]) ]  for x in p_h] ) )
     
                             break
                         else: