Updated classifier

app.py

@@ -13,7 +13,7 @@ from huggingface_hub import hf_hub_download
 
 
 class FlakeLayerClassifier(nn.Module):
-    def __init__(self, num_materials, material_dim, num_classes=4, dropout_prob=0.1, freeze_cnn=False):
+    def __init__(self, num_materials, material_dim, num_classes=2, dropout_prob=0.1, freeze_cnn=False):
         super().__init__()
         self.cnn = ResNetModel.from_pretrained("microsoft/resnet-18")
         if freeze_cnn:
@@ -72,8 +72,8 @@ print(f"Using device: {device}")
 # Load YOLO detector
 yolo = YOLO("models/uark_detector_v3.pt")
 
-# Load classifier model checkpoint with proper device mapping
-ckpt_path = "models/flake_classifier_5layer.pth"
+# Load classifier model checkpoint
+ckpt_path = "models/flake_monolayer_classifier.pth"
 ckpt = torch.load(ckpt_path, map_location=device)
 
 num_classes   = len(ckpt["class_to_idx"])

models/flake_monolayer_classifier.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cc4e98bf4dd3127970ca7c68b633d29aa523a0a66a2d6481bf346d7662dbe7b8
+size 47191055