''' This script is used to generate the semantic labels for the objects in the scene. ''' from utils.config import get_dataset, get_args from dataset.scannet import ScanNetDataset import os import numpy as np import open_clip import cv2 from PIL import Image import torch from torch.utils.data import TensorDataset, DataLoader from tqdm import tqdm from utils.config import update_args import argparse LEVELS = 3 def load_clip(device): print(f'[INFO] loading CLIP model...') model, _, preprocess = open_clip.create_model_and_transforms("ViT-H-14", pretrained="laion2b_s32b_b79k") model.to(device) model.eval() print(f'[INFO]', ' finish loading CLIP model...') return model, preprocess def box_multi_level(bbox, shape, level, expansion_ratio): left, top, right, bottom = bbox if level == 0: return left, top, right , bottom x_exp = int(abs(right - left)*expansion_ratio * level) y_exp = int(abs(bottom - top)*expansion_ratio * level) return max(0, left - x_exp), max(0, top - y_exp), min(shape[1], right + x_exp), min(shape[0], bottom + y_exp) def get_cropped_images(object, scene_id, preprocess, preloaded_images, num_images=5, expansion_ratio=0.1): croped_images = [] images = [] for frame in object['frames'][:num_images]: image = preloaded_images[frame['frame_id']] x1, y1, x2, y2 = frame['bbox'] x1, y1, x2, y2 = np.round([x1, y1, x2, y2]).astype(int) for level in range(LEVELS): x1_, y1_, x2_, y2_ = box_multi_level((x1, y1, x2, y2), np.asarray(image).shape, level, expansion_ratio) if x1_ == x2_ or y1_ == y2_: continue pil_image = pad_into_square(Image.fromarray(cv2.cvtColor(image[y1_:y2_, x1_:x2_], cv2.COLOR_BGR2RGB))) croped_images.append(preprocess(pil_image)) images.append(np.asarray(pil_image.resize((64, 64)))) if len(croped_images) == 0: return None, None return torch.stack(croped_images), np.concatenate(images, axis=1)[..., ::-1] def pad_into_square(image): width, height = image.size new_size = max(width, height) new_image = Image.new("RGB", (new_size, new_size), (255,255,255)) left = (new_size - width) // 2 top = (new_size - height) // 2 new_image.paste(image, (left, top)) return new_image def preload_images(scene, frame_ids): preloaded_images = {} for frame_id in frame_ids: img_path = f'/home/jovyan/users/lemeshko/mmdetection3d/data/scannet/posed_images/{scene}/{str(frame_id).zfill(5)}.jpg' image = cv2.imread(img_path) preloaded_images[frame_id] = image return preloaded_images def custom_probs(feature, label_text_features): object_feature = feature #np.mean(feature, axis=0, keepdims=True) print(object_feature.shape, np.mean(feature, axis=0, keepdims=True).shape) raw_similarity = object_feature @ label_text_features.T raw_similarity = np.sum(raw_similarity, axis=0, keepdims=True) exp_sim = np.exp(raw_similarity) prob = exp_sim / np.sum(exp_sim, axis=1, keepdims=True) probs = np.max(prob, axis=0) max_label_id = np.argmax(probs) prob = probs[max_label_id] return prob, max_label_id def main(args): scenes = np.loadtxt('/home/jovyan/users/bulat/workspace/3drec/Indoor/MaskClustering/splits/scannet200_subset.txt', dtype=str) # scenes = scenes[:1] # scenes = ["scene0011_00"] for scene in tqdm(scenes): frame_ids = set() args.seq_name = scene pred_dir = os.path.join('data/prediction', args.config, args.exp_name) if os.path.exists(f'{pred_dir}/{args.seq_name}.npz'): continue dataset = ScanNetDataset(scene, use_templates=True) total_point_num = dataset.get_scene_points().shape[0] label_features_dict = dataset.get_label_features() label_text_features = np.stack(list(label_features_dict.values())) descriptions = list(label_features_dict.keys()) scene_name = dataset.seq_name logs_path = f'logs/{args.exp_name}/{scene}' # Загружаем данные в новом формате (список объектов) object_list = np.load(f"/home/jovyan/users/bulat/workspace/3drec/det/OV/mask_proj/outputs/30/{scene}/mask_data.npy", allow_pickle=True) # print(f'[INFO] loaded {object_list} objects') label2id = dataset.get_label_id()[0] os.makedirs(pred_dir, exist_ok=True) num_instance = len(object_list) pred_dict = { "pred_masks": np.zeros((total_point_num, num_instance), dtype=bool), "pred_score": np.ones(num_instance), "pred_classes" : np.zeros(num_instance, dtype=np.int32) } # Собираем все frame_ids из всех объектов for object_data in object_list: frame_ids.update([frame['frame_id'] for frame in object_data['frames']]) frame_ids = list(frame_ids) preloaded_images = preload_images(scene, frame_ids) # Исправляем итерацию по объектам for idx, object_data in enumerate(object_list): croped_images, saved_images = get_cropped_images(object_data, scene_name, args.preprocess, preloaded_images=preloaded_images, num_images=args.num_images) if croped_images is None: print(f'[INFO] no croped images for object {idx}') continue bs = 32 chunks = torch.chunk(croped_images, max(1, len(croped_images) // bs)) features = [] for images in chunks: # images = images[0] images = images.to(args.device) with torch.no_grad(): image_features = args.model.encode_image(images).float() image_features /= image_features.norm(dim=-1, keepdim=True) image_features = image_features.cpu().numpy() for f in image_features: features.append(f) feature = np.stack(features) object_feature = np.mean(feature, axis=0, keepdims=True) raw_similarity = np.dot(object_feature, label_text_features.T) exp_sim = np.exp(raw_similarity * 100) prob = exp_sim / np.sum(exp_sim, axis=1, keepdims=True) probs = np.max(prob, axis=0) max_label_id = np.argmax(probs) prob = probs[max_label_id] pred_dict['pred_score'][idx] = prob label_id = label2id[descriptions[max_label_id]] pred_dict['pred_classes'][idx] = label_id point_ids = object_data['mask'] binary_mask = np.zeros(total_point_num, dtype=bool) binary_mask[list(point_ids)] = True pred_dict['pred_masks'][:, idx] = binary_mask if args.debug: os.makedirs(logs_path, exist_ok=True) cv2.imwrite(f'{logs_path}/{str(idx).zfill(5)}_{dataset.get_label_id()[1][label_id]}({prob:.2f}).jpg', saved_images) print(idx, label_id, dataset.get_label_id()[1][label_id], "confidence:", prob) pred_classes = pred_dict['pred_classes'] # pred_classes = [dataset.get_label_id()[1][i] for i in pred_classes] pred_dict['pred_masks'] = pred_dict['pred_masks'][:, pred_classes != 0] pred_dict['pred_score'] = pred_dict['pred_score'][pred_classes != 0] pred_dict['pred_classes'] = pred_dict['pred_classes'][pred_classes != 0] np.savez(f'{pred_dir}/{args.seq_name}.npz', **pred_dict) if __name__ == '__main__': parser = argparse.ArgumentParser() parser.add_argument('--config', type=str, default='scannet') parser.add_argument('--debug', action="store_true") parser.add_argument('--exp_name', type=str, default='baseline') parser.add_argument('--device', type=str, default='cuda:0') parser.add_argument('--num_images', type=int, default=5) parser.add_argument('--num_workers', '-n', type=int, default=1) parser.add_argument('--path_to_predictions', type=str, default='/home/jovyan/users/bulat/workspace/3drec/Indoor/Grounded-SAM-2/results/gsam_result/scannet200/yolo/memory_150_classes_198') parser.add_argument('--job_id', '-i', type=int, default=0) args = parser.parse_args() args = update_args(args) model, preprocess = load_clip(args.device) args.model = model args.preprocess = preprocess main(args)