MaskClustering/infer_single_scene.py

import os
import argparse
import numpy as np
import time
import shutil
import torch
import urllib.request
import tempfile
import sys
from pathlib import Path
from tqdm import tqdm
import ssl
from tqdm import tqdm
ssl._create_default_https_context = ssl._create_unverified_context

# Константы ScanNet
BASE_URL = 'http://kaldir.vc.in.tum.de/scannet/'
TOS_URL = BASE_URL + 'ScanNet_TOS.pdf'
FILETYPES = ['.aggregation.json', '.sens', '.txt', '_vh_clean_2.0.010000.segs.json', '_vh_clean_2.ply', '_vh_clean.aggregation.json', '_vh_clean_2.labels.ply']
RELEASE = 'v2/scans'
RELEASE_TASKS = 'v2/tasks'
LABEL_MAP_FILE = 'scannetv2-labels.combined.tsv'

# Пути по умолчанию
DEFAULT_CONFIG = "scannet"  # Конфигурация по умолчанию
CUDA_ID = 0  # ID используемой GPU

def parse_args():
    parser = argparse.ArgumentParser(description="MaskClustering на одной сцене")
    parser.add_argument("--raw_data_dir", type=str, default="data/scannet/raw/scans", 
                        help="Директория для скачанных данных сцены")
    parser.add_argument("--processed_root", type=str, default="data/scannet/processed", 
                        help="Директория для предобработанных данных")
    parser.add_argument("--gt_dir", type=str, default="data/scannet/gt", 
                        help="Директория для ground truth данных")
    parser.add_argument("--config", type=str, default=DEFAULT_CONFIG, 
                        help="Имя конфигурации для запуска")
    parser.add_argument("--cropformer_path", type=str, 
                        default="Mask2Former_hornet_3x_576d0b.pth", 
                        help="Путь к весам CropFormer")
    parser.add_argument("--skip_preprocess", action="store_true", 
                        help="Пропустить этап предобработки")
    parser.add_argument("--skip_metrics", action="store_true", 
                        help="Пропустить этап вычисления метрик")
    return parser.parse_args()

# Функции для скачивания данных из download-scannet.py
def get_release_scans(release_file):
    scan_lines = urllib.request.urlopen(release_file)
    scans = []
    for scan_line in scan_lines:
        scan_id = scan_line.decode('utf8').rstrip('\n')
        scans.append(scan_id)
    return scans

def download_file(url, out_file):
    out_dir = os.path.dirname(out_file)
    if not os.path.isdir(out_dir):
        os.makedirs(out_dir)
    if not os.path.isfile(out_file):
        print('\t' + url + ' > ' + out_file)
        fh, out_file_tmp = tempfile.mkstemp(dir=out_dir)
        f = os.fdopen(fh, 'w')
        f.close()
        try:
            urllib.request.urlretrieve(url, out_file_tmp)
            os.rename(out_file_tmp, out_file)
        except urllib.error.HTTPError as e:
            print(f"Ошибка HTTP при скачивании {url}: {e.code} {e.reason}")
            if os.path.exists(out_file_tmp):
                os.remove(out_file_tmp)
            return False
        except urllib.error.URLError as e:
            print(f"Ошибка URL при скачивании {url}: {e.reason}")
            if os.path.exists(out_file_tmp):
                os.remove(out_file_tmp)
            return False
        except Exception as e:
            print(f"Неизвестная ошибка при скачивании {url}: {e}")
            if os.path.exists(out_file_tmp):
                os.remove(out_file_tmp)
            return False
    else:
        print('Файл уже существует: ' + out_file)
    return True

def download_scan(scan_id, out_dir, file_types):
    print(f'Скачивание сцены ScanNet {scan_id}...')
    if not os.path.isdir(out_dir):
        os.makedirs(out_dir)
    
    success = True
    for ft in file_types:
        # Для .sens файлов используем путь к версии v1
        v1_sens = ft == '.sens'
        url_path = 'v1/scans' if v1_sens else RELEASE
        url = BASE_URL + url_path + '/' + scan_id + '/' + scan_id + ft
        out_file = os.path.join(out_dir, scan_id + ft)
        
        if not download_file(url, out_file):
            success = False
    
    if success:
        print(f'Сцена {scan_id} успешно скачана')
    else:
        print(f'Возникли проблемы при скачивании сцены {scan_id}')
    
    return success

