ARCube/detect.py at main · X-D-R/ARCube · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
import argparse
import os.path
import numpy as np
import cv2 as cv
from src.detection.detection import Detector
from src.tracking.frame import track_frame
from src.utils.draw_functions import draw_contours_of_rectangle, visualize_matches_on_photo
from rendering_CV import rendering_video, generate_obj_file

MAIN_DIR = os.path.dirname(os.path.abspath("detect.py"))


def parse_args_and_execute():
    '''Parse command-line arguments and execute the appropriate function (detect).'''

    parser = argparse.ArgumentParser(description="Detection")

    parser.add_argument('--model', type=str, help="Path to the saved model file")
    parser.add_argument('--demo', action='store_true', help="Use the default object model from the repository")
    parser.add_argument('--camera_params', type=str, required=False, help="Path to camera parameters file")
    parser.add_argument('--input', type=str, required=False, help="Path to input image or video for detection")
    parser.add_argument('--video', action='store_true', help="if you want to detect video,"
                                                             "don't use if you want to detect photo")
    parser.add_argument('--output', type=str, required=True, help="Path to output image or video after detection")
    parser.add_argument('--use_tracker', action='store_true', help="Use if you want to use tracking")
    parser.add_argument('--save_video', action='store_true', help="Use if you want to save video")
    parser.add_argument('--web_camera', action='store_true', help="Use if you want to use your camera")
    parser.add_argument('--visualize_matches', action='store_true', help="Use if you want to visualize matches with"
                                                                         " the reference image")
    parser.add_argument('--render_cv', action='store_true', help="Use if you want to use rendering by CV,"
                                                                 " only with video")
    parser.add_argument('--pyrender', type=str, help="Path to the the obj file for rendering, use if you want to use "
                                                     "rendering by pyrender. only with video")

    args = parser.parse_args()

    if args.pyrender is not None and args.render_cv:
        raise ValueError("Use only one render or None")

    camera_params_approximate = {}
    if args.camera_params is None:
        imgSize = (0, 0)
        if args.video and not args.web_camera:
            cap = cv.VideoCapture(args.input)
            ret, frame = cap.read()
            imgSize = frame.shape
            cap.release()
        elif args.web_camera:
            cap = cv.VideoCapture(0)
            ret, frame = cap.read()
            imgSize = frame.shape
            cap.release()
        else:
            img = cv.imread(args.input)
            imgSize = img.shape
        h, w = imgSize[0], imgSize[1]
        f = 0.9 * max(w, h)
        camera_params_approximate = {'mtx': np.array([[f, 0, w / 2], [0, f, h / 2], [0, 0, 1]], np.float32),
                                     'dist': np.array([0, 0, 0, 0, 0], np.float32)}

    if args.demo:
        detector = set_detector(os.path.join(MAIN_DIR, "ExampleFiles", "ModelParams", "model_test.npz"),
                                args.camera_params, camera_params_approximate=camera_params_approximate, use_flann=True)
    else:
        detector = set_detector(args.model, args.camera_params, camera_params_approximate=camera_params_approximate,
                                use_flann=True)

    if args.video:
        print('detecting object on video')
        track_length = 50 if args.use_tracker else 1
        if args.render_cv:
            coord_3d = detector.registration_params['object_corners_3d']
            w = coord_3d[2][0]
            h = coord_3d[2][1]
            generate_obj_file(os.path.join(MAIN_DIR, "ExampleFiles", "3d_models", "box_CV.npz"), w, h)
            print("OBJ file is generated")

        if args.visualize_matches:
            print('visualizing matches with the reference image')
            track_frame(detector, args.input, args.output, visualizing_matches=True, use_tracker=args.use_tracker,
                        use_web_camera=args.web_camera, save_video=args.save_video, render_cv=args.render_cv,
                        pyrender=args.pyrender)
        else:
            track_frame(detector, args.input, args.output, use_tracker=args.use_tracker, use_web_camera=args.web_camera,
                        save_video=args.save_video, render_cv=args.render_cv, pyrender=args.pyrender)

    else:
        print('detecting object on photo')
        img_points, src_pts, dst_pts, keypoints, matches = detector.detect_path(args.input)
        draw_contours_of_rectangle(args.input, args.output, img_points)

        if args.visualize_matches:
            print('visualizing matches with the reference image')
            visualize_matches_on_photo(detector.registration_params["img"],
                                       detector.registration_params["key_points_2d"],
                                       args.input, keypoints, matches)


