Beispiel #1
0
def detect_and_save(video: Video, start=0, end=None, detector=None):
    """
    Create a file containing object detection results of frames inside in_frame_folder

    :param video: Object of Video class
    :param start: Starting Frame Number (default: 1)
    :param end: Ending Frame Number (default: last frame)
    :param detector: Detector to use for object detection
    """

    in_frame_folder = os.path.abspath(video.frames_path)
    out_file_name = os.path.abspath(video.detection_path)

    # print(in_frame_folder, out_file_name)
    print(detector)
    pickler = BATRPickle(out_file=out_file_name)

    sorted_frames_filenames = sorted(os.listdir(in_frame_folder))

    for frame_n, frame_filename in enumerate(sorted_frames_filenames,
                                             start=start):
        frame_filename = os.path.join(in_frame_folder, frame_filename)
        detected_objects = detector.detect(frame_filename)

        pickler.pickle(detected_objects, f"frame{frame_n:06d}")

        # print(f"Frame # {frame_n}")
        if frame_n == end:
            break

    del pickler
Beispiel #2
0
def improved_fill_missing_detection(video: Video):
    in_detection_result_file = os.path.abspath(video.detection_path)

    pickler = BATRPickle(in_file=in_detection_result_file)

    tracking_path = os.path.join(os.path.dirname(in_detection_result_file),
                                 "improved_detections.tar.gz")
    output_pickler = BATRPickle(out_file=tracking_path)

    sorted_frames_abs_filenames = sorted([
        os.path.join(video.frames_path, filename)
        for filename in os.listdir(video.frames_path)
    ])
    last_serial = 0
    previously_detected_objects = []
    for frame_n, frame_filename in enumerate(sorted_frames_abs_filenames,
                                             start=0):
        # print(f"Frame # {frame_n}\n{'-' * 10}")

        frame = cv2.imread(frame_filename)

        detected_objects = pickler.unpickle(f"frame{frame_n:06d}")

        # Find Out which object is Tracked and which is not
        for obj in detected_objects:
            set_tracker(obj, frame)

            nn_index, nn_distance = nearest_neighbor(
                obj, previously_detected_objects)
            if nn_distance > 0.5 and not math.isinf(nn_distance):
                previously_detected_objects[
                    nn_index].tracked_in_next_frame = True
                obj.serial = previously_detected_objects[nn_index].serial
            else:
                # This is new object
                obj.serial = last_serial
                last_serial += 1

        # Seperate objects that are not tracked
        missed_obj = [
            obj for obj in previously_detected_objects
            if not hasattr(obj, "tracked_in_next_frame")
        ]

        # Track these missed object in current frame
        for obj in missed_obj:
            nn_index, nn_distance = nearest_neighbor(obj, detected_objects)
            nn_index = nn_index

            # Double Check. May be it is already being tracked.

            if not math.isinf(nn_index):
                intersect = intersection(obj.bbox(),
                                         detected_objects[nn_index].bbox())

                if obj.w * obj.h == 0 or intersect / (obj.w * obj.h) > 0.7:
                    continue
                # cv2.rectangle(frame, obj.pt_a(), obj.pt_b(), COLOR['blue'], 6)
            #
            # if nn_distance > 0.5 and not math.isinf(nn_distance):
            #     continue

            ok, bbox = obj.tracker.update(frame)
            bbox = int(bbox[0]), int(bbox[1]), int(bbox[2]), int(bbox[3])

            # Found in current frame by tracker
            if ok:
                lost_and_found = DetectedObject(obj.type,
                                                obj.probability,
                                                _x=bbox[0],
                                                _y=bbox[1],
                                                _w=bbox[2],
                                                _h=bbox[3])
                lost_and_found.correct_bbox(frame)
                lost_and_found.tracker = obj.tracker
                lost_and_found.serial = obj.serial
                detected_objects.append(lost_and_found)

        previously_detected_objects = list(detected_objects)

        # Display Rectangle around objects and checking whether their
        # bbox are correct

        detected_objects_without_trackers = []
        for obj in detected_objects:
            cv2.rectangle(frame, obj.pt_a(), obj.pt_b(), Color.GREEN.value, 4)
            write_text(frame, f"{obj.serial}", (obj.cx(), obj.cy()))
            detected_objects_without_trackers.append(
                DetectedObject(obj.type, obj.probability, obj.x, obj.y, obj.w,
                               obj.h))

        for obj in detected_objects_without_trackers:
            object_bbox_check(frame, obj)
        output_pickler.pickle(detected_objects_without_trackers,
                              f"frame{frame_n:06d}")

        cv2.imwrite(f"output/frame{frame_n:06d}.jpg", frame)