Update app.py

app.py

@@ -9,7 +9,7 @@ import matplotlib.pyplot as plt
 import numpy as np
 from sklearn.preprocessing import LabelEncoder
 from huggingface_hub import hf_hub_download
-import base64
+import torch
 from transformers import ViTImageProcessor, ViTForImageClassification
 
 # Dataset loading function with caching
@@ -17,35 +17,45 @@ from transformers import ViTImageProcessor, ViTForImageClassification
 def load_datasets():
     try:
         with st.spinner('Loading dataset...'):
-            # Load the CSV content
             original_data = pd.read_csv('CTP_Model1.csv', low_memory=False)
-            
-            # Ensure column names match the model's expectations
             original_data.columns = original_data.columns.str.strip().str.capitalize()
             return original_data
     except Exception as e:
         st.error(f"Error loading dataset: {str(e)}")
         raise e
 
-# Function definitions
-
 def load_image(image_file):
     return Image.open(image_file)
 
-
 def classify_image(image):
-    processor = ViTImageProcessor.from_pretrained("dima806/car_models_image_detection")
-    model = ViTForImageClassification.from_pretrained("dima806/car_models_image_detection")
-    
-    inputs = processor(images=image, return_tensors="pt")
-    
-    with torch.no_grad():
-        outputs = model(**inputs)
-    
-    logits = outputs.logits
-    predicted_class_idx = logits.argmax(-1).item()
+    try:
+        processor = ViTImageProcessor.from_pretrained("dima806/car_models_image_detection")
+        model = ViTForImageClassification.from_pretrained("dima806/car_models_image_detection")
+        
+        inputs = processor(images=image, return_tensors="pt")
+        
+        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 model.config.id2label[predicted_class_idx]
+    except Exception as e:
+        st.error(f"Classification error: {e}")
+        return None
+
 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:
@@ -66,7 +76,6 @@ def load_model_and_encodings():
             model_content = hf_hub_download(repo_id="EdBoy2202/car_prediction_model", filename="car_price_modelv3.pkl")
             model = joblib.load(model_content)
 
-        # Load datasets
         original_data = load_datasets()
         
         label_encoders = {}
@@ -86,19 +95,13 @@ def load_model_and_encodings():
         raise e
 
 def predict_price(model, encoders, user_input):
-    # Transform user input into model input format
     encoded_features = {feature: encoders[feature].transform([value])[0] if value in encoders[feature] else 0 
                         for feature, value in user_input.items()}
-
-    # Create a DataFrame for prediction
     input_data = pd.DataFrame([encoded_features])
-    
-    # Predict price
     predicted_price = model.predict(input_data)
     return predicted_price[0]
 
 # Streamlit App
-
 st.title("Auto Appraise")
 st.write("Capture a car image using your camera or upload an image to get its brand, model, overview, and expected price!")
 
@@ -106,8 +109,7 @@ st.write("Capture a car image using your camera or upload an image to get its br
 model, label_encoders = load_model_and_encodings()
 
 # Initialize OpenAI API key
-openai.api_key = st.secrets["GPT_TOKEN"]  # Your OpenAI API key
-HUGGINGFACE_API_KEY = st.secrets["HF_TOKEN"]  # Your Hugging Face API key
+openai.api_key = st.secrets["GPT_TOKEN"]
 
 # Camera input for taking photo
 camera_image = st.camera_input("Take a picture of the car!")
@@ -117,13 +119,23 @@ if camera_image is not None:
     st.image(image, caption='Captured Image.', use_container_width=True)
 
     # Classify the car image
-    with st.spinner('Classifying image...'):
-        car_info = classify_image(image)
+    with st.spinner('Analyzing image...'):
+        car_classifications = classify_image(image)
     
-    if car_info:
-        st.write(f"Identified Car: {car_info}")
+    if car_classifications:
+        st.subheader("Car Classification Results:")
+        for classification in car_classifications:
+            st.write(f"Model: {classification['label']}")
+            st.write(f"Confidence: {classification['probability']*100:.2f}%")
+        
+        # Use the top prediction for further processing
+        top_prediction = car_classifications[0]['label']
+        brand, model_name = top_prediction.split(' ', 1)
+        
+        st.write(f"Identified Car: {brand} {model_name}")
 
         # Find the closest match in the CSV
+        df = load_datasets()
         match = find_closest_match(df, brand, model_name)
         if match is not None:
             st.write("Closest Match Found:")
@@ -172,6 +184,6 @@ if camera_image is not None:
         else:
             st.write("No match found in the database.")
     else:
-        st.error("Could not identify the brand or model. Please try again.")
+        st.error("Could not classify the image. Please try again with a different image.")
 else:
     st.write("Please take a picture of the car to proceed.")