53 lines
1.9 KiB
Python
53 lines
1.9 KiB
Python
"""Worker-only offline SegFormer comparison; no live control authority."""
|
|
|
|
import argparse
|
|
import hashlib
|
|
import json
|
|
import time
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
import torch
|
|
from PIL import Image
|
|
from transformers import SegformerForSemanticSegmentation, SegformerImageProcessor
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--image", type=Path, required=True)
|
|
parser.add_argument("--output", type=Path, required=True)
|
|
args = parser.parse_args()
|
|
processor = SegformerImageProcessor.from_pretrained("/assets", local_files_only=True)
|
|
model = (
|
|
SegformerForSemanticSegmentation.from_pretrained("/assets", local_files_only=True)
|
|
.cuda()
|
|
.eval()
|
|
)
|
|
image = Image.open(args.image).convert("RGB").crop((100, 0, 700, 600))
|
|
tensor = processor(images=image, return_tensors="pt")["pixel_values"].cuda()
|
|
times = []
|
|
with torch.inference_mode():
|
|
for _ in range(6):
|
|
torch.cuda.synchronize()
|
|
start = time.monotonic()
|
|
logits = model(tensor).logits
|
|
logits = torch.nn.functional.interpolate(
|
|
logits, size=(512, 512), mode="bilinear", align_corners=False
|
|
)
|
|
labels = logits.argmax(1)[0].cpu().numpy().astype(np.uint8)
|
|
times.append((time.monotonic() - start) * 1000)
|
|
args.output.mkdir(parents=True, exist_ok=True)
|
|
Image.fromarray(labels).save(args.output / "labels.png")
|
|
ids, counts = np.unique(labels, return_counts=True)
|
|
report = {
|
|
"source_sha256": hashlib.sha256(args.image.read_bytes()).hexdigest(),
|
|
"classes": {model.config.id2label[int(i)]: int(counts[n]) for n, i in enumerate(ids)},
|
|
"inference_ms": times,
|
|
}
|
|
(args.output / "report.json").write_text(json.dumps(report, indent=2), encoding="utf-8")
|
|
print(json.dumps(report))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|