geetu040/DepthPro

1from PIL import Image2import torch3 4# custom installation from this PR: https://github.com/huggingface/transformers/pull/345835# !pip install git+https://github.com/geetu040/transformers.git@depth-pro-projects#egg=transformers6from transformers import DepthProImageProcessorFast, DepthProForDepthEstimation7 8# initialize processor and model9checkpoint = "geetu040/DepthPro"10revision = "project"11image_processor = DepthProImageProcessorFast.from_pretrained(checkpoint, revision=revision)12model = DepthProForDepthEstimation.from_pretrained(checkpoint, revision=revision)13device = torch.device("cuda" if torch.cuda.is_available() else "cpu")14model = model.to(device)15 16def predict(image):17	# inference18 19	# prepare image for the model20	inputs = image_processor(images=image, return_tensors="pt")21	inputs = {k: v.to(device) for k, v in inputs.items()}22 23	with torch.no_grad():24		outputs = model(**inputs)25 26	# interpolate to original size27	post_processed_output = image_processor.post_process_depth_estimation(28		outputs, target_sizes=[(image.height, image.width)],29	)30 31	# visualize the prediction32	depth = post_processed_output[0]["predicted_depth"]33	depth = (depth - depth.min()) / depth.max()34	depth = depth * 255.35	depth = depth.detach().cpu().numpy()36	depth = Image.fromarray(depth.astype("uint8"))37 38	return depth39