Upload 3 files

.env

@@ -1,6 +1,11 @@
-# Name of the first model (e.g., your original classifier)
-MODEL_1_NAME=best_new_EP382.pt
-
-# Name of the second model (for Tyre/Alloy classification)
-#MODEL_2_NAME=best_TA_377EP.pt
-MODEL_2_NAME=best_parts_EP336.pt
+# # Name of the first model (e.g., your original classifier)
+# MODEL_1_NAME=best_new_EP382.pt
+
+# # Name of the second model (for Tyre/Alloy classification)
+# #MODEL_2_NAME=best_TA_377EP.pt
+# MODEL_2_NAME=best_parts_EP336.pt
+
+
+
+PARTS_MODEL_NAME=best_parts_EP336.pt
+DAMAGE_MODEL_NAME=best_new_EP382.pt

app.py

@@ -8,6 +8,9 @@ from werkzeug.utils import secure_filename
 from ultralytics import YOLO
 from dotenv import load_dotenv
 
+# Import the new processing logic
+from processing import process_images
+
 # Load environment variables from .env file
 load_dotenv()
 
@@ -21,12 +24,13 @@ UPLOAD_FOLDER = 'static/uploads'
 MODELS_FOLDER = 'models'
 ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg'}
 
-# --- 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
+# --- Load model names from .env file ---
+# Updated names to be more descriptive
+PARTS_MODEL_NAME = os.getenv('PARTS_MODEL_NAME', 'best_parts_EP336.pt')
+DAMAGE_MODEL_NAME = os.getenv('DAMAGE_MODEL_NAME', 'best_new_EP382.pt')
 
-MODEL_1_PATH = os.path.join(MODELS_FOLDER, MODEL_1_NAME)
-MODEL_2_PATH = os.path.join(MODELS_FOLDER, MODEL_2_NAME)
+PARTS_MODEL_PATH = os.path.join(MODELS_FOLDER, PARTS_MODEL_NAME)
+DAMAGE_MODEL_PATH = os.path.join(MODELS_FOLDER, DAMAGE_MODEL_NAME)
 
 app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
 os.makedirs(app.config['UPLOAD_FOLDER'], exist_ok=True)
@@ -37,30 +41,30 @@ os.makedirs('templates', exist_ok=True)
 device = "cuda" if torch.cuda.is_available() else "cpu"
 print(f"Using device: {device}")
 
-# --- NEW: Load multiple YOLO Models ---
-model1, model2 = None, None
+# --- Load YOLO Models ---
+parts_model, damage_model = None, None
 
-# Load Model 1
+# Load Parts Model
 try:
-    if not os.path.exists(MODEL_1_PATH):
-        print(f"Warning: Model file not found at {MODEL_1_PATH}")
+    if not os.path.exists(PARTS_MODEL_PATH):
+        print(f"Warning: Parts model file not found at {PARTS_MODEL_PATH}")
     else:
-        model1 = YOLO(MODEL_1_PATH)
-        model1.to(device)
-        print(f"Successfully loaded model '{MODEL_1_NAME}' on {device}.")
+        parts_model = YOLO(PARTS_MODEL_PATH)
+        parts_model.to(device)
+        print(f"Successfully loaded parts model '{PARTS_MODEL_NAME}' on {device}.")
 except Exception as e:
-    print(f"Error loading Model 1 ({MODEL_1_NAME}): {e}")
+    print(f"Error loading Parts Model ({PARTS_MODEL_NAME}): {e}")
 
-# Load Model 2
+# Load Damage Model
 try:
-    if not os.path.exists(MODEL_2_PATH):
-        print(f"Warning: Model file not found at {MODEL_2_PATH}")
+    if not os.path.exists(DAMAGE_MODEL_PATH):
+        print(f"Warning: Damage model file not found at {DAMAGE_MODEL_PATH}")
     else:
-        model2 = YOLO(MODEL_2_PATH)
-        model2.to(device)
-        print(f"Successfully loaded model '{MODEL_2_NAME}' on {device}.")
+        damage_model = YOLO(DAMAGE_MODEL_PATH)
+        damage_model.to(device)
+        print(f"Successfully loaded damage model '{DAMAGE_MODEL_NAME}' on {device}.")
 except Exception as e:
-    print(f"Error loading Model 2 ({MODEL_2_NAME}): {e}")
+    print(f"Error loading Damage Model ({DAMAGE_MODEL_NAME}): {e}")
 
 
 def allowed_file(filename):
@@ -68,23 +72,6 @@ def allowed_file(filename):
     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."""
@@ -93,61 +80,51 @@ def home():
 @app.route('/predict', methods=['POST'])
 def predict():
     """