def set_detector(model_params_file: str, camera_params_file: str, use_flann: bool = True,
                 camera_params_approximate: dict = {}) -> Detector:
    '''
    This function set detector using model params
    :param model_params_file: str, path to where model should be saved
    :param camera_params_file: str, path to saved Camera parameters
    :param use_flann: bool, use FLANN-based matching or not
    :return: Detector, object of Detector class
    '''

    detector = Detector()
    # detector.set_detector_by_model("CameraParams/CameraParams.npz", model, True)
    detector.set_detector(camera_params_file, model_params_file, use_flann, camera_params_approximate)

    return detector


def detect_photo(detector: Detector, input_file: str, output_file: str):
    '''
    This function detects object on the photo.
    :param detector: Detector, object of Detector class
    :param input_file: str, path to image where we use detector
    :param output_file: str, path to store the image with detected photo
    :return:
    '''
    img_points, src_pts, dst_pts, keypoints, matches = detector.detect_path(input_file)
    draw_contours_of_rectangle(input_file, output_file, img_points)


if __name__ == "__main__":
    # example of the detection
    # detector = set_detector(os.path.join(MAIN_DIR, "ExampleFiles", "ModelParams", "model_test.npz"),
    #                         os.path.join(MAIN_DIR, "ExampleFiles", "CameraParams", "CameraParams.npz"), True)
    # # Photo detection
    # detect_photo(detector, os.path.join(MAIN_DIR, "ExampleFiles", "examples", "images", "new_book_check.png"),
    #              os.path.join(MAIN_DIR, "ExampleFiles", "OutputFiles", "OutputImages", "contours_drawn.png"))
    #
    # # Video detection
    # track_frame(detector, os.path.join(MAIN_DIR, "ExampleFiles", "new_book_check", "new_book_video_main.mp4"),
    #             os.path.join(MAIN_DIR,
    #                          "ExampleFiles", "OutputFiles", "OutputVideos", "new_book_video_main_result_new_color.mp4"), 60,
    #             30, (0, 0, 255))

    # or
    parse_args_and_execute()
    '''
    python detect.py --model "ExampleFiles/ModelParams/model_test.npz" --camera_params "ExampleFiles/CameraParams/CameraParams.npz" --input "ExampleFiles/new_book_check/new_book_video_main.mp4" --video --output "ExampleFiles/OutputFiles/OutputVideos/new_book_video_main_result_new_color.mp4"

    python detect.py --demo --camera_params "ExampleFiles/CameraParams/CameraParams.npz" --input "ExampleFiles/new_book_check/new_book_video_main.mp4" --video --output "ExampleFiles/OutputFiles/OutputVideos/new_book_video_main_result_new_color.mp4" --use_tracker

    python detect.py --model "ExampleFiles/ModelParams/model_varior_book.npz" --video --output "ExampleFiles/OutputFiles/OutputVideos/varior_book_result.mp4" --use_tracker --web_camera

    python detect.py --model "ExampleFiles/ModelParams/model_varior_book_iphone.npz" --video --output "ExampleFiles/OutputFiles/OutputVideos/varior_book_result_with_pyrender_without_tracker.mp4" --use_tracker --web_camera --pyrender "ExampleFiles/3d_models/colored_box_varior.obj" --save_video

    python detect.py --model "ExampleFiles/ModelParams/model_varior_book_iphone.npz" --video --output "ExampleFiles/OutputFiles/OutputVideos/varior_book_result_with_render_cv_without_tracker.mp4" --web_camera --render_cv --save_video


    '''