Upload 2 files

.env

@@ -1 +1,5 @@
- MODEL_NAME=best_new_EP382.pt
+# Name of the first model (e.g., your original classifier)
+MODEL_1_NAME=best_88E.pt
+
+# Name of the second model (for Tyre/Alloy classification)
+MODEL_2_NAME=best_TA_377EP.pt

app.py

@@ -18,41 +18,73 @@ CORS(app)
 
 # --- Configuration ---
 UPLOAD_FOLDER = 'static/uploads'
-MODELS_FOLDER = 'models' # New folder for models
+MODELS_FOLDER = 'models'
 ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg'}
 
-# Load model name from .env file, with a fallback default
-MODEL_NAME = os.getenv('MODEL_NAME', 'best.pt')
-MODEL_PATH = os.path.join(MODELS_FOLDER, MODEL_NAME)
+# --- NEW: Load model names from .env file, with fallback defaults ---
+MODEL_1_NAME = os.getenv('MODEL_1_NAME', 'best.pt')
+MODEL_2_NAME = os.getenv('MODEL_2_NAME', 'tyre_alloy.pt') # New model for Tyre/Alloy
+
+MODEL_1_PATH = os.path.join(MODELS_FOLDER, MODEL_1_NAME)
+MODEL_2_PATH = os.path.join(MODELS_FOLDER, MODEL_2_NAME)
 
 app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
 os.makedirs(app.config['UPLOAD_FOLDER'], exist_ok=True)
-os.makedirs(MODELS_FOLDER, exist_ok=True) # Ensure models folder exists
-os.makedirs('templates', exist_ok=True) # Ensure templates folder exists
+os.makedirs(MODELS_FOLDER, exist_ok=True)
+os.makedirs('templates', exist_ok=True)
 
-# --- Determine Device and Load YOLO Model ---
-# Use CUDA if available, otherwise use CPU
+# --- Determine Device ---
 device = "cuda" if torch.cuda.is_available() else "cpu"
 print(f"Using device: {device}")
 
-# Load the model once when the application starts for efficiency.
-model = None
+# --- NEW: Load multiple YOLO Models ---
+model1, model2 = None, None
+
+# Load Model 1
 try:
-    if not os.path.exists(MODEL_PATH):
-        print(f"Error: Model file not found at {MODEL_PATH}")
-        print("Please make sure the model file exists and the MODEL_NAME in your .env file is correct.")
+    if not os.path.exists(MODEL_1_PATH):
+        print(f"Warning: Model file not found at {MODEL_1_PATH}")
     else:
-        model = YOLO(MODEL_PATH)
-        model.to(device) # Move model to the selected device
-        print(f"Successfully loaded model '{MODEL_NAME}' on {device}.")
+        model1 = YOLO(MODEL_1_PATH)
+        model1.to(device)
+        print(f"Successfully loaded model '{MODEL_1_NAME}' on {device}.")
 except Exception as e:
-    print(f"Error loading YOLO model: {e}")
+    print(f"Error loading Model 1 ({MODEL_1_NAME}): {e}")
+
+# Load Model 2
+try:
+    if not os.path.exists(MODEL_2_PATH):
+        print(f"Warning: Model file not found at {MODEL_2_PATH}")
+    else:
+        model2 = YOLO(MODEL_2_PATH)
+        model2.to(device)
+        print(f"Successfully loaded model '{MODEL_2_NAME}' on {device}.")
+except Exception as e:
+    print(f"Error loading Model 2 ({MODEL_2_NAME}): {e}")
+
 
 def allowed_file(filename):
     """Checks if a file's extension is in the ALLOWED_EXTENSIONS set."""
     return '.' in filename and \
            filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS
 
+def run_inference(model, filepath):
+    """Helper function to run inference and format the result."""
+    if model is None:
+        return None # Return None if the model isn't loaded
+
+    results = model(filepath)
+    result = results[0]
+    probs = result.probs
+    top1_index = probs.top1
+    top1_confidence = float(probs.top1conf)
+    class_name = model.names[top1_index]
+    
+    return {
+        "class": class_name,
+        "confidence": top1_confidence
+    }
+
 @app.route('/')
 def home():
     """Serve the main HTML page."""
@@ -61,59 +93,59 @@ def home():
 @app.route('/predict', methods=['POST'])
 def predict():
     """
-    Endpoint to receive an image, run YOLO classification, and return the single best prediction.
+    Endpoint to receive an image and run classification based on the requested model type.
     """
-    if model is None:
-        return jsonify({"error": "Model could not be loaded. Please check server logs."}), 500
-        
     # 1. --- File Validation ---
     if 'file' not in request.files:
         return jsonify({"error": "No file part in the request"}), 400
-
     file = request.files['file']
     if file.filename == '':
         return jsonify({"error": "No selected file"}), 400
-
     if not file or not allowed_file(file.filename):
         return jsonify({"error": "File type not allowed"}), 400
 
+    # --- NEW: Get the model type from the form data ---
+    model_type = request.form.get('model_type', 'model1') # default to model1
+
     # 2. --- Save the File Temporarily ---
     filename = secure_filename(file.filename)
     filepath = os.path.join(app.config['UPLOAD_FOLDER'], filename)
     file.save(filepath)
 
-    # 3. --- Perform Inference ---
+    # 3. --- Perform Inference based on model_type ---
     try:
-        # Run the YOLO model on the uploaded image. The model is already on the correct device.
-        results = model(filepath)
-
-        # 4. --- Process Results to Get ONLY the Top Prediction ---
-        # Get the first result object from the list
-        result = results[0]
-        
-        # Access the probabilities object
-        probs = result.probs
-        
-        # Get the index and confidence of the top prediction
-        top1_index = probs.top1
-        top1_confidence = float(probs.top1conf) # Convert tensor to Python float
-        
-        # Get the class name from the model's 'names' dictionary
-        class_name = model.names[top1_index]
-        
-        # Create the final prediction object
-        prediction = {
-            "class": class_name,
-            "confidence": top1_confidence
-        }
-
-        # Return the single prediction object as JSON
-        return jsonify(prediction)
+        if model_type == 'model1':
+            if model1 is None:
+                return jsonify({"error": f"Model '{MODEL_1_NAME}' is not loaded. Check server logs."}), 500
+            prediction = run_inference(model1, filepath)
+            return jsonify(prediction)
+
+        elif model_type == 'model2':
+            if model2 is None:
+                return jsonify({"error": f"Model '{MODEL_2_NAME}' is not loaded. Check server logs."}), 500
+            prediction = run_inference(model2, filepath)
+            return jsonify(prediction)
+
+        elif model_type == 'combined':
+            if model1 is None or model2 is None:
+                 return jsonify({"error": "One or more models required for combined mode are not loaded. Check server logs."}), 500
+            
+            pred1 = run_inference(model1, filepath)
+            pred2 = run_inference(model2, filepath)
+            
+            combined_prediction = {
+                "model1_result": pred1,
+                "model2_result": pred2
+            }
+            return jsonify(combined_prediction)
+
+        else:
+            return jsonify({"error": "Invalid model type specified"}), 400
 
     except Exception as e:
         return jsonify({"error": f"An error occurred during inference: {str(e)}"}), 500
     finally:
-        # 5. --- Cleanup ---
+        # 4. --- Cleanup ---
         if os.path.exists(filepath):
             os.remove(filepath)