first update for only images extraction

qurator/eynollah/eynollah.py

@@ -195,6 +195,7 @@ class Eynollah:
         self.allow_scaling = allow_scaling
         self.headers_off = headers_off
         self.light_version = light_version
+        self.extract_only_images = True
         self.ignore_page_extraction = ignore_page_extraction
         self.pcgts = pcgts
         if not dir_in:
@@ -225,6 +226,7 @@ class Eynollah:
         self.model_page_dir = dir_models + "/eynollah-page-extraction_20210425"
         self.model_region_dir_p_ens = dir_models + "/eynollah-main-regions-ensembled_20210425"
         self.model_region_dir_p_ens_light = dir_models + "/eynollah-main-regions_20220314"
+        self.model_region_dir_p_ens_light_only_images_extraction = dir_models + "/eynollah-main-regions_20231127_672_org_ens_11_13_16_17_18"
         if self.textline_light:
             self.model_textline_dir = dir_models + "/eynollah-textline_light_20210425"
         else:
@@ -249,7 +251,23 @@ class Eynollah:
             
             self.ls_imgs  = os.listdir(self.dir_in)
             
-        if dir_in and not light_version:
+        if dir_in and self.extract_only_images:
+            config = tf.compat.v1.ConfigProto()
+            config.gpu_options.allow_growth = True
+            session = tf.compat.v1.Session(config=config)
+            set_session(session)
+            
+            self.model_page = self.our_load_model(self.model_page_dir)
+            self.model_classifier = self.our_load_model(self.model_dir_of_col_classifier)
+            #self.model_bin = self.our_load_model(self.model_dir_of_binarization)
+            #self.model_textline = self.our_load_model(self.model_textline_dir)
+            self.model_region = self.our_load_model(self.model_region_dir_p_ens_light_only_images_extraction)
+            #self.model_region_fl_np = self.our_load_model(self.model_region_dir_fully_np)
+            #self.model_region_fl = self.our_load_model(self.model_region_dir_fully)
+            
+            self.ls_imgs  = os.listdir(self.dir_in)
+            
+        if dir_in and not (light_version or self.extract_only_images):
             config = tf.compat.v1.ConfigProto()
             config.gpu_options.allow_growth = True
             session = tf.compat.v1.Session(config=config)
@@ -268,6 +286,7 @@ class Eynollah:
             self.ls_imgs  = os.listdir(self.dir_in)
             
             
+        
     def _cache_images(self, image_filename=None, image_pil=None):
         ret = {}
         if image_filename:
@@ -463,6 +482,27 @@ class Eynollah:
 
         return img_new, num_column_is_classified
     
+    def calculate_width_height_by_columns_extract_only_images(self, img, num_col, width_early, label_p_pred):
+        self.logger.debug("enter calculate_width_height_by_columns")
+        if num_col == 1:
+            img_w_new = 700
+        elif num_col == 2:
+            img_w_new = 900
+        elif num_col == 3:
+            img_w_new = 1500
+        elif num_col == 4:
+            img_w_new = 1800
+        elif num_col == 5:
+            img_w_new = 2200
+        elif num_col == 6:
+            img_w_new = 2500
+        img_h_new = int(img.shape[0] / float(img.shape[1]) * img_w_new)
+
+        img_new = resize_image(img, img_h_new, img_w_new)
+        num_column_is_classified = True
+
+        return img_new, num_column_is_classified
+
     def resize_image_with_column_classifier(self, is_image_enhanced, img_bin):
         self.logger.debug("enter resize_image_with_column_classifier")
         if self.input_binary:
@@ -570,6 +610,7 @@ class Eynollah:
         
         self.logger.info("Found %d columns (%s)", num_col, np.around(label_p_pred, decimals=5))
         
+        if not self.extract_only_images:
             if dpi < DPI_THRESHOLD:
                 img_new, num_column_is_classified = self.calculate_width_height_by_columns(img, num_col, width_early, label_p_pred)
                 if light_version:
@@ -581,6 +622,10 @@ class Eynollah:
                 num_column_is_classified = True
                 image_res = np.copy(img)
                 is_image_enhanced = False
+        else:
+            img_new, num_column_is_classified = self.calculate_width_height_by_columns_extract_only_images(img, num_col, width_early, label_p_pred)
+            image_res = np.copy(img_new)
+            is_image_enhanced = False
 
         self.logger.debug("exit resize_and_enhance_image_with_column_classifier")
         return is_image_enhanced, img, image_res, num_col, num_column_is_classified, img_bin
@@ -867,12 +912,14 @@ class Eynollah:
                     seg_not_base = label_p_pred[0,:,:,4]
                     ##seg2 = -label_p_pred[0,:,:,2]
                     
-                    
+                    if self.extract_only_images:
+                        seg_not_base[seg_not_base>0.3] =1
+                        seg_not_base[seg_not_base<1] =0
+                    else:
                         seg_not_base[seg_not_base>0.03] =1
                         seg_not_base[seg_not_base<1] =0
 
                     
-                    
                     seg_test = label_p_pred[0,:,:,1]
                     ##seg2 = -label_p_pred[0,:,:,2]
                     
