Update app.py

app.py

@@ -10,6 +10,7 @@ import numpy as np
 from sklearn.preprocessing import LabelEncoder
 from huggingface_hub import hf_hub_download
 import base64
+from transformers import ViTImageProcessor, ViTForImageClassification
 
 # Dataset loading function with caching
 @st.cache_data
@@ -33,37 +34,18 @@ def load_image(image_file):
 
 
 def classify_image(image):
-    # Convert PIL Image to bytes
-    img_byte_arr = BytesIO()
-    image.save(img_byte_arr, format='PNG')
-    img_byte_arr = img_byte_arr.getvalue()
-
-    # Encode image to base64
-    encoded_image = base64.b64encode(img_byte_arr).decode('ascii')
-
-    headers = {
-        "Authorization": f"Bearer {HUGGINGFACE_API_KEY}",
-        "Content-Type": "application/json"
-    }
+    processor = ViTImageProcessor.from_pretrained("dima806/car_models_image_detection")
+    model = ViTForImageClassification.from_pretrained("dima806/car_models_image_detection")
     
-    payload = {
-        "inputs": encoded_image
-    }
-
-    try:
-        response = requests.post(
-            'https://api-inference.huggingface.co/models/dima806/car_models_image_detection',
-            headers=headers,
-            json=payload
-        )
-        response.raise_for_status()  # Raises an HTTPError for bad responses
-        return response.json()
-    except requests.exceptions.RequestException as e:
-        st.error(f"Image classification failed: {e}")
-        if response.text:
-            st.error(f"API Response: {response.text}")
-        return None
-
+    inputs = processor(images=image, return_tensors="pt")
+    
+    with torch.no_grad():
+        outputs = model(**inputs)
+    
+    logits = outputs.logits
+    predicted_class_idx = logits.argmax(-1).item()
+    
+    return model.config.id2label[predicted_class_idx]
 def find_closest_match(df, brand, model):
     match = df[(df['Make'].str.contains(brand, case=False)) & (df['Model'].str.contains(model, case=False))]
     if not match.empty:
@@ -139,11 +121,7 @@ if camera_image is not None:
         car_info = classify_image(image)
     
     if car_info:
-        brand = car_info.get('brand', None)  # Adjust according to the response structure
-        model_name = car_info.get('model', None)
-        
-        if brand and model_name:
-            st.write(f"Identified Car: {brand} {model_name}")
+        st.write(f"Identified Car: {car_info}")
 
             # Find the closest match in the CSV
             match = find_closest_match(df, brand, model_name)