fix: refuse unavailable gallery and pin head trainer
This commit is contained in:
parent
d5a7a56dff
commit
ac315eb48e
4 changed files with 77 additions and 6 deletions
|
|
@ -132,12 +132,14 @@ def enrollment_report_text(report: dict) -> str:
|
|||
|
||||
|
||||
def build_model(image_size: int = 518):
|
||||
import timm
|
||||
import torch
|
||||
|
||||
from gallery import load_dinov2
|
||||
|
||||
device = "mps" if torch.backends.mps.is_available() else "cpu"
|
||||
model = timm.create_model("vit_small_patch14_dinov2.lvd142m", pretrained=True,
|
||||
num_classes=0, img_size=image_size, dynamic_img_size=True)
|
||||
# Same hash-pinned local file as inference. pretrained=True would fetch a
|
||||
# different copy even with YOLO_OFFLINE=1.
|
||||
model = load_dinov2(image_size)
|
||||
return model.eval().to(device), device
|
||||
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue