Fix yaml import

Author: lorenmt

experts/generate_depth.py

@@ -6,8 +6,11 @@
 
 import torch
 import os
-import ruamel_yaml as yaml
-
+try:
+    import ruamel_yaml as yaml
+except ModuleNotFoundError:
+    import ruamel.yaml as yaml
+    
 from experts.model_bank import load_expert_model
 from experts.depth.generate_dataset import Dataset
 import PIL.Image as Image

experts/generate_edge.py

@@ -6,7 +6,10 @@
 
 import torch
 import os
-import ruamel_yaml as yaml
+try:
+    import ruamel_yaml as yaml
+except ModuleNotFoundError:
+    import ruamel.yaml as yaml
 
 from experts.model_bank import load_expert_model
 from experts.edge.generate_dataset import Dataset

experts/generate_normal.py

@@ -6,8 +6,11 @@
 
 import torch
 import os
-import ruamel_yaml as yaml
-
+try:
+    import ruamel_yaml as yaml
+except ModuleNotFoundError:
+    import ruamel.yaml as yaml
+    
 from experts.model_bank import load_expert_model
 from experts.normal.generate_dataset import CustomDataset
 import PIL.Image as Image

experts/generate_objdet.py

@@ -6,11 +6,14 @@
 
 import torch
 import os
-import ruamel_yaml as yaml
 import json
 import copy
 import PIL.Image as Image
-
+try:
+    import ruamel_yaml as yaml
+except ModuleNotFoundError:
+    import ruamel.yaml as yaml
+    
 from experts.model_bank import load_expert_model
 from experts.obj_detection.generate_dataset import Dataset, collate_fn
 from accelerate import Accelerator

experts/generate_ocrdet.py

@@ -11,8 +11,10 @@
 import cv2
 import clip
 import pickle as pk
-import ruamel_yaml as yaml
-
+try:
+    import ruamel_yaml as yaml
+except ModuleNotFoundError:
+    import ruamel.yaml as yaml
 
 from experts.model_bank import load_expert_model
 from experts.ocr_detection.generate_dataset import Dataset

experts/generate_segmentation.py

@@ -7,8 +7,11 @@
 import torch
 import os
 import PIL.Image as Image
-import ruamel_yaml as yaml
-
+try:
+    import ruamel_yaml as yaml
+except ModuleNotFoundError:
+    import ruamel.yaml as yaml
+    
 from experts.model_bank import load_expert_model
 from experts.segmentation.generate_dataset import Dataset, collate_fn
 from accelerate import Accelerator