Spaces:

codedad
/

Tile_Defect_Inspector

Running

codedad commited on 30 days ago

Commit

9fba7e2

verified ·

1 Parent(s): e3303f6

Update app.py

Files changed (1) hide show

app.py CHANGED Viewed

@@ -8,7 +8,7 @@ from PIL import Image
 from torchvision import models
 model = models.resnet18()
 model.fc = torch.nn.Linear(model.fc.in_features, 2)
-model.load_state_dict(torch.load("tile_model.pt", map_location="cpu"))
 model.eval()
 # 2. Define labels based on your dataset folders

 from torchvision import models
 model = models.resnet18()
 model.fc = torch.nn.Linear(model.fc.in_features, 2)
+model.load_state_dict(torch.load("fine_tuned_model.pt", map_location="cpu"))
 model.eval()
 # 2. Define labels based on your dataset folders