Merge pull request #24 from LCAS/rgb_image_only_pred

Run on rgb only if no depth

Author: yilmazabdurrah

scripts/detectron_predictor/detectron_predictor.py

@@ -99,6 +99,57 @@ def _get_catalog(self,catalog_file):
             if(__debug__): print(traceback.format_exc())
             raise Exception(e)
         return data
+    
+    def get_rgb_predictions_image(self, rgb_image, output_json_file_path='',prediction_output_dir='',image_file_name='',sample_no=1,fruit_type=FruitTypes.Strawberry):
+        rgb_image = rgb_image[:, :, :3].astype(np.uint8)
+        image_size = rgb_image.shape
+        image_size = tuple(reversed(image_size[:-1]))
+        save_json_file = True
+        try:
+            outputs = self.predictor(rgb_image)
+            predictions = outputs["instances"].to("cpu")
+            vis_aoc = AOCVisualizer(rgb_image,
+                                    metadata=self.metadata[0],
+                                    scale=self.scale,
+                                    instance_mode=self.instance_mode,
+                                    colours=self.colours,
+                                    category_ids=self.list_category_ids,
+                                    masks=self.masks,
+                                    bbox=self.bbox,
+                                    show_orientation=self.show_orientation,
+                                    fruit_type=fruit_type
+                                    )
+            start_time = datetime.now()
+            drawn_predictions = vis_aoc.draw_instance_predictions(outputs["instances"].to("cpu"))
+            end_time = datetime.now()
+            predicted_image = drawn_predictions.get_image()[:, :, ::-1].copy()            
+            
+            pred_image_dir = os.path.join(prediction_output_dir, 'predicted_images')
+            if not os.path.exists(pred_image_dir):
+                os.makedirs(pred_image_dir)
+
+            if (self.rename_pred_images):
+                f_name=f'img_{str(sample_no).zfill(6)}.png'
+                overlay_fName = os.path.join(pred_image_dir, f_name)
+                file_dir, f = os.path.split(output_json_file_path)
+                image_file_name = f_name
+                f_name = f'img_{str(sample_no).zfill(6)}.json'
+                output_json_file_path = os.path.join(file_dir, f_name)
+
+            else:
+                file_dir, f_name = os.path.split(image_file_name)
+                overlay_fName = os.path.join(pred_image_dir, f_name)
+            cv2.imwrite(overlay_fName, cv2.cvtColor(predicted_image, cv2.COLOR_BGR2RGB))
+            delta=str(end_time - start_time)
+            print(f"Predicted image saved in output folder for file {overlay_fName}, Duration: {delta}")
+            json_writer = JSONWriter(rgb_image, self.metadata[0])
+            categories_info=self.metadata[1] # category info is saved as second list
+            predicted_json_ann=json_writer.create_prediction_json(predictions, output_json_file_path, image_file_name,categories_info,image_size,1,save_json_file)
+            return predicted_json_ann,predicted_image,[]
+        except Exception as e:
+            logging.error(e)
+            if(__debug__): print(traceback.format_exc())
+            raise Exception(e)
 
     def get_predictions_image(self, rgbd_image,output_json_file_path='',prediction_output_dir='',image_file_name='',sample_no=1,fruit_type=FruitTypes.Strawberry):
 

scripts/fruit_detection.py

@@ -169,7 +169,16 @@ def __init__(self, non_ros_config_path):
                         prediction_json_output_file = os.path.join(self.prediction_json_dir, filename)+'.json'
                     self.det_predictor.get_predictions_image(rgbd_image, prediction_json_output_file, self.prediction_output_dir, image_file_name, sample_no, self.fruit_type)
                 else:
-                    self.get_logger().warn(f"Warning: No corresponding depth file: {corr_depth_file} for rgb file: {rgb_file}")
+                    self.get_logger().warn(f"Warning: No corresponding depth file: {corr_depth_file} for rgb file: {rgb_file}.\nPredicting using rgb only.")
+                    image_file_name=os.path.join(self.image_dir, rgb_file)
+                    rgb_image = cv2.imread(image_file_name)  # bgr8
+                    filename, extension = os.path.splitext(rgb_file)
+                    if (self.prediction_json_dir!=""):
+                        os.makedirs(self.prediction_json_dir, exist_ok=True)
+                        prediction_json_output_file = os.path.join(self.prediction_json_dir, filename)+'.json'
+                    
+                    self.det_predictor.get_rgb_predictions_image(rgb_image, prediction_json_output_file, self.prediction_output_dir, image_file_name, sample_no, self.fruit_type)
+                    
                 sample_no += 1
 
     def compute_pose2d(self, annotation_id, pose_dict):