phucd commited on
Commit ·
3bd068b
1
Parent(s): 92b5cbe
Minor bug fixed
Browse files
app.py
CHANGED
|
@@ -1,5 +1,3 @@
|
|
| 1 |
-
import os
|
| 2 |
-
import matplotlib.pyplot as plt
|
| 3 |
import cv2
|
| 4 |
import torch
|
| 5 |
import torch.nn as nn
|
|
@@ -35,7 +33,7 @@ class U2NETP_DocSeg(nn.Module):
|
|
| 35 |
docseg = U2NETP_DocSeg(num_classes=1)
|
| 36 |
# Load pretrained weights
|
| 37 |
docseg_weight_path = './weights/u2netp_docseg_epoch_225_date_2026-01-02.pth'
|
| 38 |
-
checkpoint = torch.load(docseg_weight_path)
|
| 39 |
docseg.load_state_dict(checkpoint[f"model_state_dict"])
|
| 40 |
docseg.eval()
|
| 41 |
|
|
|
|
|
|
|
|
|
|
| 1 |
import cv2
|
| 2 |
import torch
|
| 3 |
import torch.nn as nn
|
|
|
|
| 33 |
docseg = U2NETP_DocSeg(num_classes=1)
|
| 34 |
# Load pretrained weights
|
| 35 |
docseg_weight_path = './weights/u2netp_docseg_epoch_225_date_2026-01-02.pth'
|
| 36 |
+
checkpoint = torch.load(docseg_weight_path, map_location=torch.device('cpu'))
|
| 37 |
docseg.load_state_dict(checkpoint[f"model_state_dict"])
|
| 38 |
docseg.eval()
|
| 39 |
|