upload app

.gitattributes

@@ -33,3 +33,6 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+*.jpg filter=lfs diff=lfs merge=lfs -text
+*.jpeg filter=lfs diff=lfs merge=lfs -text
+*.png filter=lfs diff=lfs merge=lfs -text

app.py

@@ -0,0 +1,40 @@
+import gradio as gr
+from model import predict
+
+color_maps = [
+    'viridis', 'plasma', 'inferno', 'magma', 'cividis',
+    'Greys', 'Purples', 'Blues', 'Greens', 'Oranges', 'Reds',
+    'YlOrBr', 'YlOrRd', 'OrRd', 'PuRd', 'RdPu', 'BuPu',
+    'GnBu', 'PuBu', 'YlGnBu', 'PuBuGn', 'BuGn', 'YlGn',
+    'binary', 'gist_yarg', 'gist_gray', 'gray', 'bone',
+    'pink', 'spring', 'summer', 'autumn', 'winter', 'cool',
+    'Wistia', 'hot', 'afmhot', 'gist_heat', 'copper',
+    'PiYG', 'PRGn', 'BrBG', 'PuOr', 'RdGy', 'RdBu', 'RdYlBu',
+    'RdYlGn', 'Spectral', 'coolwarm', 'bwr', 'seismic',
+    'twilight', 'twilight_shifted', 'hsv',
+    'Pastel1', 'Pastel2', 'Paired', 'Accent', 'Dark2',
+    'Set1', 'Set2', 'Set3', 'tab10', 'tab20', 'tab20b', 'tab20c',
+    'flag', 'prism', 'ocean', 'gist_earth', 'terrain',
+    'gist_stern', 'gnuplot', 'gnuplot2', 'CMRmap',
+    'cubehelix', 'brg', 'gist_rainbow', 'rainbow', 'jet',
+    'turbo', 'nipy_spectral', 'gist_ncar',
+]
+examples = [
+    ["assets/examples/myself.jpeg", "afmhot"],
+    ["assets/examples/myself.jpeg", "inferno"],
+]
+
+interface = gr.Interface(
+    fn=predict,
+    inputs=[
+        gr.Image(type="pil"),
+        gr.Dropdown(choices=color_maps),
+    ],
+    outputs=gr.Image(type="pil"),
+    title="DepthPro: Colorify",
+    description="Applies segmentation on the input image, then creates the depth map and finally colorizes it.",
+    examples=examples,
+)
+
+if __name__ == "__main__":
+    interface.launch()

model/__init__.py

@@ -0,0 +1,28 @@
+from PIL import Image
+import numpy as np
+import matplotlib.pyplot as plt
+from segmentation import predict as segmentation_predict
+from depth_estimation import predict as depth_estimation_predict
+
+def predict(image, color_map):
+	# inference
+
+	mask_image = segmentation_predict(image)
+
+	segmented_image = Image.composite(
+		image,
+		Image.new("RGB", image.size, (0, 0, 0)),
+		mask_image.convert("L")
+	)
+
+	depth_image = depth_estimation_predict(segmented_image)
+
+	# apply matplotlib colormap (e.g., viridis)
+	depth_array = np.array(depth_image)  # Convert PIL image to NumPy array
+	colormap = plt.get_cmap(color_map)   # Choose a colormap
+	depth_colored = colormap(depth_array[:, :, 0] / 255.0)  # Normalize and apply colormap
+	depth_colored = (depth_colored * 255).astype(np.uint8)  # Convert to RGB (discard alpha)
+
+	depth_colored = Image.fromarray(depth_colored)
+
+	return depth_colored

model/depth_estimation.py

