Added function to crop detections and save them as new images

Author: theAIGuysCode

core/functions.py

@@ -28,4 +28,25 @@ def count_objects(data, by_class = False):
     else:
         counts['total object'] = num_objects
 
-    return counts
+    return counts
+
+# function for cropping each detection and saving as new image
+def crop_objects(img, data, path, allowed_classes = None):
+    boxes, scores, classes, num_objects = data
+    class_names = read_class_names(cfg.YOLO.CLASSES)
+    #create dictionary to hold count of objects for image name
+    counts = dict()
+    for i in range(num_objects):
+        # get count of class for part of image name
+        class_index = int(classes[i])
+        class_name = class_names[class_index]
+        counts[class_name] = counts.get(class_name, 0) + 1
+        # get box coords
+        xmin, ymin, xmax, ymax = boxes[i]
+        # crop detection from image (take an additional 5 pixels around all edges)
+        cropped_img = img[int(ymin)-5:int(ymax)+5, int(xmin)-5:int(xmax)+5]
+        # construct image name and join it to path for saving crop properly
+        img_name = class_name + '_' + str(counts[class_name]) + '.png'
+        img_path = os.path.join(path, img_name )
+        # save image
+        cv2.imwrite(img_path, cropped_img)

detect.py

@@ -26,10 +26,11 @@
 flags.DEFINE_list('images', './data/images/kite.jpg', 'path to input image')
 flags.DEFINE_string('output', './detections/', 'path to output folder')
 flags.DEFINE_float('iou', 0.45, 'iou threshold')
-flags.DEFINE_float('score', 0.25, 'score threshold')
+flags.DEFINE_float('score', 0.50, 'score threshold')
 flags.DEFINE_boolean('count', False, 'count objects within images')
 flags.DEFINE_boolean('dont_show', False, 'dont show image output')
 flags.DEFINE_boolean('info', False, 'print info on detections')
+flags.DEFINE_boolean('crop', False, 'crop detections from images')
 
 def main(_argv):
     config = ConfigProto()
@@ -52,6 +53,10 @@ def main(_argv):
 
         image_data = cv2.resize(original_image, (input_size, input_size))
         image_data = image_data / 255.
+        
+        # get image name by using split method
+        image_name = image_path.split('/')[-1]
+        image_name = image_name.split('.')[0]
 
         images_data = []
         for i in range(1):
@@ -95,6 +100,15 @@ def main(_argv):
         # hold all detection data in one variable
         pred_bbox = [bboxes, scores.numpy()[0], classes.numpy()[0], valid_detections.numpy()[0]]
 
+        # if crop flag is enabled, crop each detection and save it as new image
+        if FLAGS.crop:
+            crop_path = os.path.join(os.getcwd(), 'detections', 'crop', image_name)
+            try:
+                os.mkdir(crop_path)
+            except FileExistsError:
+                pass
+            crop_objects(cv2.cvtColor(original_image, cv2.COLOR_BGR2RGB), pred_bbox, crop_path)
+
         if FLAGS.count:
             # count objects found
             counted_classes = count_objects(pred_bbox, by_class = False)

detect_video.py

@@ -10,7 +10,7 @@
 from absl.flags import FLAGS
 import core.utils as utils
 from core.yolov4 import filter_boxes
-from core.functions import count_objects
+from core.functions import *
 from tensorflow.python.saved_model import tag_constants
 from PIL import Image
 import cv2
@@ -32,6 +32,7 @@
 flags.DEFINE_boolean('count', False, 'count objects within video')
 flags.DEFINE_boolean('dont_show', False, 'dont show video output')
 flags.DEFINE_boolean('info', False, 'print info on detections')
+flags.DEFINE_boolean('crop', False, 'crop detections from images')
 
 def main(_argv):
     config = ConfigProto()
@@ -40,7 +41,9 @@ def main(_argv):
     STRIDES, ANCHORS, NUM_CLASS, XYSCALE = utils.load_config(FLAGS)
     input_size = FLAGS.size
     video_path = FLAGS.video
-
+    # get video name by using split method
+    video_name = video_path.split('/')[-1]
+    video_name = video_name.split('.')[0]
     if FLAGS.framework == 'tflite':
         interpreter = tf.lite.Interpreter(model_path=FLAGS.weights)
         interpreter.allocate_tensors()
@@ -68,10 +71,12 @@ def main(_argv):
         codec = cv2.VideoWriter_fourcc(*FLAGS.output_format)
         out = cv2.VideoWriter(FLAGS.output, codec, fps, (width, height))
 
+    frame_num = 0
     while True:
         return_value, frame = vid.read()
         if return_value:
             frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+            frame_num += 1
             image = Image.fromarray(frame)
         else:
             print('Video has ended or failed, try a different video format!')
@@ -116,6 +121,24 @@ def main(_argv):
 
         pred_bbox = [bboxes, scores.numpy()[0], classes.numpy()[0], valid_detections.numpy()[0]]
 
+        # if crop flag is enabled, crop each detection and save it as new image
+        if FLAGS.crop:
+            crop_rate = 150 # capture images every so many frames (ex. crop photos every 150 frames)
+            crop_path = os.path.join(os.getcwd(), 'detections', 'crop', video_name)
+            try:
+                os.mkdir(crop_path)
+            except FileExistsError:
+                pass
+            if frame_num % crop_rate == 0:
+                final_path = os.path.join(crop_path, 'frame_' + str(frame_num))
+                try:
+                    os.mkdir(final_path)
+                except FileExistsError:
+                    pass          
+                crop_objects(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB), pred_bbox, final_path)
+            else:
+                pass
+
         if FLAGS.count:
             # count objects found
             counted_classes = count_objects(pred_bbox, by_class = False)