Skip to content

Commit a5f612c

Browse files
committed
different approach for yolov12
1 parent 5c08cb7 commit a5f612c

File tree

1 file changed

+61
-14
lines changed

1 file changed

+61
-14
lines changed

roboflow/util/model_processor.py

Lines changed: 61 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -56,6 +56,9 @@ def _get_processor_function(model_type: str) -> Callable:
5656

5757
if "yolonas" in model_type:
5858
return _process_yolonas
59+
60+
if "yolov12" in model_type:
61+
return _process_yolov12
5962

6063
return _process_yolo
6164

@@ -107,18 +110,6 @@ def _process_yolo(model_type: str, model_path: str, filename: str) -> str:
107110

108111
print_warn_for_wrong_dependencies_versions([("ultralytics", ">=", "8.3.0")], ask_to_continue=True)
109112

110-
elif "yolov12" in model_type:
111-
try:
112-
import torch
113-
import ultralytics
114-
except ImportError:
115-
raise RuntimeError(
116-
"The ultralytics python package is required to deploy yolov12"
117-
" models. Please install it with `pip install ultralytics`"
118-
)
119-
120-
print_warn_for_wrong_dependencies_versions([("ultralytics", ">=", "8.3.78")], ask_to_continue=True)
121-
122113
model = torch.load(os.path.join(model_path, filename))
123114

124115
if isinstance(model["model"].names, list):
@@ -130,13 +121,12 @@ def _process_yolo(model_type: str, model_path: str, filename: str) -> str:
130121
class_names.sort(key=lambda x: x[0])
131122
class_names = [x[1] for x in class_names]
132123

133-
if "yolov8" in model_type or "yolov10" in model_type or "yolov11" in model_type or "yolov12" in model_type:
124+
if "yolov8" in model_type or "yolov10" in model_type or "yolov11" in model_type:
134125
# try except for backwards compatibility with older versions of ultralytics
135126
if (
136127
"-cls" in model_type
137128
or model_type.startswith("yolov10")
138129
or model_type.startswith("yolov11")
139-
or model_type.startswith("yolov12")
140130
):
141131
nc = model["model"].yaml["nc"]
142132
args = model["train_args"]
@@ -210,6 +200,63 @@ def _process_yolo(model_type: str, model_path: str, filename: str) -> str:
210200
return zip_file_name
211201

212202

203+
def _process_yolov12(model_type: str, model_path: str, filename: str) -> str:
204+
# For YOLOv12, since it uses a special Ultralytics version,
205+
# state dict extraction and model artifacts are handled during model conversion
206+
207+
print(
208+
"Note: Model must be trained using ultralytics from https://github.com/sunsmarterjie/yolov12 "
209+
"or through the Roboflow platform"
210+
)
211+
212+
# Check if model_path exists
213+
if not os.path.exists(model_path):
214+
raise FileNotFoundError(f"Model path {model_path} does not exist.")
215+
216+
# Find any .pt file in model path
217+
model_files = os.listdir(model_path)
218+
pt_file = next((f for f in model_files if f.endswith('.pt')), None)
219+
220+
if pt_file is None:
221+
raise RuntimeError("No .pt model file found in the provided path")
222+
223+
# Copy the .pt file to weights.pt if not already named weights.pt
224+
if pt_file != "weights.pt":
225+
shutil.copy(
226+
os.path.join(model_path, pt_file),
227+
os.path.join(model_path, "weights.pt")
228+
)
229+
230+
required_files = [
231+
"weights.pt"
232+
]
233+
234+
optional_files = [
235+
"results.csv",
236+
"results.png",
237+
"model_artifacts.json"
238+
]
239+
240+
zip_file_name = "roboflow_deploy.zip"
241+
with zipfile.ZipFile(os.path.join(model_path, zip_file_name), "w") as zipMe:
242+
for file in required_files:
243+
zipMe.write(
244+
os.path.join(model_path, file),
245+
arcname=file,
246+
compress_type=zipfile.ZIP_DEFLATED
247+
)
248+
249+
for file in optional_files:
250+
if os.path.exists(os.path.join(model_path, file)):
251+
zipMe.write(
252+
os.path.join(model_path, file),
253+
arcname=file,
254+
compress_type=zipfile.ZIP_DEFLATED
255+
)
256+
257+
return zip_file_name
258+
259+
213260
def _process_huggingface(
214261
model_type: str, model_path: str, filename: str = "fine-tuned-paligemma-3b-pt-224.f16.npz"
215262
) -> str:

0 commit comments

Comments
 (0)