Update app.py

app.py

@@ -12,6 +12,22 @@ model_configs = {
     'vitl': {'encoder': 'vitl', 'features': 256, 'out_channels': [256, 512, 1024, 1024]}
 }
 
+def get_image_intensity(img, gamma_correction=1.0):
+    """
+    Extract intensity map from an image using HSV color space
+    """
+    # Convert to HSV color space
+    result = cv2.cvtColor(img, cv2.COLOR_RGB2HSV)
+    # Extract Value channel (intensity)
+    result = result[:, :, 2].astype(np.float32) / 255.0
+    # Apply gamma correction
+    result = result ** gamma_correction
+    # Convert back to 0-255 range
+    result = (result * 255.0).clip(0, 255).astype(np.uint8)
+    # Convert to RGB (still grayscale but in RGB format)
+    result = cv2.cvtColor(result, cv2.COLOR_GRAY2RGB)
+    return result
+
 # Initialize model globally
 def initialize_model():
     encoder = 'vitl'
@@ -42,21 +58,21 @@ MODEL = initialize_model()
 @spaces.GPU
 def process_image(input_image):
     """
-    Process the input image and return depth maps
+    Process the input image and return depth maps and intensity map
     """
     if input_image is None:
-        return None, None
+        return None, None, None
     
     # Move model to GPU for processing
     MODEL.to('cuda')
     MODEL.eval()
     
-    # Convert from RGB to BGR
-    input_image = cv2.cvtColor(np.array(input_image), cv2.COLOR_RGB2BGR)
+    # Convert from RGB to BGR for depth processing
+    input_bgr = cv2.cvtColor(np.array(input_image), cv2.COLOR_RGB2BGR)
     
     with torch.no_grad():
         # Get depth map
-        depth = MODEL.infer_image(input_image)
+        depth = MODEL.infer_image(input_bgr)
         
         # Normalize depth for visualization (0-255)
         depth_normalized = ((depth - depth.min()) / (depth.max() - depth.min()) * 255).astype(np.uint8)
@@ -68,16 +84,19 @@ def process_image(input_image):
     # Move model back to CPU after processing
     MODEL.to('cpu')
     
-    return depth_normalized, depth_colormap
+    # Get intensity map
+    intensity_map = get_image_intensity(np.array(input_image), gamma_correction=1.0)
+    
+    return depth_normalized, depth_colormap, intensity_map
 
 @spaces.GPU
 def gradio_interface(input_img):
     try:
-        depth_raw, depth_colored = process_image(input_img)
-        return [input_img, depth_raw, depth_colored]
+        depth_raw, depth_colored, intensity = process_image(input_img)
+        return [input_img, depth_raw, depth_colored, intensity]
     except Exception as e:
         print(f"Error processing image: {str(e)}")
-        return [input_img, None, None]
+        return [input_img, None, None, None]
 
 # Define interface
 iface = gr.Interface(
@@ -86,10 +105,11 @@ iface = gr.Interface(
     outputs=[
         gr.Image(label="Original Image"),
         gr.Image(label="Raw Depth Map"),
-        gr.Image(label="Colored Depth Map")
+        gr.Image(label="Colored Depth Map"),
+        gr.Image(label="Intensity Map")
     ],
-    title="Depth Estimation",
-    description="Upload an image to generate its depth map.",
+    title="Depth and Intensity Estimation",
+    description="Upload an image to generate its depth map and intensity map.",
     examples=["image.jpg"]
 )