|
|
@@ -14,8 +14,8 @@ YOLO_DIR = '/workspace/yolov5'
|
|
|
# WEIGHTS = '/data/yolov5/runs/train/yolov5x_layout_reuslt37/weights/best.pt'
|
|
|
WEIGHTS = '/workspace/best.pt'
|
|
|
|
|
|
-device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
|
|
|
-
|
|
|
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
|
+print(device)
|
|
|
|
|
|
app = FastAPI()
|
|
|
templates = Jinja2Templates(directory = 'templates')
|
