|
|
@@ -26,8 +26,8 @@ model_dict = {model_name: None for model_name in model_selection_options} #set u
|
|
|
colors = [tuple([random.randint(0, 255) for _ in range(3)]) for _ in range(100)] #for bbox plotting
|
|
|
|
|
|
|
|
|
-if model_dict[model_name] is None:
|
|
|
- model_dict[model_name] = model = torch.hub.load(YOLO_DIR, 'custom', path=WEIGHTS, source='local').to(device)
|
|
|
+if model_dict['ocr-layout'] is None:
|
|
|
+ model_dict['ocr-layout'] = model = torch.hub.load(YOLO_DIR, 'custom', path=WEIGHTS, source='local').to(device)
|
|
|
print("========>加载成功")
|
|
|
|
|
|
|
