Update app.py

app.py

@@ -1,9 +1,7 @@
-from ultralytics import YOLO
 import cv2
 import gradio as gr
 import numpy as np
 import os
-import torch
 import utils
 import plotly.graph_objects as go
 import spaces
@@ -15,24 +13,26 @@ from point_cloud_generator import display_pcd
 # params
 CANCEL_PROCESSING = False
 
-# Initialize models (but don't load them yet)
-img_seg = ImageSegmenter(model_type="yolov8s-seg")
-depth_estimator = MonocularDepthEstimator(model_type="midas_v21_small_256")
+# 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")
 
 @spaces.GPU
 def process_image(image):
     try:
         print("Starting image processing")
-        image = utils.resize(image)
-        print("Image resized")
+        initialize_models()
         
-        # Models will be loaded here when needed
+        image = utils.resize(image)
         image_segmentation, objects_data = img_seg.predict(image)
-        print("Segmentation complete")
-        
         depthmap, depth_colormap = depth_estimator.make_prediction(image)
-        print("Depth estimation complete")
-        
         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)
@@ -43,37 +43,47 @@ def process_image(image):
         print(traceback.format_exc())
         raise
 
-@spaces.GPU(duration=30)
+@spaces.GPU
 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
 
-@spaces.GPU(duration=60)  # Longer duration for video processing
+@spaces.GPU
 def process_video(vid_path=None):
-    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)
-    
-    return 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)
+        
+        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
 
-@spaces.GPU(duration=10)  # Short duration for model loading
+@spaces.GPU
 def model_selector(model_type):
     global img_seg, depth_estimator
     
@@ -94,7 +104,6 @@ def cancel():
     CANCEL_PROCESSING = True
 
 if __name__ == "__main__":
-    # gradio gui app
     with gr.Blocks() as my_app:
         # title
         gr.Markdown("<h1><center>Simultaneous Segmentation and Depth Estimation</center></h1>")
@@ -185,5 +194,4 @@ if __name__ == "__main__":
         options_checkbox_vid.change(update_segmentation_options, options_checkbox_vid, [])
         conf_thres_vid.change(update_confidence_threshold, conf_thres_vid, [])       
 
-    # Launch with appropriate queue settings for ZeroGPU
-    my_app.queue(max_size=20).launch()
+    my_app.queue(max_size=10).launch()