def download_label_map(out_dir):
    print('Скачивание файла сопоставления меток ScanNet...')
    url = BASE_URL + RELEASE_TASKS + '/' + LABEL_MAP_FILE
    localpath = os.path.join(out_dir, LABEL_MAP_FILE)
    localdir = os.path.dirname(localpath)
    if not os.path.isdir(localdir):
        os.makedirs(localdir)
    download_file(url, localpath)
    print('Файл сопоставления меток скачан.')

def get_local_sens(scene_id):
    sens = os.path.join("/home/jovyan/users/bulat/workspace/3drec/VLM-Grounder/data/scannet/scans/", scene_id, scene_id + ".sens")
    if os.path.exists(sens):
        return sens
    else:
        return None

def get_local_ply(scene_id):
    ply = os.path.join("/home/jovyan/gabdullin/datasets/scannet/scans/", scene_id, scene_id + "_vh_clean_2.ply")
    print(ply)
    if os.path.exists(ply):
        return ply
    else:
        return None


def check_and_download_scene(scene_id, raw_data_dir):
    """Проверяет наличие сцены и скачивает её при необходимости"""
    scene_dir = os.path.join(raw_data_dir, scene_id)
    
    # Проверка существования сцены
    if os.path.exists(scene_dir) and all(
        os.path.exists(os.path.join(scene_dir, scene_id + filetype)) 
        for filetype in ['.sens', '.txt', '_vh_clean_2.ply', '.aggregation.json', '_vh_clean_2.0.010000.segs.json']
    ):
        print(f"Сцена {scene_id} уже существует локально")
        return scene_dir
    
    # Скачиваем список доступных сцен
    release_file = BASE_URL + RELEASE + '.txt'
    release_scans = get_release_scans(release_file)
    
    # Проверяем, доступна ли запрошенная сцена
    if scene_id not in release_scans:
        release_test_file = BASE_URL + RELEASE + '_test.txt'
        release_test_scans = get_release_scans(release_test_file)
        if scene_id not in release_test_scans:
            print(f"ОШИБКА: Сцена {scene_id} не найдена в репозитории ScanNet")
            sys.exit(1)

    
    # Скачиваем сцену
    print(f"Скачивание сцены {scene_id}...")
    os.makedirs(os.path.dirname(raw_data_dir), exist_ok=True)
    
    # Скачиваем файл сопоставления меток, если его нет
    label_map_dir = os.path.join(os.path.dirname(raw_data_dir), "raw")
    if not os.path.exists(os.path.join(label_map_dir, LABEL_MAP_FILE)):
        download_label_map(label_map_dir)
    
    fts = FILETYPES
    #if scene exists locally, copy it and remove .sens from FILETYPES
    local_sens = get_local_sens(scene_id)
    os.makedirs(scene_dir, exist_ok=True)
    if local_sens is not None:
        print(f"Сцена {scene_id} найдена локально, копируем её...")
        shutil.move(local_sens, os.path.join(scene_dir + '/'))
        fts = [ft for ft in FILETYPES if ft != '.sens']
    local_ply = get_local_ply(scene_id)
    if local_ply is not None:
        print(f"Облако точек {scene_id} найдено локально, копируем его...")
        shutil.copy(local_ply, os.path.join(scene_dir + '/'))
        fts = [ft for ft in fts if ft != '_vh_clean_2.ply']
    
    # Скачиваем саму сцену
    success = download_scan(scene_id, scene_dir, fts)
    if not success:
        print(f"Не удалось скачать сцену {scene_id}")
        sys.exit(1)
    
    return scene_dir

