File size: 10,489 Bytes
8166792
 
 
 
661e202
00ab2e7
8dda9f5
c5c62a8
8166792
 
 
661e202
8166792
 
 
 
f407854
 
 
 
 
 
 
 
 
 
8166792
c5c62a8
 
 
 
 
 
 
 
 
 
 
 
 
 
 
8166792
55bd8a1
 
f407854
55bd8a1
f407854
55bd8a1
 
 
 
 
 
 
 
 
 
 
661e202
c5c62a8
661e202
f407854
661e202
 
 
 
8166792
c5c62a8
8166792
f407854
 
 
 
 
 
 
 
 
 
 
 
 
c5c62a8
f407854
 
 
 
 
 
8166792
 
f407854
8166792
 
 
 
 
f407854
8166792
 
c5c62a8
661e202
8dda9f5
 
661e202
 
 
 
 
 
 
 
 
 
 
 
8166792
8dda9f5
8166792
 
 
c5c62a8
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
8166792
 
00ab2e7
 
 
8166792
 
 
 
 
 
661e202
 
 
 
 
 
8166792
 
 
 
 
 
 
 
661e202
 
 
00ab2e7
8166792
 
00ab2e7
 
 
 
 
8166792
00ab2e7
8166792
 
 
 
 
 
 
 
661e202
 
 
 
 
00ab2e7
661e202
8166792
 
 
 
 
 
 
 
00ab2e7
 
661e202
 
 
8166792
 
00ab2e7
 
 
 
8166792
 
 
 
00ab2e7
8166792
 
661e202
8166792
 
661e202
00ab2e7
8166792
 
00ab2e7
8166792
f407854
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
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
import cv2
import gradio as gr
import numpy as np
import os
import utils
import plotly.graph_objects as go
import spaces
import torch

from image_segmenter import ImageSegmenter
from monocular_depth_estimator import MonocularDepthEstimator
from point_cloud_generator import display_pcd

# params
CANCEL_PROCESSING = False

# Initialize classes without loading models
img_seg = None
depth_estimator = None

def initialize_models():
    global img_seg, depth_estimator
    if img_seg is None:
        img_seg = ImageSegmenter(model_type="yolov8s-seg")
    if depth_estimator is None:
        depth_estimator = MonocularDepthEstimator(model_type="midas_v21_small_256")

def safe_gpu_decorator(func):
    """Custom decorator to handle GPU operations safely"""
    def wrapper(*args, **kwargs):
        try:
            return func(*args, **kwargs)
        except RuntimeError as e:
            if "cudaGetDeviceCount" in str(e):
                print("GPU initialization failed, falling back to CPU")
                # Set environment variable to force CPU
                os.environ['CUDA_VISIBLE_DEVICES'] = ''
                return func(*args, **kwargs)
            raise
    return wrapper

@safe_gpu_decorator
def process_image(image):
    try:
        print("Starting image processing")
        initialize_models()
        
        image = utils.resize(image)
        image_segmentation, objects_data = img_seg.predict(image)
        depthmap, depth_colormap = depth_estimator.make_prediction(image)
        dist_image = utils.draw_depth_info(image, depthmap, objects_data)
        objs_pcd = utils.generate_obj_pcd(depthmap, objects_data)
        plot_fig = display_pcd(objs_pcd)
        return image_segmentation, depth_colormap, dist_image, plot_fig
    except Exception as e:
        print(f"Error in process_image: {str(e)}")
        import traceback
        print(traceback.format_exc())
        raise

@safe_gpu_decorator
def test_process_img(image):
    initialize_models()
    image = utils.resize(image)
    image_segmentation, objects_data = img_seg.predict(image)
    depthmap, depth_colormap = depth_estimator.make_prediction(image)
    return image_segmentation, objects_data, depthmap, depth_colormap

@safe_gpu_decorator
def process_video(vid_path=None):
    try:
        initialize_models()
        vid_cap = cv2.VideoCapture(vid_path)
        while vid_cap.isOpened():
            ret, frame = vid_cap.read()        
            if ret:
                print("making predictions ....")
                frame = utils.resize(frame)
                image_segmentation, objects_data = img_seg.predict(frame)
                depthmap, depth_colormap = depth_estimator.make_prediction(frame)
                dist_image = utils.draw_depth_info(frame, depthmap, objects_data)
                yield cv2.cvtColor(image_segmentation, cv2.COLOR_BGR2RGB), depth_colormap, cv2.cvtColor(dist_image, cv2.COLOR_BGR2RGB)
        
        vid_cap.release()
        return None
    except Exception as e:
        print(f"Error in process_video: {str(e)}")
        import traceback
        print(traceback.format_exc())
        raise

def update_segmentation_options(options):
    initialize_models()
    img_seg.is_show_bounding_boxes = True if 'Show Boundary Box' in options else False
    img_seg.is_show_segmentation = True if 'Show Segmentation Region' in options else False
    img_seg.is_show_segmentation_boundary = True if 'Show Segmentation Boundary' in options else False

def update_confidence_threshold(thres_val):
    initialize_models()
    img_seg.confidence_threshold = thres_val/100

@safe_gpu_decorator
def model_selector(model_type):
    global img_seg, depth_estimator
    
    if "Small - Better performance and less accuracy" == model_type:
        midas_model, yolo_model = "midas_v21_small_256", "yolov8s-seg"
    elif "Medium - Balanced performance and accuracy" == model_type:
        midas_model, yolo_model = "dpt_hybrid_384", "yolov8m-seg"
    elif "Large - Slow performance and high accuracy" == model_type:
        midas_model, yolo_model = "dpt_large_384", "yolov8l-seg"
    else:
        midas_model, yolo_model = "midas_v21_small_256", "yolov8s-seg"

    img_seg = ImageSegmenter(model_type=yolo_model)
    depth_estimator = MonocularDepthEstimator(model_type=midas_model)

