Spaces:
Runtime error
Runtime error
| import torch | |
| import torch.nn.functional as F | |
| import numpy as np | |
| import cv2 | |
| from PIL import Image | |
| from config import SAPIENS_LITE_MODELS_PATH | |
| def load_model(task, version): | |
| try: | |
| model_path = SAPIENS_LITE_MODELS_PATH[task][version] | |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| if torch.cuda.is_available() and torch.cuda.get_device_properties(0).major >= 8: | |
| torch.backends.cuda.matmul.allow_tf32 = True | |
| torch.backends.cudnn.allow_tf32 = True | |
| model = torch.jit.load(model_path) | |
| model.eval() | |
| model.to(device) | |
| return model, device | |
| except KeyError as e: | |
| print(f"Error: Tarea o versión inválida. {e}") | |
| return None, None | |
| def preprocess_image(image, input_shape): | |
| img = cv2.resize(image, (input_shape[2], input_shape[1]), interpolation=cv2.INTER_LINEAR).transpose(2, 0, 1) | |
| img = torch.from_numpy(img) | |
| img = img[[2, 1, 0], ...].float() | |
| mean = torch.tensor([123.5, 116.5, 103.5]).view(-1, 1, 1) | |
| std = torch.tensor([58.5, 57.0, 57.5]).view(-1, 1, 1) | |
| img = (img - mean) / std | |
| return img.unsqueeze(0) | |
| def post_process_normal(result, original_shape): | |
| if result.dim() == 3: | |
| result = result.unsqueeze(0) | |
| elif result.dim() == 4: | |
| pass | |
| else: | |
| raise ValueError(f"Unexpected result dimension: {result.dim()}") | |
| seg_logits = F.interpolate(result, size=original_shape, mode="bilinear", align_corners=False).squeeze(0) | |
| normal_map = seg_logits.float().cpu().numpy().transpose(1, 2, 0) # H x W x 3 | |
| return normal_map | |
| def visualize_normal(normal_map): | |
| normal_map_norm = np.linalg.norm(normal_map, axis=-1, keepdims=True) | |
| normal_map_normalized = normal_map / (normal_map_norm + 1e-5) # Add a small epsilon to avoid division by zero | |
| normal_map_vis = ((normal_map_normalized + 1) / 2 * 255).astype(np.uint8) | |
| normal_map_vis = normal_map_vis[:, :, ::-1] # RGB to BGR | |
| return normal_map_vis | |
| def process_image_or_video(input_data, task='normal', version='sapiens_0.3b'): | |
| model, device = load_model(task, version) | |
| if model is None or device is None: | |
| return None | |
| input_shape = (3, 1024, 768) | |
| def process_frame(frame): | |
| if isinstance(frame, Image.Image): | |
| frame = np.array(frame) | |
| if frame.shape[2] == 4: # RGBA | |
| frame = cv2.cvtColor(frame, cv2.COLOR_RGBA2RGB) | |
| img = preprocess_image(frame, input_shape) | |
| with torch.no_grad(): | |
| result = model(img.to(device)) | |
| normal_map = post_process_normal(result, (frame.shape[0], frame.shape[1])) | |
| normal_image = visualize_normal(normal_map) | |
| return Image.fromarray(cv2.cvtColor(normal_image, cv2.COLOR_BGR2RGB)) | |
| if isinstance(input_data, np.ndarray): # Video frame | |
| return process_frame(input_data) | |
| elif isinstance(input_data, Image.Image): # Imagen | |
| return process_frame(input_data) | |
| else: | |
| print("Tipo de entrada no soportado. Por favor, proporcione una imagen PIL o un frame de video numpy.") | |
| return None |