Update app.py

app.py

@@ -9,8 +9,7 @@ import matplotlib.pyplot as plt
 import numpy as np
 from sklearn.preprocessing import LabelEncoder
 from huggingface_hub import hf_hub_download
-import torch
-from transformers import ViTImageProcessor, ViTForImageClassification
+from transformers import pipeline
 
 # Dataset loading function with caching
 @st.cache_data
@@ -29,28 +28,14 @@ def load_image(image_file):
 
 def classify_image(image):
     try:
-        processor = ViTImageProcessor.from_pretrained("dima806/car_models_image_detection")
-        model = ViTForImageClassification.from_pretrained("dima806/car_models_image_detection")
+        # Create a pipeline for image classification
+        classifier = pipeline('image-classification', model="dima806/car_models_image_detection", device=-1)  # Use -1 for CPU, or 0 for GPU if available
         
-        inputs = processor(images=image, return_tensors="pt")
+        # Classify the image
+        results = classifier(image)
         
-        with torch.no_grad():
-            outputs = model(**inputs)
-        
-        logits = outputs.logits
-        probabilities = torch.softmax(logits, dim=-1)
-        
-        top_predictions = torch.topk(probabilities, k=3)
-        
-        predicted_classes = [
-            {
-                'label': model.config.id2label[idx.item()],
-                'probability': prob.item()
-            }
-            for idx, prob in zip(top_predictions.indices[0], top_predictions.values[0])
-        ]
-        
-        return predicted_classes
+        # Return top 5 predictions
+        return results[:5]
     
     except Exception as e:
         st.error(f"Classification error: {e}")
@@ -126,7 +111,7 @@ if camera_image is not None:
         st.subheader("Car Classification Results:")
         for classification in car_classifications:
             st.write(f"Model: {classification['label']}")
-            st.write(f"Confidence: {classification['probability']*100:.2f}%")
+            st.write(f"Confidence: {classification['score']*100:.2f}%")
         
         # Use the top prediction for further processing
         top_prediction = car_classifications[0]['label']