def cancel():
    global CANCEL_PROCESSING
    CANCEL_PROCESSING = True

if __name__ == "__main__":
    # Try to initialize CUDA early to catch any issues
    try:
        if torch.cuda.is_available():
            print("CUDA is available. Using GPU.")
            # Test CUDA initialization
            torch.cuda.init()
            device = torch.device("cuda")
        else:
            print("CUDA is not available. Using CPU.")
            os.environ['CUDA_VISIBLE_DEVICES'] = ''
            device = torch.device("cpu")
    except RuntimeError as e:
        print(f"CUDA initialization failed: {e}")
        print("Falling back to CPU mode")
        os.environ['CUDA_VISIBLE_DEVICES'] = ''
        device = torch.device("cpu")

    with gr.Blocks() as my_app:
        # title
        gr.Markdown("<h1><center>Simultaneous Segmentation and Depth Estimation</center></h1>")
        gr.Markdown("<h3><center>Created by Vaishanth</center></h3>")
        gr.Markdown("<h3><center>This model estimates the depth of segmented objects.</center></h3>")

        # tabs
        with gr.Tab("Image"):
            with gr.Row():
                with gr.Column(scale=1):
                    img_input = gr.Image()
                    model_type_img = gr.Dropdown(
                        ["Small - Better performance and less accuracy", 
                         "Medium - Balanced performance and accuracy", 
                         "Large - Slow performance and high accuracy"], 
                        label="Model Type", value="Small - Better performance and less accuracy",
                        info="Select the inference model before running predictions!")
                    options_checkbox_img = gr.CheckboxGroup(["Show Boundary Box", "Show Segmentation Region", "Show Segmentation Boundary"], label="Options")
                    conf_thres_img = gr.Slider(1, 100, value=60, label="Confidence Threshold", info="Choose the threshold above which objects should be detected")
                    submit_btn_img = gr.Button(value="Predict")                    

                with gr.Column(scale=2):
                    with gr.Row():
                        segmentation_img_output = gr.Image(height=300, label="Segmentation")
                        depth_img_output = gr.Image(height=300, label="Depth Estimation")
                    
                    with gr.Row():
                        dist_img_output = gr.Image(height=300, label="Distance")
                        pcd_img_output = gr.Plot(label="Point Cloud")
            
            gr.Examples(
                examples=[os.path.join(os.path.dirname(__file__), "assets/images/baggage_claim.jpg"),
                          os.path.join(os.path.dirname(__file__), "assets/images/kitchen_2.png"),
                          os.path.join(os.path.dirname(__file__), "assets/images/soccer.jpg"),
                          os.path.join(os.path.dirname(__file__), "assets/images/room_2.png"),
                          os.path.join(os.path.dirname(__file__), "assets/images/living_room.jpg")],
                inputs=img_input,
                outputs=[segmentation_img_output, depth_img_output, dist_img_output, pcd_img_output],
                fn=process_image,
                cache_examples=True,
            )

        with gr.Tab("Video"):
            with gr.Row():
                with gr.Column(scale=1):
                    vid_input = gr.Video()
                    model_type_vid = gr.Dropdown(
                        ["Small - Better performance and less accuracy", 
                         "Medium - Balanced performance and accuracy", 
                         "Large - Slow performance and high accuracy"], 
                        label="Model Type", value="Small - Better performance and less accuracy",
                        info="Select the inference model before running predictions!")
                    
                    options_checkbox_vid = gr.CheckboxGroup(["Show Boundary Box", "Show Segmentation Region", "Show Segmentation Boundary"], label="Options")
                    conf_thres_vid = gr.Slider(1, 100, value=60, label="Confidence Threshold", info="Choose the threshold above which objects should be detected")
                    with gr.Row():
                        cancel_btn = gr.Button(value="Cancel")
                        submit_btn_vid = gr.Button(value="Predict")
            
                with gr.Column(scale=2):
                    with gr.Row():
                        segmentation_vid_output = gr.Image(height=300, label="Segmentation")
                        depth_vid_output = gr.Image(height=300, label="Depth Estimation")
                    
                    with gr.Row():
                        dist_vid_output = gr.Image(height=300, label="Distance")
            
            gr.Examples(
                examples=[os.path.join(os.path.dirname(__file__), "assets/videos/input_video.mp4"),
                          os.path.join(os.path.dirname(__file__), "assets/videos/driving.mp4"),
                          os.path.join(os.path.dirname(__file__), "assets/videos/overpass.mp4"),
                          os.path.join(os.path.dirname(__file__), "assets/videos/walking.mp4")],
                inputs=vid_input,
            )

        # image tab logic
        submit_btn_img.click(process_image, inputs=img_input, outputs=[segmentation_img_output, depth_img_output, dist_img_output, pcd_img_output])
        options_checkbox_img.change(update_segmentation_options, options_checkbox_img, [])
        conf_thres_img.change(update_confidence_threshold, conf_thres_img, [])
        model_type_img.change(model_selector, model_type_img, [])

        # video tab logic
        submit_btn_vid.click(process_video, inputs=vid_input, outputs=[segmentation_vid_output, depth_vid_output, dist_vid_output])
        model_type_vid.change(model_selector, model_type_vid, [])
        cancel_btn.click(cancel, inputs=[], outputs=[])
        options_checkbox_vid.change(update_segmentation_options, options_checkbox_vid, [])
        conf_thres_vid.change(update_confidence_threshold, conf_thres_vid, [])       

    my_app.queue(max_size=10).launch()