Update model.py

model.py

@@ -6,37 +6,49 @@ import json
 
 # Define image transformation
 transform_image = T.Compose([
-    T.Resize(224),  # Corrected to 224 to match CenterCrop
+    T.Resize(224),
     T.CenterCrop(224),
     T.ToTensor(),
     T.Normalize([0.5], [0.5])
 ])
 
-def load_image(img: str) -> torch.Tensor:
+def load_image(img_path: str) -> torch.Tensor:
     """
-    Load an image and return a tensor that can be used as an input to DINOv2.
+    Load an image and return a tensor that can be used as an input to the model.
     """
-    img = Image.open(img)
-    transformed_img = transform_image(img)[:3].unsqueeze(0)
+    img = Image.open(img_path).convert("RGB")
+    transformed_img = transform_image(img).unsqueeze(0)
     return transformed_img
 
-# Load models for inference
-dinov2_vits14 = torch.hub.load("facebookresearch/dinov2", "dinov2_vits14")
+# Load DINOv2 model for feature extraction
+dinov2_model = torch.hub.load('facebookresearch/dino:main', 'dino_vits16')
 device = torch.device('cuda' if torch.cuda.is_available() else "cpu")
-dinov2_vits14.to(device)
-dinov2_vits14.eval()  # Set the model to evaluation mode
+dinov2_model.to(device)
+dinov2_model.eval()  # Set the model to evaluation mode
 
-# Load the classifier
+# Load the pre-trained SVM classifier
 clf = joblib.load('svm_model.joblib')
 
-# Load the embeddings
-with open('all_embeddings.json', 'r') as f:
-    embeddings = json.load(f)
-
-# Predict class for a new image
-def predict(image_path):
+# Function to predict the class for a new image
+def predict(image_path: str):
+    # Load and transform the image
     new_image = load_image(image_path).to(device)
+    
+    # Extract features using DINOv2
     with torch.no_grad():
-        embedding = dinov2_vits14(new_image).cpu().numpy().reshape(1, -1)
+        features = dinov2_model(new_image)
+    
+    # Flatten features to 2D for SVM input
+    embedding = features.cpu().numpy().reshape(1, -1)
+    
+    # Predict the class using the SVM classifier
     prediction = clf.predict(embedding)
+    
     return prediction[0]
+
+# If running as a script
+if __name__ == "__main__":
+    import sys
+    image_path = sys.argv[1]  # Get image path from command line arguments
+    predicted_class = predict(image_path)
+    print("Predicted class:", predicted_class)