Update app.py

app.py

@@ -6,6 +6,7 @@ import os
 import torch
 import utils
 import plotly.graph_objects as go
+import spaces
 
 from image_segmenter import ImageSegmenter
 from monocular_depth_estimator import MonocularDepthEstimator
@@ -14,9 +15,11 @@ from point_cloud_generator import display_pcd
 # params
 CANCEL_PROCESSING = False
 
+# Initialize models (but actual loading happens in decorated functions)
 img_seg = ImageSegmenter(model_type="yolov8s-seg")
 depth_estimator = MonocularDepthEstimator(model_type="midas_v21_small_256")
 
+@spaces.GPU(duration=30)  # Adjust duration based on your needs
 def process_image(image):
     image = utils.resize(image)
     image_segmentation, objects_data = img_seg.predict(image)
@@ -26,12 +29,14 @@ def process_image(image):
     plot_fig = display_pcd(objs_pcd)
     return image_segmentation, depth_colormap, dist_image, plot_fig
 
+@spaces.GPU(duration=30)
 def test_process_img(image):
     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
 
+@spaces.GPU(duration=60)  # Longer duration for video processing
 def process_video(vid_path=None):
     vid_cap = cv2.VideoCapture(vid_path)
     while vid_cap.isOpened():
@@ -54,8 +59,10 @@ def update_segmentation_options(options):
 def update_confidence_threshold(thres_val):
     img_seg.confidence_threshold = thres_val/100
 
+@spaces.GPU(duration=10)  # Short duration for model loading
 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:
@@ -69,30 +76,12 @@ def model_selector(model_type):
     depth_estimator = MonocularDepthEstimator(model_type=midas_model)
 
 def cancel():
+    global CANCEL_PROCESSING
     CANCEL_PROCESSING = True
 
 if __name__ == "__main__":
-
-    # testing
-    # img_1 = cv2.imread("assets/images/bus.jpg")
-    # img_1 = utils.resize(img_1)
-
-    # image_segmentation, objects_data, depthmap, depth_colormap = test_process_img(img_1)
-    # final_image = utils.draw_depth_info(image_segmentation, depthmap, objects_data)
-    # objs_pcd = utils.generate_obj_pcd(depthmap, objects_data)
-    # # print(objs_pcd[0][0])
-    # display_pcd(objs_pcd, use_matplotlib=True)
-
-    # cv2.imshow("Segmentation", image_segmentation)
-    # cv2.imshow("Depth", depthmap*objects_data[2][3])
-    # cv2.imshow("Final", final_image)
-
-    # cv2.waitKey(0)
-    # cv2.destroyAllWindows()
-
     # gradio gui app
     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>")
@@ -167,10 +156,7 @@ if __name__ == "__main__":
                           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,
-                # outputs=vid_output,
-                # fn=vid_segmenation,
             )
-            
 
         # 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])
@@ -185,5 +171,5 @@ if __name__ == "__main__":
         options_checkbox_vid.change(update_segmentation_options, options_checkbox_vid, [])
         conf_thres_vid.change(update_confidence_threshold, conf_thres_vid, [])       
 
-
-    my_app.queue(concurrency_count=5, max_size=20).launch()
+    # Launch with appropriate queue settings for ZeroGPU
+    my_app.queue(concurrency_count=1, max_size=10).launch()