Spaces:
Running on Zero
Running on Zero
File size: 4,731 Bytes
fb5b18e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 | 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
@torch.no_grad()
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()
|