@@ -888,11 +935,8 @@ class Eynollah:
                     seg_line[seg_line>0.1] =1
                     seg_line[seg_line<1] =0
                     
-                    
+                    if not self.extract_only_images:
                         seg_background = label_p_pred[0,:,:,0]
-                    ##seg2 = -label_p_pred[0,:,:,2]
-                    
-                    
                         seg_background[seg_background>0.25] =1
                         seg_background[seg_background<1] =0
                     ##seg = seg+seg2
@@ -908,6 +952,7 @@ class Eynollah:
                     #seg[seg==1]=0
                     #seg[seg_test==1]=1
                     seg[seg_not_base==1]=4
+                    if not self.extract_only_images:
                         seg[seg_background==1]=0
                     seg[(seg_line==1) & (seg==0)]=3
                     seg_color = np.repeat(seg[:, :, np.newaxis], 3, axis=2)
@@ -1573,6 +1618,60 @@ class Eynollah:
         q.put(slopes_sub)
         poly.put(poly_sub)
         box_sub.put(boxes_sub_new)
+        
+    def get_regions_light_v_extract_only_images(self,img,is_image_enhanced, num_col_classifier):
+        self.logger.debug("enter get_regions_light_v")
+        erosion_hurts = False
+        img_org = np.copy(img)
+        img_height_h = img_org.shape[0]
+        img_width_h = img_org.shape[1]
+
+        #model_region, session_region = self.start_new_session_and_model(self.model_region_dir_p_ens)
+
+
+        img_resized = np.copy(img)
+        
+
+
+        if not self.dir_in:
+            model_region, session_region = self.start_new_session_and_model(self.model_region_dir_p_ens_light_only_images_extraction)
+            prediction_regions_org = self.do_prediction_new_concept(True, img_resized, model_region)
+        else:
+            prediction_regions_org = self.do_prediction_new_concept(True, img_resized, self.model_region)
+            
+        #plt.imshow(prediction_regions_org[:,:,0])
+        #plt.show()
+            
+        prediction_regions_org = resize_image(prediction_regions_org,img_height_h, img_width_h )
+
+        
+        prediction_regions_org=prediction_regions_org[:,:,0]
+            
+        mask_lines_only = (prediction_regions_org[:,:] ==3)*1
+        
+        mask_texts_only = (prediction_regions_org[:,:] ==1)*1
+        
+        mask_images_only=(prediction_regions_org[:,:] ==2)*1
+        
+        polygons_lines_xml, hir_lines_xml = return_contours_of_image(mask_lines_only)
+        polygons_lines_xml = textline_con_fil = filter_contours_area_of_image(mask_lines_only, polygons_lines_xml, hir_lines_xml, max_area=1, min_area=0.00001)
+        
+        
+        polygons_of_only_texts = return_contours_of_interested_region(mask_texts_only,1,0.00001)
+        
+        polygons_of_only_lines = return_contours_of_interested_region(mask_lines_only,1,0.00001)
+        
+        text_regions_p_true = np.zeros(prediction_regions_org.shape)
+        
+        text_regions_p_true = cv2.fillPoly(text_regions_p_true, pts = polygons_of_only_lines, color=(3,3,3))
+        
+        text_regions_p_true[:,:][mask_images_only[:,:] == 1] = 2
+        
+        text_regions_p_true = cv2.fillPoly(text_regions_p_true, pts = polygons_of_only_texts, color=(1,1,1))
+        
+        polygons_of_images = return_contours_of_interested_region(text_regions_p_true, 2)
+        
+        return text_regions_p_true, erosion_hurts, polygons_lines_xml, polygons_of_images
     def get_regions_light_v(self,img,is_image_enhanced, num_col_classifier):
         self.logger.debug("enter get_regions_light_v")
         erosion_hurts = False
@@ -2825,6 +2924,8 @@ class Eynollah:
         """
         self.logger.debug("enter run")
         
+        self.extract_only_images = True
+
 
         t0_tot = time.time()
 
@@ -2836,6 +2937,20 @@ class Eynollah:
             if self.dir_in:
                 self.reset_file_name_dir(os.path.join(self.dir_in,img_name))
             
+            
+            if self.extract_only_images:
+                img_res, is_image_enhanced, num_col_classifier, num_column_is_classified = self.run_enhancement(self.light_version)
+                self.logger.info("Enhancing took %.1fs ", time.time() - t0)
+                
+                text_regions_p_1 ,erosion_hurts, polygons_lines_xml,polygons_of_images = self.get_regions_light_v_extract_only_images(img_res, is_image_enhanced, num_col_classifier)
+                #self.logger.info("Textregion detection took %.1fs ", time.time() - t1t)
+                
+                if self.plotter:
+                    self.plotter.write_images_into_directory(polygons_of_images, img_res)
+                #plt.imshow(text_regions_p_1)
+                #plt.show()
+
+            else:
                 img_res, is_image_enhanced, num_col_classifier, num_column_is_classified = self.run_enhancement(self.light_version)
                 self.logger.info("Enhancing took %.1fs ", time.time() - t0)