mirror of
https://github.com/vladmandic/sdnext.git
synced 2026-01-27 15:02:48 +03:00
48 lines
1.8 KiB
Python
48 lines
1.8 KiB
Python
from PIL import Image
|
|
import numpy as np
|
|
import torch
|
|
from transformers import AutoImageProcessor, DPTForDepthEstimation
|
|
from modules import devices
|
|
from modules.shared import opts
|
|
|
|
|
|
image_processor: AutoImageProcessor = None
|
|
|
|
|
|
class DPTDetector:
|
|
def __init__(self, model=None, processor=None, model_path=None):
|
|
self.model = model
|
|
self.processor = processor
|
|
self.model_path = model_path or "Intel/dpt-large"
|
|
|
|
def __call__(self, input_image=None, model_path=None):
|
|
from modules.control.processors import cache_dir
|
|
if model_path is not None and model_path != self.model_path:
|
|
self.model_path = model_path
|
|
self.processor = None
|
|
self.model = None
|
|
if self.processor is None:
|
|
self.processor = AutoImageProcessor.from_pretrained(self.model_path, cache_dir=cache_dir)
|
|
if self.model is None:
|
|
self.model = DPTForDepthEstimation.from_pretrained(self.model_path, cache_dir=cache_dir)
|
|
|
|
self.model.to(devices.device)
|
|
with devices.inference_context():
|
|
inputs = self.processor(images=input_image, return_tensors="pt")
|
|
inputs.to(devices.device)
|
|
outputs = self.model(**inputs)
|
|
predicted_depth = outputs.predicted_depth
|
|
prediction = torch.nn.functional.interpolate(
|
|
predicted_depth.unsqueeze(1),
|
|
size=input_image.size[::-1],
|
|
mode="bicubic",
|
|
align_corners=False,
|
|
)
|
|
output = prediction.squeeze().cpu().numpy()
|
|
formatted = (output * 255 / np.max(output)).astype("uint8")
|
|
if opts.control_move_processor:
|
|
self.model.to('cpu')
|
|
depth = Image.fromarray(formatted)
|
|
depth = depth.convert('RGB')
|
|
return depth
|