-    Endpoint to receive an image and run classification based on the requested model type.
+    Endpoint to receive one or more images, run the two-step prediction,
+    and return the combined results.
     """
-    # 1. --- File Validation ---
+    # 1. --- File Validation for Multiple Files ---
     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
+    
+    files = request.files.getlist('file')
+    
+    if not files or all(f.filename == '' for f in files):
+        return jsonify({"error": "No selected files"}), 400
 
-    # 2. --- Save the File Temporarily ---
-    filename = secure_filename(file.filename)
-    filepath = os.path.join(app.config['UPLOAD_FOLDER'], filename)
-    file.save(filepath)
+    saved_filepaths = []
+    
+    for file in files:
+        if file and allowed_file(file.filename):
+            filename = secure_filename(file.filename)
+            filepath = os.path.join(app.config['UPLOAD_FOLDER'], filename)
+            file.save(filepath)
+            saved_filepaths.append(filepath)
+        else:
+            # You might want to log this or inform the user about skipped files
+            print(f"Skipped invalid file: {file.filename}")
 
-    # 3. --- Perform Inference based on model_type ---
+    if not saved_filepaths:
+        return jsonify({"error": "No valid files were uploaded. Allowed types: png, jpg, jpeg"}), 400
+        
+    # 2. --- Perform Inference ---
     try:
-        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
+        # Pass the models and file paths to the processing function
+        results = process_images(parts_model, damage_model, saved_filepaths)
+        return jsonify(results)
 
     except Exception as e:
-        return jsonify({"error": f"An error occurred during inference: {str(e)}"}), 500
+        # Log the full error for debugging
+        print(f"An error occurred during processing: {e}")
+        import traceback
+        traceback.print_exc()
+        return jsonify({"error": f"An error occurred during processing: {str(e)}"}), 500
     finally:
-        # 4. --- Cleanup ---
-        if os.path.exists(filepath):
-            os.remove(filepath)
+        # 3. --- Cleanup ---
+        for filepath in saved_filepaths:
+            if os.path.exists(filepath):
+                os.remove(filepath)
 
 if __name__ == '__main__':
+    # Setting debug=False is recommended for production
     app.run(host='0.0.0.0', port=7860, debug=True)

processing.py

@@ -0,0 +1,83 @@
+# processing.py
+
+import os
+from ultralytics import YOLO
+
+# --- Configuration ---
+# These are the specific parts that require a subsequent damage check.
+DAMAGE_CHECK_PARTS = {
+    'driver_front_side',
+    'driver_rear_side',
+    'passenger_front_side',
+    'passenger_rear_side',
+}
+
+def run_single_inference(model, filepath):
+    """
+    Helper function to run inference for a single model and format the result.
+    """
+    if model is None:
+        return None  # Return None if the model isn't loaded
+
+    results = model(filepath, verbose=False) # verbose=False to keep logs clean
+    result = results[0]
+    
+    # Check if it's a classification model with probabilities
+    if result.probs is not None:
+        probs = result.probs
+        top1_index = probs.top1
+        top1_confidence = float(probs.top1conf)
+        class_name = model.names[top1_index]
+    else: # Fallback for detection models or if probs are not available
+        # Assuming the top prediction is what we need
+        top1_index = result.boxes.cls[0].int() if len(result.boxes) > 0 else 0
+        top1_confidence = float(result.boxes.conf[0]) if len(result.boxes) > 0 else 0.0
+        class_name = model.names[top1_index] if len(result.boxes) > 0 else "unknown"
+
+    return {
+        "class": class_name,
+        "confidence": round(top1_confidence, 4)
+    }
+
+def process_images(parts_model, damage_model, image_paths):
+    """
+    Processes a list of images.
+    1. Runs the 'parts_model' on every image.
+    2. If the detected part is in DAMAGE_CHECK_PARTS, it then runs the 'damage_model'.
+    3. Otherwise, the damage status defaults to 'correct'.
+    """
+    if parts_model is None or damage_model is None:
+        raise RuntimeError("One or more models are not loaded. Check server logs.")
+
+    final_results = []
+
+    for filepath in image_paths:
+        filename = os.path.basename(filepath)
+        print(f"Processing {filename}...")
+
+        # 1. First, predict the part
+        part_prediction = run_single_inference(parts_model, filepath)
+        predicted_part = part_prediction.get("class") if part_prediction else "unknown"
+
+        damage_prediction = None
+        # 2. Conditionally predict the damage
+        if predicted_part in DAMAGE_CHECK_PARTS:
+            print(f"  -> Part '{predicted_part}' requires damage check. Running damage model...")
+            damage_prediction = run_single_inference(damage_model, filepath)
+        else:
+            print(f"  -> Part '{predicted_part}' does not require damage check. Defaulting to 'correct'.")
+            # 3. For other parts, default to 'correct'
+            damage_prediction = {
+                "class": "correct",
+                "confidence": 1.0,
+                "note": "Result by default, not by model inference."
+            }
+
+        # Assemble the final result for this image
+        final_results.append({
+            "filename": filename,
+            "part_prediction": part_prediction,
+            "damage_prediction": damage_prediction
+        })
+
+    return final_results