Update app.py

app.py

@@ -1,41 +1,53 @@
 import cv2
 import torch
 import numpy as np
-from transformers import DPTForDepthEstimation, DPTImageProcessor
+from transformers import DPTImageProcessor
 import gradio as gr
-import torch.nn.utils.prune as prune
 import matplotlib.pyplot as plt
 from mpl_toolkits.mplot3d import Axes3D
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-model = DPTForDepthEstimation.from_pretrained("Intel/dpt-swinv2-tiny-256", torch_dtype=torch.float32)
-model.eval()
-
-# Apply global unstructured pruning
-parameters_to_prune = [
-    (module, "weight") for module in filter(lambda m: isinstance(m, (torch.nn.Conv2d, torch.nn.Linear)), model.modules())
-]
-prune.global_unstructured(
-    parameters_to_prune,
-    pruning_method=prune.L1Unstructured,
-    amount=0.2,  # Prune 20% of weights
-)
+# Load your custom trained model
+class CompressedStudentModel(nn.Module):
+    def __init__(self):
+        super(CompressedStudentModel, self).__init__()
+        self.encoder = nn.Sequential(
+            nn.Conv2d(3, 64, kernel_size=3, padding=1),
+            nn.ReLU(),
+            nn.Conv2d(64, 64, kernel_size=3, padding=1),
+            nn.ReLU(),
+            nn.MaxPool2d(2),
+            nn.Conv2d(64, 128, kernel_size=3, padding=1),
+            nn.ReLU(),
+            nn.Conv2d(128, 128, kernel_size=3, padding=1),
+            nn.ReLU(),
+            nn.MaxPool2d(2),
+            nn.Conv2d(128, 256, kernel_size=3, padding=1),
+            nn.ReLU(),
+            nn.Conv2d(256, 256, kernel_size=3, padding=1),
+            nn.ReLU(),
+        )
+        self.decoder = nn.Sequential(
+            nn.ConvTranspose2d(256, 128, kernel_size=3, stride=2, padding=1, output_padding=1),
+            nn.ReLU(),
+            nn.ConvTranspose2d(128, 64, kernel_size=3, stride=2, padding=1, output_padding=1),
+            nn.ReLU(),
+            nn.Conv2d(64, 1, kernel_size=3, padding=1),
+        )
 
-for module, _ in parameters_to_prune:
-    prune.remove(module, "weight")
+    def forward(self, x):
+        features = self.encoder(x)
+        depth = self.decoder(features)
+        return depth
 
-model = torch.quantization.quantize_dynamic(
-    model, {torch.nn.Linear, torch.nn.Conv2d}, dtype=torch.qint8
-)
-
-model = model.to(device)
+# Initialize and load weights into the student model
+model = CompressedStudentModel().to(device)
+model.load_state_dict(torch.load("huntrezz_depth_v2.pt", map_location=device))
+model.eval()
 
 processor = DPTImageProcessor.from_pretrained("Intel/dpt-swinv2-tiny-256")
 
-color_map = cv2.applyColorMap(np.arange(256, dtype=np.uint8), cv2.COLORMAP_INFERNO)
-color_map = torch.from_numpy(color_map).to(device)
-
 def preprocess_image(image):
     image = cv2.resize(image, (128, 72))
     image = torch.from_numpy(image).permute(2, 0, 1).unsqueeze(0).float().to(device)
@@ -46,15 +58,13 @@ def plot_depth_map(depth_map, original_image):
     ax = fig.add_subplot(111, projection='3d')
     x, y = np.meshgrid(range(depth_map.shape[1]), range(depth_map.shape[0]))
     
-    # Resize original image to match depth map dimensions
     original_image_resized = cv2.resize(original_image, (depth_map.shape[1], depth_map.shape[0]))
     colors = original_image_resized.reshape(depth_map.shape[0], depth_map.shape[1], 3) / 255.0
     
     ax.plot_surface(x, y, depth_map, facecolors=colors, shade=False)
     ax.set_zlim(0, 1)
     
-    # Adjust the view to look down at an angle from a higher position
-    ax.view_init(elev=45, azim=180)  # 180-degree rotation and a higher angle
+    ax.view_init(elev=45, azim=180)
     plt.axis('off')
     plt.close(fig)
     
@@ -69,14 +79,11 @@ def process_frame(image):
     if image is None:
         return None
     preprocessed = preprocess_image(image)
-    predicted_depth = model(preprocessed).predicted_depth
-    depth_map = predicted_depth.squeeze().cpu().numpy()
+    predicted_depth = model(preprocessed).squeeze().cpu().numpy()
     
-    # Normalize depth map
-    depth_map = (depth_map - depth_map.min()) / (depth_map.max() - depth_map.min())
+    depth_map = (predicted_depth - predicted_depth.min()) / (predicted_depth.max() - predicted_depth.min())
     
-    # Convert BGR to RGB if necessary
-    if image.shape[2] == 3:  # Check if it's a color image
+    if image.shape[2] == 3:
         image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
     
     return plot_depth_map(depth_map, image)