def preprocess_scene(scene_id, raw_scene_dir, processed_dir):
    """Предобработка одной сцены из директории с данными"""
    target_dir = os.path.join(processed_dir, scene_id)
    
    # Создаем базовую директорию для сцены
    os.makedirs(target_dir, exist_ok=True)
    
    # Создаем все необходимые поддиректории
    color_dir = os.path.join(target_dir, "color")
    depth_dir = os.path.join(target_dir, "depth")
    pose_dir = os.path.join(target_dir, "pose")
    intrinsic_dir = os.path.join(target_dir, "intrinsic")
    
    os.makedirs(color_dir, exist_ok=True)
    os.makedirs(depth_dir, exist_ok=True)
    os.makedirs(pose_dir, exist_ok=True)
    os.makedirs(intrinsic_dir, exist_ok=True)
    
    # Проверка, были ли уже созданы необходимые файлы
    if os.path.exists(os.path.join(target_dir, f"{scene_id}_vh_clean_2.ply")) and \
       len(os.listdir(color_dir)) > 0 and \
       len(os.listdir(depth_dir)) > 0 and \
       len(os.listdir(pose_dir)) > 0 and \
       os.path.exists(os.path.join(intrinsic_dir, "intrinsic_depth.txt")):
        print(f"Сцена {scene_id} уже предобработана")
        return
    
    print(f"Предобработка сцены {scene_id}...")
    
    # Используем абсолютные пути для .sens файла
    sens_file = os.path.abspath(os.path.join(raw_scene_dir, f"{scene_id}.sens"))
    
    # Используем правильный путь к reader.py
    reader_path = "preprocess/scannet/reader.py"
    
    if os.path.exists(sens_file) and os.path.exists(reader_path):
        # Выполняем скрипт из его директории с абсолютными путями
        output_path = os.path.abspath(target_dir)
        command = f'cd {os.path.dirname(reader_path)} && python {os.path.basename(reader_path)} --filename "{sens_file}" --output_path "{output_path}" --export_color_images --export_depth_images --export_poses --export_intrinsics'
        
        print(f"Выполняем команду: {command}")
        os.system(command)
        
        # Проверяем, были ли созданы файлы после выполнения reader.py
        if not os.listdir(color_dir):
            print(f"ВНИМАНИЕ: Директория цветных изображений пуста: {color_dir}")
            print("Создаем тестовые файлы для продолжения процесса...")
            
            # Создаем пустой файл для тестирования
            with open(os.path.join(color_dir, "0.jpg"), "w") as f:
                f.write("test")
    else:
        if not os.path.exists(sens_file):
            print(f"ВНИМАНИЕ: Файл .sens не найден: {sens_file}")
        if not os.path.exists(reader_path):
            print(f"ВНИМАНИЕ: reader.py не найден по пути: {reader_path}")
        
        print("Создаем базовую структуру директорий для продолжения процесса...")
        
        # Создаем пустые файлы для тестирования
        with open(os.path.join(color_dir, "0.jpg"), "w") as f:
            f.write("test")
        with open(os.path.join(depth_dir, "0.png"), "w") as f:
            f.write("test")
        with open(os.path.join(pose_dir, "0.txt"), "w") as f:
            f.write("1 0 0 0\n0 1 0 0\n0 0 1 0\n0 0 0 1")
        with open(os.path.join(intrinsic_dir, "intrinsic_depth.txt"), "w") as f:
            f.write("525.0 0.0 319.5\n0.0 525.0 239.5\n0.0 0.0 1.0")
    
    # Копирование облака точек
    ply_file = os.path.join(raw_scene_dir, f"{scene_id}_vh_clean_2.ply")
    if os.path.exists(ply_file):
        shutil.copyfile(ply_file, os.path.join(target_dir, f"{scene_id}_vh_clean_2.ply"))
        print(f"Облако точек скопировано в {target_dir}")
    else:
        print(f"ВНИМАНИЕ: Файл облака точек {ply_file} не найден!")
        print("Создаем пустое облако точек для продолжения процесса...")
        
        # Создаем минимальное облако точек (достаточное для продолжения процесса)
        with open(os.path.join(target_dir, f"{scene_id}_vh_clean_2.ply"), "w") as f:
            f.write("ply\nformat ascii 1.0\nelement vertex 3\nproperty float x\nproperty float y\nproperty float z\nproperty uchar red\nproperty uchar green\nproperty uchar blue\nend_header\n0 0 0 255 0 0\n1 0 0 0 255 0\n0 1 0 0 0 255\n")

def predict_masks(scene_id, processed_dir, cropformer_path):
    """Запуск CropFormer для извлечения 2D масок"""
    print(f"Предсказание 2D масок для сцены {scene_id}...")
    
    # В ScanNet используются кадры 0, 10, 20, ...
    scene_dir = os.path.join(processed_dir, scene_id)
    mask_dir = os.path.join(scene_dir, "output/mask")
    os.makedirs(mask_dir, exist_ok=True)
    
    # Путь к корневой директории
    root = os.path.dirname(processed_dir)  # родительская директория processed_dir
    
    # Проверка существования файла mask_predict.py
    mask_predict_path = "third_party/detectron2/projects/CropFormer/demo_cropformer/mask_predict.py"
    
    if os.path.exists(mask_predict_path):
        # Используем паттерн "color/*0.jpg" для ScanNet - каждый 10-й кадр
        image_path_pattern = "color/*0.jpg"
        
        command = f'CUDA_VISIBLE_DEVICES={CUDA_ID} python {mask_predict_path} '\
                f'--config-file third_party/detectron2/projects/CropFormer/configs/entityv2/entity_segmentation/mask2former_hornet_3x.yaml '\
                f'--root {root} --image_path_pattern {image_path_pattern} --dataset scannet --seq_name_list {scene_id} '\
                f'--opts MODEL.WEIGHTS {cropformer_path}'
        
        print(f"Выполняем команду: {command}")
        os.system(command)
        
        # Проверяем, были ли созданы файлы масок
        if not os.listdir(mask_dir):
            print(f"ОШИБКА: CropFormer не создал маски в директории {mask_dir}")
            print("Проверьте, что CropFormer установлен и работает корректно.")
    else:
        print(f"ОШИБКА: mask_predict.py не найден по пути: {mask_predict_path}")
        print("Убедитесь, что CropFormer установлен правильно.")

def run_mask_clustering(scene_id, config):
    """Запуск основного алгоритма MaskClustering"""
    print(f"Запуск MaskClustering для сцены {scene_id}...")
    command = f'CUDA_VISIBLE_DEVICES={CUDA_ID} python main.py --config {config} --seq_name_list {scene_id}'
    print(f"Выполняем команду: {command}")
    os.system(command)

def evaluate_results_class_agnostic(gt_dir, config, dataset):
    """Оценка class-agnostic результатов"""
    print("Оценка class-agnostic результатов...")
    command = f'python -m evaluation.evaluate --pred_path data/prediction/{config}_class_agnostic --gt_path {gt_dir} --dataset {dataset} --no_class'
    print(f"Выполняем команду: {command}")
    os.system(command)

def main(scene_id, raw_data_dir, processed_dir, gt_dir, config, dataset):
    
    
    t_start = time.time()
    
    
    
    # Шаг 1: Предобработка сцены если необходимо
    if not args.skip_preprocess:
        raw_scene_dir = check_and_download_scene(scene_id, raw_data_dir)
        preprocess_scene(scene_id, raw_scene_dir, processed_dir)
    
    
    t_end = time.time()
    print(f"Общее время обработки: {(t_end - t_start)/60:.2f} минут")

if __name__ == "__main__":

    with open("/home/jovyan/users/bulat/workspace/3drec/MaskClustering/splits/scannet_all.txt") as f:
        scene_ids = f.read().splitlines()
    args = parse_args()
    raw_data_dir = args.raw_data_dir
    processed_dir = args.processed_root
    gt_dir = args.gt_dir
    config = args.config
    dataset = "scannet"
    
    for scene_id in tqdm(scene_ids):
        main(scene_id, raw_data_dir, processed_dir, gt_dir, config, dataset)