Tharan-J commited on
Commit
35d6185
·
1 Parent(s): 70349b9
Files changed (2) hide show
  1. model.py +5 -1
  2. requirements.txt +1 -0
model.py CHANGED
@@ -3,8 +3,12 @@ import gradio as gr
3
  from transformers import ConvNextForImageClassification, AutoImageProcessor
4
  from PIL import Image
5
 
 
 
 
 
 
6
  # Load model configuration and weights manually
7
- model = ConvNextForImageClassification.from_pretrained("facebook/convnext-base-224") # Load the base model
8
  model.load_state_dict(torch.load("convnext_base_finetuned.pth", map_location="cpu")) # Load your finetuned weights
9
  model.eval()
10
 
 
3
  from transformers import ConvNextForImageClassification, AutoImageProcessor
4
  from PIL import Image
5
 
6
+ model = ConvNextForImageClassification.from_pretrained("facebook/convnext-base-224")
7
+
8
+ # Redefine classifier for 23 classes
9
+ model.classifier = torch.nn.Linear(in_features=1024, out_features=23)
10
+
11
  # Load model configuration and weights manually
 
12
  model.load_state_dict(torch.load("convnext_base_finetuned.pth", map_location="cpu")) # Load your finetuned weights
13
  model.eval()
14
 
requirements.txt CHANGED
@@ -1,3 +1,4 @@
1
  torch==1.13.1
2
  transformers==4.27.0
3
  gradio==3.0.1
 
 
1
  torch==1.13.1
2
  transformers==4.27.0
3
  gradio==3.0.1
4
+ numpy<2