Update app.py

app.py

@@ -6,24 +6,32 @@ import gradio as gr
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model = DPTForDepthEstimation.from_pretrained("./", local_files_only=True, torch_dtype=torch.float16).to(device)
+model = torch.jit.script(model)
 processor = DPTImageProcessor.from_pretrained("Intel/dpt-swinv2-tiny-256")
 
 color_map = cv2.applyColorMap(np.arange(256, dtype=np.uint8), cv2.COLORMAP_INFERNO)
 
+input_tensor = torch.zeros((1, 3, 128, 128), dtype=torch.float16, device=device)
+depth_map = np.zeros((128, 128), dtype=np.float32)
+depth_map_colored = np.zeros((128, 128, 3), dtype=np.uint8)
+
+def preprocess_image(image):
+    return cv2.resize(image, (128, 128), interpolation=cv2.INTER_AREA).transpose(2, 0, 1).astype(np.float32) / 255.0
+
 @torch.inference_mode()
 def process_frame(image):
-    rgb_frame = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
-    resized_frame = cv2.resize(rgb_frame, (128, 128), interpolation=cv2.INTER_AREA)
+    preprocessed = preprocess_image(image)
+    input_tensor[0] = torch.from_numpy(preprocessed).to(device)
 
-    inputs = processor(images=resized_frame, return_tensors="pt").to(device)
-    inputs = {k: v.to(torch.float16) for k, v in inputs.items()}
+    if torch.cuda.is_available():
+        torch.cuda.synchronize()
 
-    predicted_depth = model(**inputs).predicted_depth
-    depth_map = predicted_depth.squeeze().cpu().numpy()
-    
-    depth_map = (depth_map - depth_map.min()) / (depth_map.max() - depth_map.min())
-    depth_map = (depth_map * 255).astype(np.uint8)
-    depth_map_colored = cv2.applyColorMap(depth_map, color_map)
+    predicted_depth = model(input_tensor).predicted_depth
+    np.subtract(predicted_depth.squeeze().cpu().numpy(), predicted_depth.min().item(), out=depth_map)
+    np.divide(depth_map, depth_map.max(), out=depth_map)
+    np.multiply(depth_map, 255, out=depth_map)
+    depth_map = depth_map.astype(np.uint8)
+    cv2.applyColorMap(depth_map, color_map, dst=depth_map_colored)
     
     return depth_map_colored