@@ -0,0 +1,38 @@
+from PIL import Image
+import torch
+
+# custom installation from this PR: https://github.com/huggingface/transformers/pull/34583
+# !pip install git+https://github.com/geetu040/transformers.git@depth-pro-projects#egg=transformers
+from transformers import DepthProImageProcessorFast, DepthProForDepthEstimation
+
+# initialize processor and model
+checkpoint = "geetu040/DepthPro"
+revision = "project"
+image_processor = DepthProImageProcessorFast.from_pretrained(checkpoint, revision=revision)
+model = DepthProForDepthEstimation.from_pretrained(checkpoint, revision=revision)
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model = model.to(device)
+
+def predict(image):
+	# inference
+
+	# prepare image for the model
+	inputs = image_processor(images=image, return_tensors="pt")
+	inputs = {k: v.to(device) for k, v in inputs.items()}
+
+	with torch.no_grad():
+		outputs = model(**inputs)
+
+	# interpolate to original size
+	post_processed_output = image_processor.post_process_depth_estimation(
+		outputs, target_sizes=[(image.height, image.width)],
+	)
+
+	# visualize the prediction
+	depth = post_processed_output[0]["predicted_depth"]
+	depth = (depth - depth.min()) / depth.max()
+	depth = depth * 255.
+	depth = depth.detach().cpu().numpy()
+	depth = Image.fromarray(depth.astype("uint8"))
+
+	return depth

model/segmentation.py

@@ -0,0 +1,65 @@
+from PIL import Image
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from huggingface_hub import hf_hub_download
+
+# custom installation from this PR: https://github.com/huggingface/transformers/pull/34583
+# !pip install git+https://github.com/geetu040/transformers.git@depth-pro-projects#egg=transformers
+from transformers import DepthProConfig, DepthProImageProcessorFast, DepthProForDepthEstimation
+
+# initialize model
+config = DepthProConfig(use_fov_model=False)
+model = DepthProForDepthEstimation(config)
+features = config.fusion_hidden_size
+semantic_classifier_dropout = 0.1
+num_labels = 1
+model.head.head = nn.Sequential(
+    nn.Conv2d(features, features, kernel_size=3, padding=1, bias=False),
+    nn.BatchNorm2d(features),
+    nn.ReLU(),
+    nn.Dropout(semantic_classifier_dropout),
+    nn.Conv2d(features, features, kernel_size=1),
+    nn.ConvTranspose2d(features, num_labels, kernel_size=2, stride=2, padding=0, bias=True),
+)
+
+# load weights
+weights_path = hf_hub_download(repo_id="geetu040/DepthPro_Segmentation_Human", filename="model_weights.pth")
+model.load_state_dict(torch.load(weights_path, map_location=torch.device('cpu'), weights_only=True))
+
+# load to device
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model = model.to(device)
+
+# load image processor
+image_processor = DepthProImageProcessorFast()
+
+def predict(image):
+	# inference
+
+	image = image.convert("RGB")
+
+	# prepare image for the model
+	inputs = image_processor(images=image, return_tensors="pt")
+	inputs = {k: v.to(device) for k, v in inputs.items()}
+
+	# inference
+	with torch.no_grad():
+		output = model(**inputs)
+
+	# convert tensors to PIL.Image
+	output = output[0]                 # get output logits
+	output = F.interpolate(
+		output.unsqueeze(0),
+		size=(image.height, image.width)
+	)                                  # interpolate to match size
+	output = output.squeeze()          # get first and only batch and channel
+	output = output.sigmoid()          # apply sigmoid for binary segmentation
+	output = (output > 0.5).float()    # threshold to create binary mask
+	output = output.cpu()              # unload from cuda if used
+	output = output * 255              # convert [0, 1] to [0, 255]
+	output = output.numpy()            # convert to numpy
+	output = output.astype('uint8')    # convert to PIL.Image compatible format
+	output = Image.fromarray(output)   # create PIL.Image object
+
+	return output

requirements.txt

@@ -0,0 +1,7 @@
+gradio
+numpy
+matplotlib
+pillow
+torch
+torchvision
+git+https://github.com/geetu040/transformers.git@depth-pro-projects#egg=transformers