Skip to content

Commit d1ee455

Browse files
JacobsolawetzJacobsolawetz
committed
not dry yolov5 upload start
1 parent e0b8ce3 commit d1ee455

File tree

1 file changed

+119
-74
lines changed

1 file changed

+119
-74
lines changed

‎roboflow/core/version.py‎

Lines changed: 119 additions & 74 deletions
Original file line numberDiff line numberDiff line change
@@ -291,7 +291,7 @@ def deploy(self, model_type: str, model_path: str) -> None:
291291
model_path (str): File path to model weights to be uploaded
292292
"""
293293

294-
supported_models = ["yolov8"]
294+
supported_models = ["yolov8", "yolov5"]
295295

296296
if model_type not in supported_models:
297297
raise (
@@ -300,16 +300,112 @@ def deploy(self, model_type: str, model_path: str) -> None:
300300
)
301301
)
302302

303-
try:
304-
import torch
305-
import ultralytics
306-
except ImportError as e:
307-
raise (
308-
"The ultralytics python package is required to deploy yolov8 models. Please install it with `pip install ultralytics`"
303+
if model_type == "yolov8":
304+
try:
305+
import torch
306+
import ultralytics
307+
except ImportError as e:
308+
raise (
309+
"The ultralytics python package is required to deploy yolov8 models. Please install it with `pip install ultralytics`"
310+
)
311+
312+
# add logic to save torch state dict safely
313+
if model_type == "yolov8":
314+
model = torch.load(model_path + "weights/best.pt")
315+
316+
class_names = []
317+
for i, val in enumerate(model["model"].names):
318+
class_names.append((val, model["model"].names[val]))
319+
class_names.sort(key=lambda x: x[0])
320+
class_names = [x[1] for x in class_names]
321+
322+
try:
323+
model_artifacts = {
324+
"names": class_names,
325+
"yaml": model["model"].yaml,
326+
"nc": model["model"].nc,
327+
"args": {
328+
k: val
329+
for k, val in model["model"].args.items()
330+
if ((k == "model") or (k == "imgsz") or (k == "batch"))
331+
},
332+
"ultralytics_version": ultralytics.__version__,
333+
"model_type": model_type,
334+
}
335+
except:
336+
model_artifacts = {
337+
"names": class_names,
338+
"yaml": model["model"].yaml,
339+
"nc": model["model"].nc,
340+
"args": {
341+
k: val
342+
for k, val in model["model"].args.__dict__.items()
343+
if ((k == "model") or (k == "imgsz") or (k == "batch"))
344+
},
345+
"ultralytics_version": ultralytics.__version__,
346+
"model_type": model_type,
347+
}
348+
349+
with open(model_path + "model_artifacts.json", "w") as fp:
350+
json.dump(model_artifacts, fp)
351+
352+
torch.save(model["model"].state_dict(), model_path + "state_dict.pt")
353+
354+
lista_files = [
355+
"results.csv",
356+
"results.png",
357+
"model_artifacts.json",
358+
"state_dict.pt",
359+
]
360+
with zipfile.ZipFile(model_path + "roboflow_deploy.zip", "w") as zipMe:
361+
for file in lista_files:
362+
zipMe.write(
363+
model_path + file,
364+
arcname=file,
365+
compress_type=zipfile.ZIP_DEFLATED,
366+
)
367+
368+
res = requests.get(
369+
f"{API_URL}/{self.workspace}/{self.project}/{self.version}/uploadModel?api_key={self.__api_key}"
370+
)
371+
try:
372+
if res.status_code == 429:
373+
raise RuntimeError(
374+
f"This version already has a trained model. Please generate and train a new version in order to upload model to Roboflow."
375+
)
376+
else:
377+
res.raise_for_status()
378+
except Exception as e:
379+
print(f"An error occured when getting the model upload URL: {e}")
380+
return
381+
382+
res = requests.put(
383+
res.json()["url"], data=open(model_path + "roboflow_deploy.zip", "rb")
309384
)
385+
try:
386+
res.raise_for_status()
387+
388+
if self.public:
389+
print(
390+
f"View the status of your deployment at: {APP_URL}/{self.workspace}/{self.project}/deploy/{self.version}"
391+
)
392+
print(
393+
f"Share your model with the world at: {UNIVERSE_URL}/{self.workspace}/{self.project}/model/{self.version}"
394+
)
395+
else:
396+
print(
397+
f"View the status of your deployment at: {APP_URL}/{self.workspace}/{self.project}/deploy/{self.version}"
398+
)
399+
400+
except Exception as e:
401+
print(f"An error occured when uploading the model: {e}")
402+
elif model_type == "yolov5":
403+
try:
404+
import torch
405+
import yaml
406+
except ImportError as e:
407+
raise ("PyTorch must be installed to use this feature.")
310408

311-
# add logic to save torch state dict safely
312-
if model_type == "yolov8":
313409
model = torch.load(os.path.join(model_path + "weights/best.pt"))
314410

315411
class_names = []
@@ -318,34 +414,18 @@ def deploy(self, model_type: str, model_path: str) -> None:
318414
class_names.sort(key=lambda x: x[0])
319415
class_names = [x[1] for x in class_names]
320416

321-
try:
322-
model_artifacts = {
323-
"names": class_names,
324-
"yaml": model["model"].yaml,
325-
"nc": model["model"].nc,
326-
"args": {
327-
k: val
328-
for k, val in model["model"].args.items()
329-
if ((k == "model") or (k == "imgsz") or (k == "batch"))
330-
},
331-
"ultralytics_version": ultralytics.__version__,
332-
"model_type": model_type,
333-
}
334-
except:
335-
model_artifacts = {
336-
"names": class_names,
337-
"yaml": model["model"].yaml,
338-
"nc": model["model"].nc,
339-
"args": {
340-
k: val
341-
for k, val in model["model"].args.__dict__.items()
342-
if ((k == "model") or (k == "imgsz") or (k == "batch"))
343-
},
344-
"ultralytics_version": ultralytics.__version__,
345-
"model_type": model_type,
346-
}
347-
348-
with open(os.path.join(model_path + "model_artifacts.json", "w")) as fp:
417+
with open(os.path.join(model_path + "opt.yaml"), "r") as stream:
418+
opts = yaml.safe_load(stream)
419+
420+
model_artifacts = {
421+
"names": class_names,
422+
"yaml": model["model"].yaml,
423+
"nc": model["model"].nc,
424+
"args": {"imgsz": opts["imgsz"], "batch": opts["batch_size"]},
425+
"model_type": model_type,
426+
}
427+
428+
with open(os.path.join(model_path + "model_artifacts.json"), "w") as fp:
349429
json.dump(model_artifacts, fp)
350430

351431
torch.save(
@@ -358,8 +438,9 @@ def deploy(self, model_type: str, model_path: str) -> None:
358438
"model_artifacts.json",
359439
"state_dict.pt",
360440
]
441+
361442
with zipfile.ZipFile(
362-
os.path.join(model_path + "roboflow_deploy.zip", "w")
443+
os.path.join(model_path + "roboflow_deploy.zip"), "w"
363444
) as zipMe:
364445
for file in lista_files:
365446
zipMe.write(
@@ -368,42 +449,6 @@ def deploy(self, model_type: str, model_path: str) -> None:
368449
compress_type=zipfile.ZIP_DEFLATED,
369450
)
370451

371-
res = requests.get(
372-
f"{API_URL}/{self.workspace}/{self.project}/{self.version}/uploadModel?api_key={self.__api_key}"
373-
)
374-
try:
375-
if res.status_code == 429:
376-
raise RuntimeError(
377-
f"This version already has a trained model. Please generate and train a new version in order to upload model to Roboflow."
378-
)
379-
else:
380-
res.raise_for_status()
381-
except Exception as e:
382-
print(f"An error occured when getting the model upload URL: {e}")
383-
return
384-
385-
res = requests.put(
386-
res.json()["url"],
387-
data=open(os.path.join(model_path + "roboflow_deploy.zip", "rb")),
388-
)
389-
try:
390-
res.raise_for_status()
391-
392-
if self.public:
393-
print(
394-
f"View the status of your deployment at: {APP_URL}/{self.workspace}/{self.project}/deploy/{self.version}"
395-
)
396-
print(
397-
f"Share your model with the world at: {UNIVERSE_URL}/{self.workspace}/{self.project}/model/{self.version}"
398-
)
399-
else:
400-
print(
401-
f"View the status of your deployment at: {APP_URL}/{self.workspace}/{self.project}/deploy/{self.version}"
402-
)
403-
404-
except Exception as e:
405-
print(f"An error occured when uploading the model: {e}")
406-
407452
# torch.load("state_dict.pt", weights_only=True)
408453

409454
def __download_zip(self, link, location, format):

0 commit comments

Comments
 (0)