first commit

Pretrained_vit_feature_extractor_RDDV2AMC.pth

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:f1cca05ced71d79538362045c346a0a223fd53057a45770798a83324a99c006b
+oid sha256:15852098f6b673a52c5dfba214406c21d66039b96e95eedf8bc06b7f03fccc82
 size 343273361

app.py

@@ -13,7 +13,7 @@ class_names = ["D00-Longitudinal Crack", "D10-Transverse Crack", "D20-Aligator C
 
 ### 2. Model and transforms preparation ###
 
-# Create EffNetB2 model
+# Create ViT model
 vit, vit_transforms = create_vit_model(
     num_classes=len(class_names), # len(class_names) would also work
 )
@@ -58,7 +58,7 @@ def predict(img) -> Tuple[Dict, float]:
 # Create title, description and article strings
 title = "Road Damages Detection 🚧"
 description = "An ViT feature extractor computer vision model to classify images of common road damages."
-article = "Created at [Road Damages Detection](https:/xxx/)."
+article = "Created at https://huggingface.co/spaces/erictam/Road-Damages-Detection_Multi-Class"
 
 # Create examples list from "examples/" directory
 example_list = [["examples/" + example] for example in os.listdir("examples")]

model.py

@@ -28,9 +28,7 @@ def create_vit_model(num_classes:int=7,
 
     # Change classifier head with random seed for reproducibility
     torch.manual_seed(seed)
-    model.classifier = nn.Sequential(
-        nn.Dropout(p=0.3, inplace=True),
-        nn.Linear(in_features=768, out_features=num_classes),
-    )
-    
+    model.heads = nn.Sequential(nn.Dropout(p=0.3, inplace=True),
+                                nn.Linear(in_features=768, # keep this the same as original model
+                                          out_features=num_classes)) # update to reflect target number of classes
     return model, transforms