|
2 | 2 | # Code written by Fatih C Akyon and Devrim Cavusoglu, 2022.
|
3 | 3 |
|
4 | 4 | import logging
|
| 5 | +import os |
5 | 6 | from typing import Any, Dict, List, Optional, Tuple, Union
|
6 | 7 |
|
7 | 8 | import numpy as np
|
@@ -69,18 +70,19 @@ def num_categories(self) -> int:
|
69 | 70 | def load_model(self):
|
70 | 71 | from transformers import AutoModelForObjectDetection, AutoProcessor
|
71 | 72 |
|
72 |
| - model = AutoModelForObjectDetection.from_pretrained(self.model_path, token=self._token) |
| 73 | + hf_token = os.getenv("HF_TOKEN", self._token) |
| 74 | + model = AutoModelForObjectDetection.from_pretrained(self.model_path, token=hf_token) |
73 | 75 | if self.image_size is not None:
|
74 | 76 | if model.base_model_prefix == "rt_detr_v2":
|
75 | 77 | size = {"height": self.image_size, "width": self.image_size}
|
76 | 78 | else:
|
77 | 79 | size = {"shortest_edge": self.image_size, "longest_edge": None}
|
78 | 80 | # use_fast=True raises error: AttributeError: 'SizeDict' object has no attribute 'keys'
|
79 | 81 | processor = AutoProcessor.from_pretrained(
|
80 |
| - self.model_path, size=size, do_resize=True, use_fast=False, token=self._token |
| 82 | + self.model_path, size=size, do_resize=True, use_fast=False, token=hf_token |
81 | 83 | )
|
82 | 84 | else:
|
83 |
| - processor = AutoProcessor.from_pretrained(self.model_path, use_fast=False, token=self._token) |
| 85 | + processor = AutoProcessor.from_pretrained(self.model_path, use_fast=False, token=hf_token) |
84 | 86 | self.set_model(model, processor)
|
85 | 87 |
|
86 | 88 | def set_model(self, model: Any, processor: Any = None):
|
|
0 commit comments