Spaces:
Running on Zero
Running on Zero
| from __future__ import annotations | |
| import argparse | |
| import os | |
| from pathlib import Path | |
| from typing import Dict, Tuple | |
| import sys | |
| import torch | |
| from tqdm import tqdm | |
| THIS_DIR = Path(__file__).resolve().parent | |
| REPO_ROOT = THIS_DIR.parents[2] | |
| if str(REPO_ROOT) not in sys.path: | |
| sys.path.insert(0, str(REPO_ROOT)) | |
| from quantization.escnet.data.calib import build_calib_loader | |
| from quantization.escnet.methods.registry import METHODS | |
| from quantization.escnet.quant.config import QuantConfig | |
| from quantization.escnet.utils.io import load_quant_checkpoint | |
| from quantization.escnet.utils.seed import set_seed | |
| def _import_load_config(): | |
| from model_sources.escnet.ESCNet_model.config import load_config | |
| return load_config | |
| def _import_escnet(): | |
| try: | |
| from model_sources.escnet.ESCNet_model.models.ESCNet import ESCNet # type: ignore | |
| return ESCNet | |
| except Exception as e: | |
| raise RuntimeError( | |
| "Failed to import ESCNet. Run this from repo root and ensure ./model_sources/escnet/ESCNet_model exists." | |
| ) from e | |
| def _try_import_dataset(): | |
| from model_sources.escnet.ESCNet_model.dataset import MyData # type: ignore | |
| return MyData | |
| def _try_import_save_tensor_img(): | |
| try: | |
| from model_sources.escnet.ESCNet_model.utils import save_tensor_img | |
| return save_tensor_img | |
| except Exception: | |
| # fallback | |
| from PIL import Image | |
| import numpy as np | |
| def save_tensor_img(t: torch.Tensor, out_path: str) -> None: | |
| arr = (t.detach().float().cpu().numpy().squeeze() * 255.0).clip(0, 255).astype("uint8") | |
| Image.fromarray(arr).save(out_path) | |
| return save_tensor_img | |
| def run_infer(model: torch.nn.Module, loader, pred_root: str, device: torch.device) -> None: | |
| import cv2 | |
| save_tensor_img = _try_import_save_tensor_img() | |
| model.eval() | |
| Path(pred_root).mkdir(parents=True, exist_ok=True) | |
| for batch in tqdm(loader, total=len(loader)): | |
| inputs = batch[0].to(device, non_blocking=True) | |
| label_paths = batch[-1] | |
| out_edge, scaled_preds = model(inputs) | |
| pred_lvl = (scaled_preds[-1].sigmoid() >= 0.5).float() | |
| for idx_sample in range(pred_lvl.shape[0]): | |
| res = torch.nn.functional.interpolate( | |
| pred_lvl[idx_sample].unsqueeze(0), | |
| size=cv2.imread(label_paths[idx_sample], cv2.IMREAD_GRAYSCALE).shape[:2], | |
| mode="bilinear", | |
| align_corners=True, | |
| ) | |
| out_path = os.path.join(pred_root, os.path.basename(label_paths[idx_sample])) | |
| save_tensor_img(res, out_path) | |
| def main() -> None: | |
| parser = argparse.ArgumentParser("ESCNet quant inference") | |
| parser.add_argument("--config", type=str, default="configs/escnet/main/config.yaml") | |
| parser.add_argument("--quant_ckpt", type=str, required=True, help="Quant checkpoint exported by quantize_escnet.py") | |
| parser.add_argument("--pred_root", type=str, default="preds_quant") | |
| parser.add_argument("--test_dir", type=str, default="", help="Override dataset dir (must contain Imgs/GT)") | |
| parser.add_argument("--device", type=str, default="cuda:0") | |
| parser.add_argument("--num_workers", type=int, default=4) | |
| parser.add_argument("--batch_size", type=int, default=-1, help="Override config.batch_size_valid") | |
| parser.add_argument("--seed", type=int, default=42) | |
| parser.add_argument("--verbose", action="store_true") | |
| args = parser.parse_args() | |
| set_seed(args.seed) | |
| load_config = _import_load_config() | |
| ESCNet = _import_escnet() | |
| MyData = _try_import_dataset() | |
| cfg = load_config(args.config) | |
| state, qcfg, meta = load_quant_checkpoint(args.quant_ckpt) | |
| device = torch.device(args.device if torch.cuda.is_available() else "cpu") | |
| model = ESCNet(cfg, pretrained=False) | |
| model.eval() | |
| # wrap with same qcfg then load quant state | |
| METHODS["ptq_qdq"](model, qcfg) | |
| missing, unexpected = model.load_state_dict(state, strict=False) | |
| if args.verbose: | |
| print(f"[load] missing={len(missing)}, unexpected={len(unexpected)}") | |
| model.to(device) | |
| test_dir = args.test_dir or getattr(cfg, "test_dir") | |
| bs = args.batch_size if args.batch_size > 0 else int(getattr(cfg, "batch_size_valid", 1)) | |
| loader = torch.utils.data.DataLoader( | |
| dataset=MyData(cfg, test_dir, image_size=int(getattr(cfg, "img_size", 384)), is_train=False), | |
| batch_size=bs, | |
| shuffle=False, | |
| num_workers=int(args.num_workers), | |
| pin_memory=True, | |
| ) | |
| run_infer(model, loader, pred_root=args.pred_root, device=device) | |
| print(f"[done] preds -> {args.pred_root}") | |
| if __name__ == "__main__": | |
| main() | |