Update app.py

app.py

@@ -2,7 +2,7 @@ import os
 import cv2
 import numpy as np
 import base64
-from flask import Flask, render_template_string, request, redirect, flash
+from flask import Flask, render_template_string, request, redirect, flash, url_for
 import roboflow
 import torch
 from collections import Counter
@@ -33,8 +33,6 @@ except Exception as e:
 # --- YOLOv5 Pretrained Model for Persons & Cars ---
 try:
     yolov5_model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True)
-    # Filter YOLO detections to only include persons and cars.
-    YOLO_FILTER_CLASSES = {"person", "car"}
     print("YOLOv5 model loaded successfully.")
 except Exception as e:
     print("Error loading YOLOv5 model:", e)
@@ -71,7 +69,7 @@ def custom_nms(preds, iou_threshold=0.3):
             filtered_preds.append(pred)
     return filtered_preds
 
-def process_image(image_path):
+def process_image(image_path, object_type):
     # Load image
     image = cv2.imread(image_path)
     if image is None:
@@ -79,124 +77,131 @@ def process_image(image_path):
         return None, "Could not read the image."
 
     img_height, img_width = image.shape[:2]
-    detection_info = []  # List to hold all detection results for display
+    detection_info = []  # List to hold detection results for display
 
-    # --- (a) Roboflow Box Detection & Measurement ---
-    if box_model is None:
-        print("DEBUG: Roboflow model is not initialized.")
-        return None, "Roboflow model is not available."
-    try:
-        results = box_model.predict(image_path, confidence=50, overlap=30).json()
-        predictions = results.get("predictions", [])
-    except Exception as e:
-        print("DEBUG: Error during Roboflow prediction:", e)
-        return None, "Error during Roboflow prediction."
-
-    processed_preds = []
-    for prediction in predictions:
+    if object_type == "box":
+        # --- Roboflow Box Detection & Measurement ---
+        if box_model is None:
+            print("DEBUG: Roboflow model is not initialized.")
+            return None, "Roboflow model is not available."
         try:
-            x, y, width, height = prediction["x"], prediction["y"], prediction["width"], prediction["height"]
-            x1 = int(round(x - width / 2))
-            y1 = int(round(y - height / 2))
-            x2 = int(round(x + width / 2))
-            y2 = int(round(y + height / 2))
-            # Clamp coordinates to image dimensions
-            x1 = max(0, min(x1, img_width - 1))
-            y1 = max(0, min(y1, img_height - 1))
-            x2 = max(0, min(x2, img_width - 1))
-            y2 = max(0, min(y2, img_height - 1))
-            processed_preds.append({
-                "box": (x1, y1, x2, y2),
-                "class": prediction["class"],
-                "confidence": prediction["confidence"]
-            })
+            results = box_model.predict(image_path, confidence=50, overlap=30).json()
+            predictions = results.get("predictions", [])
         except Exception as e:
-            print("DEBUG: Error processing a prediction:", e)
-            continue
+            print("DEBUG: Error during Roboflow prediction:", e)
+            return None, "Error during Roboflow prediction."
 
-    box_detections = custom_nms(processed_preds, iou_threshold=0.3)
+        processed_preds = []
+        for prediction in predictions:
+            try:
+                x, y, width, height = prediction["x"], prediction["y"], prediction["width"], prediction["height"]
+                x1 = int(round(x - width / 2))
+                y1 = int(round(y - height / 2))
+                x2 = int(round(x + width / 2))
+                y2 = int(round(y + height / 2))
+                # Clamp coordinates to image dimensions
+                x1 = max(0, min(x1, img_width - 1))
+                y1 = max(0, min(y1, img_height - 1))
+                x2 = max(0, min(x2, img_width - 1))
+                y2 = max(0, min(y2, img_height - 1))
+                processed_preds.append({
+                    "box": (x1, y1, x2, y2),
+                    "class": prediction["class"],
+                    "confidence": prediction["confidence"]
+                })
+            except Exception as e:
+                print("DEBUG: Error processing a prediction:", e)
+                continue
 
-    # Detect ArUco marker for measurement (only applicable for boxes)
-    marker_real_width_cm = 10.0  # The marker is 10cm x 10cm
-    try:
-        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
-        aruco_dict = cv2.aruco.getPredefinedDictionary(cv2.aruco.DICT_6X6_250)
-        aruco_params = cv2.aruco.DetectorParameters()
-        corners, ids, _ = cv2.aruco.detectMarkers(gray, aruco_dict, parameters=aruco_params)
-        if ids is not None and len(corners) > 0:
-            marker_corners = corners[0].reshape((4, 2))
-            cv2.aruco.drawDetectedMarkers(image, corners, ids)
-            marker_width_pixels = np.linalg.norm(marker_corners[0] - marker_corners[1])
-            marker_height_pixels = np.linalg.norm(marker_corners[1] - marker_corners[2])
-            marker_pixel_size = (marker_width_pixels + marker_height_pixels) / 2.0
-            conversion_factor = marker_real_width_cm / marker_pixel_size
-        else:
+        box_detections = custom_nms(processed_preds, iou_threshold=0.3)
+
+        # Detect ArUco marker for measurement (only applicable for boxes)
+        marker_real_width_cm = 10.0  # The marker is 10cm x 10cm
+        try:
+            gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
+            aruco_dict = cv2.aruco.getPredefinedDictionary(cv2.aruco.DICT_6X6_250)
+            aruco_params = cv2.aruco.DetectorParameters()
+            corners, ids, _ = cv2.aruco.detectMarkers(gray, aruco_dict, parameters=aruco_params)
+            if ids is not None and len(corners) > 0:
+                marker_corners = corners[0].reshape((4, 2))
+                cv2.aruco.drawDetectedMarkers(image, corners, ids)
+                marker_width_pixels = np.linalg.norm(marker_corners[0] - marker_corners[1])
+                marker_height_pixels = np.linalg.norm(marker_corners[1] - marker_corners[2])
+                marker_pixel_size = (marker_width_pixels + marker_height_pixels) / 2.0
+                conversion_factor = marker_real_width_cm / marker_pixel_size
+            else:
+                conversion_factor = None
+        except Exception as e:
+            print("DEBUG: Error during ArUco detection:", e)
             conversion_factor = None
-    except Exception as e:
-        print("DEBUG: Error during ArUco detection:", e)
-        conversion_factor = None
 
-    # Draw box detections and record measurement info (only for boxes)
-    for pred in box_detections:
-        x1, y1, x2, y2 = pred["box"]
-        label = pred["class"]
-        confidence = pred["confidence"]
-        cv2.rectangle(image, (x1, y1), (x2, y2), (0, 255, 0), 2)
-        if conversion_factor is not None:
-            box_width_pixels = x2 - x1
-            box_height_pixels = y2 - y1
-            box_width_cm = box_width_pixels * conversion_factor
-            box_height_cm = box_height_pixels * conversion_factor
-            size_text = f"{box_width_cm:.1f}x{box_height_cm:.1f} cm"
-            detection_info.append({
-                "class": label,
-                "confidence": f"{confidence:.2f}",
-                "width_cm": f"{box_width_cm:.1f}",
-                "height_cm": f"{box_height_cm:.1f}"
-            })
-        else:
-            size_text = ""
-            detection_info.append({
-                "class": label,
-                "confidence": f"{confidence:.2f}",
-                "width_cm": "N/A",
-                "height_cm": "N/A"
-            })
-        text = f"{label} ({confidence:.2f}) {size_text}"
-        (text_width, text_height), baseline = cv2.getTextSize(text, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)
-        cv2.rectangle(image, (x1, y1 - text_height - baseline - 5), (x1 + text_width, y1 - 5), (0, 255, 0), -1)
-        cv2.putText(image, text, (x1, y1 - 5 - baseline), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 1)
+        # Draw box detections and record measurement info
+        for pred in box_detections:
+            x1, y1, x2, y2 = pred["box"]
+            label = pred["class"]
+            confidence = pred["confidence"]
+            cv2.rectangle(image, (x1, y1), (x2, y2), (0, 255, 0), 2)
+            if conversion_factor is not None:
+                box_width_pixels = x2 - x1
+                box_height_pixels = y2 - y1
+                box_width_cm = box_width_pixels * conversion_factor
+                box_height_cm = box_height_pixels * conversion_factor
+                size_text = f"{box_width_cm:.1f}x{box_height_cm:.1f} cm"
+                detection_info.append({
+                    "class": label,
+                    "confidence": f"{confidence:.2f}",
+                    "width_cm": f"{box_width_cm:.1f}",
+                    "height_cm": f"{box_height_cm:.1f}"
+                })
+            else:
+                size_text = ""
+                detection_info.append({
+                    "class": label,
+                    "confidence": f"{confidence:.2f}",
+                    "width_cm": "N/A",
+                    "height_cm": "N/A"
+                })
+            text = f"{label} ({confidence:.2f}) {size_text}"
+            (text_width, text_height), baseline = cv2.getTextSize(text, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)
+            cv2.rectangle(image, (x1, y1 - text_height - baseline - 5), (x1 + text_width, y1 - 5), (0, 255, 0), -1)
+            cv2.putText(image, text, (x1, y1 - 5 - baseline), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 1)
 
-    # --- (b) YOLOv5 for Persons & Cars ---
-    if yolov5_model is None:
-        print("DEBUG: YOLOv5 model is not initialized.")
+    elif object_type in {"person", "car"}:
+        # --- YOLOv5 for Persons or Cars (filtering for the selected type) ---
+        if yolov5_model is None:
+            print("DEBUG: YOLOv5 model is not initialized.")
+            return None, "YOLOv5 model is not available."
+        else:
+            try:
+                img_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+                yolo_results = yolov5_model(img_rgb)
+                df = yolo_results.pandas().xyxy[0]
+                for _, row in df.iterrows():
+                    if row['name'] == object_type:
+                        xmin = int(row['xmin'])
+                        ymin = int(row['ymin'])
+                        xmax = int(row['xmax'])
+                        ymax = int(row['ymax'])
+                        conf = row['confidence']
+                        label = row['name']
+                        cv2.rectangle(image, (xmin, ymin), (xmax, ymax), (255, 0, 0), 2)
+                        text = f"{label} ({conf:.2f})"
+                        (text_width, text_height), baseline = cv2.getTextSize(text, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)
+                        cv2.rectangle(image, (xmin, ymin - text_height - baseline - 5), (xmin + text_width, ymin - 5), (255, 0, 0), -1)
+                        cv2.putText(image, text, (xmin, ymin - 5 - baseline), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 1)
+                        detection_info.append({
+                            "class": label,
+                            "confidence": f"{conf:.2f}",
+                            "width_cm": "N/A",
+                            "height_cm": "N/A"
+                        })
+            except Exception as e:
+                print("DEBUG: Error during YOLOv5 inference:", e)
+                return None, "Error during YOLOv5 inference."
     else:
-        try:
-            img_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
-            yolo_results = yolov5_model(img_rgb)
-            df = yolo_results.pandas().xyxy[0]
-            for _, row in df.iterrows():
-                if row['name'] in YOLO_FILTER_CLASSES:
-                    xmin = int(row['xmin'])
-                    ymin = int(row['ymin'])
-                    xmax = int(row['xmax'])
-                    ymax = int(row['ymax'])
-                    conf = row['confidence']
-                    label = row['name']
-                    cv2.rectangle(image, (xmin, ymin), (xmax, ymax), (255, 0, 0), 2)
-                    text = f"{label} ({conf:.2f})"
-                    (text_width, text_height), baseline = cv2.getTextSize(text, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)
-                    cv2.rectangle(image, (xmin, ymin - text_height - baseline - 5), (xmin + text_width, ymin - 5), (255, 0, 0), -1)
-                    cv2.putText(image, text, (xmin, ymin - 5 - baseline), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 1)
-                    detection_info.append({
-                        "class": label,
-                        "confidence": f"{conf:.2f}",
-                        "width_cm": "N/A",
-                        "height_cm": "N/A"
-                    })
-        except Exception as e:
-            print("DEBUG: Error during YOLOv5 inference:", e)
-            return None, "Error during YOLOv5 inference."
+        err_msg = f"Unrecognized object type: {object_type}"
+        print("DEBUG:", err_msg)
+        return None, err_msg
 
     # --- Build Top Summary Text ---
     detection_counts = Counter(det["class"] for det in detection_info)
@@ -209,40 +214,73 @@ def process_image(image_path):
     return image, detection_info
 
 #########################################
-# 3. Flask Routes
+# 3. HTML Templates
 #########################################
 
-@app.route('/', methods=['GET', 'POST'])
-def index():
-    image_data = None
-    detection_info = None
-    if request.method == 'POST':
-        if 'file' not in request.files:
-            flash('No file part')
-            return redirect(request.url)
-        file = request.files['file']
-        if file.filename == '':
-            flash('No selected file')
-            return redirect(request.url)
-        upload_path = "uploaded.jpg"
-        try:
-            file.save(upload_path)
-        except Exception as e:
-            print("DEBUG: Error saving uploaded file:", e)
-            flash("Error saving uploaded file.")
-            return redirect(request.url)
-        processed_image, detection_info = process_image(upload_path)
-        if processed_image is None:
-            flash("Error Processing Image: " + detection_info)
-        else:
-            retval, buffer = cv2.imencode('.jpg', processed_image)
-            image_data = base64.b64encode(buffer).decode('utf-8')
-        try:
-            os.remove(upload_path)
-        except Exception as e:
-            print("DEBUG: Error removing uploaded file:", e)
-    return render_template_string('''
- <!DOCTYPE html>
+# --- Landing Page Template (New) ---
+landing_template = '''
+<!DOCTYPE html>
+<html lang="en">
+<head>
+  <meta charset="UTF-8">
+  <meta name="viewport" content="width=device-width, initial-scale=1.0">
+  <title>MathLens</title>
+  <link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/bootstrap@5.3.0/dist/css/bootstrap.min.css">
+  <style>
+    @import url('https://fonts.googleapis.com/css2?family=Share+Tech+Mono&display=swap');
+    body {
+      background-color: #fff;
+      color: #000;
+      font-family: "Share Tech Mono", monospace;
+      text-align: center;
+      display: flex;
+      flex-direction: column;
+      justify-content: center;
+      align-items: center;
+      min-height: 100vh;
+      padding: 20px;
+      overflow: auto;
+    }
+    h1 {
+      font-size: 2.5rem;
+      margin-bottom: 20px;
+    }
+    p {
+      font-size: 1.5rem;
+      margin-bottom: 40px;
+    }
+    .btn {
+      display: inline-block;
+      margin: 10px;
+      padding: 15px 30px;
+      font-size: 1.2rem;
+      text-decoration: none;
+      border: 2px solid #000;
+      color: #000;
+      cursor: pointer;
+      transition: background-color 0.3s, color 0.3s;
+    }
+    .btn:hover {
+      background-color: #000;
+      color: #fff;
+    }
+  </style>
+</head>
+<body>
+  <h1>MathLens</h1>
+  <p>What do you want to count?</p>
+  <div>
+    <a href="{{ url_for('upload') }}?object_type=person" class="btn">People</a>
+    <a href="{{ url_for('upload') }}?object_type=car" class="btn">Cars</a>
+    <a href="{{ url_for('upload') }}?object_type=box" class="btn">Boxes</a>
+  </div>
+</body>
+</html>
+'''
+
+# --- Upload & Result Page Template (Original Design with Home Button and Progress Bar Loading Overlay) ---
+upload_template = '''
+<!DOCTYPE html>
 <html lang="en">
 <head>
     <meta charset="UTF-8">
@@ -291,6 +329,54 @@ def index():
             color: #000;
             font-family: "Share Tech Mono", monospace;
         }
+        .home-btn {
+            display: inline-block;
+            margin: 10px auto 20px;
+            padding: 10px 20px;
+            border: 2px solid #000;
+            color: #000;
+            text-decoration: none;
+            font-family: "Share Tech Mono", monospace;
+            cursor: pointer;
+            transition: background-color 0.3s, color 0.3s;
+        }
+        .home-btn:hover {
+            background-color: #000;
+            color: #fff;
+        }
+        /* Progress bar overlay styles */
+        #loading {
+            position: fixed;
+            top: 0;
+            left: 0;
+            width: 100%;
+            height: 100%;
+            background: rgba(255,255,255,0.95);
+            display: none;
+            flex-direction: column;
+            align-items: center;
+            justify-content: center;
+            z-index: 9999;
+        }
+        #progressContainer {
+            width: 80%;
+            max-width: 600px;
+            background: #ddd;
+            border: 2px solid #000;
+            border-radius: 5px;
+            overflow: hidden;
+            margin-bottom: 20px;
+        }
+        #progressBar {
+            width: 0%;
+            height: 30px;
+            background: #000;
+            transition: width 0.2s;
+        }
+        #progressText {
+            font-size: 1.2rem;
+            font-weight: bold;
+        }
         .content-wrapper {
             display: flex;
             flex-direction: row;
@@ -339,9 +425,21 @@ def index():
     </style>
 </head>
 <body>
+    <!-- Home button -->
+    <a href="{{ url_for('landing') }}" class="home-btn">Home</a>
+    <!-- Progress bar loading overlay -->
+    <div id="loading">
+        <div id="progressContainer">
+            <div id="progressBar"></div>
+        </div>
+        <div id="progressText">Initializing... 🔄</div>
+    </div>
     <div class="typing-effect" id="typing"></div>
-    <form method="post" enctype="multipart/form-data">
+    <!-- onsubmit calls the showLoading() function -->
+    <form method="post" enctype="multipart/form-data" onsubmit="showLoading()">
         <input type="file" name="file" accept="image/*" required>
+        <!-- Hidden field to pass the selected object type -->
+        <input type="hidden" name="object_type" value="{{ object_type }}">
         <button type="submit">Analyze Image</button>
     </form>
     {% if image_data or detection_info %}
@@ -373,6 +471,40 @@ def index():
     {% endif %}
     <div class="footer">&copy; 2024 MathLens AI Detection App. All rights reserved.</div>
     <script>
+        // Function to show the loading overlay with a progress bar and cool phrases
+        function showLoading() {
+            document.getElementById("loading").style.display = "flex";
+            let progress = 0;
+            const progressBar = document.getElementById("progressBar");
+            const progressText = document.getElementById("progressText");
+            // Array of phrases with their upper progress limits
+            const phrases = [
+                {limit: 20, text: "Writing scripts... ✍️"},
+                {limit: 40, text: "Calculating formulas... 🔢"},
+                {limit: 60, text: "Mixing up magic... ✨"},
+                {limit: 80, text: "Almost there... 🚀"},
+                {limit: 100, text: "Finalizing details... ✅"}
+            ];
+            // Update progress every 50ms (simulation)
+            const interval = setInterval(() => {
+                // Increase progress by a random value (to add some variability)
+                progress += Math.floor(Math.random() * 3) + 1;
+                if (progress > 100) progress = 100;
+                progressBar.style.width = progress + "%";
+                // Update phrase based on current progress
+                for (let i = 0; i < phrases.length; i++) {
+                    if (progress <= phrases[i].limit) {
+                        progressText.textContent = phrases[i].text;
+                        break;
+                    }
+                }
+                // If progress reaches 100, clear the interval
+                if (progress >= 100) {
+                    clearInterval(interval);
+                }
+            }, 50);
+        }
+        // Existing typing effect code
         const textArray = ["MathLens", "Smart Counting with Maths"];
         let textIndex = 0;
         let charIndex = 0;
@@ -401,15 +533,65 @@ def index():
     </script>
 </body>
 </html>
+'''
 
+#########################################
+# 4. Flask Routes
+#########################################
 
-    ''', image_data=image_data, detection_info=detection_info)
+# Landing page: displays only the title and three buttons
+@app.route('/')
+def landing():
+    return render_template_string(landing_template)
+
+# Upload page: uses the original upload page design and processes the image.
+@app.route('/upload', methods=['GET', 'POST'])
+def upload():
+    if request.method == 'GET':
+        # Get the selected object type from the query parameter
+        object_type = request.args.get('object_type', '').lower()
+        if object_type not in {"person", "car", "box"}:
+            flash("Please select a valid object type.")
+            return redirect(url_for('landing'))
+        return render_template_string(upload_template, object_type=object_type)
+    else:
+        # POST: process the uploaded image
+        if 'file' not in request.files:
+            flash('No file part')
+            return redirect(request.url)
+        file = request.files['file']
+        if file.filename == '':
+            flash('No selected file')
+            return redirect(request.url)
+        object_type = request.form.get('object_type', '').lower()
+        if object_type not in {"person", "car", "box"}:
+            flash("Invalid object type selected.")
+            return redirect(url_for('landing'))
+        upload_path = "uploaded.jpg"
+        try:
+            file.save(upload_path)
+        except Exception as e:
+            print("DEBUG: Error saving uploaded file:", e)
+            flash("Error saving uploaded file.")
+            return redirect(request.url)
+        processed_image, detection_info = process_image(upload_path, object_type)
+        if processed_image is None:
+            flash("Error Processing Image: " + detection_info)
+            return redirect(request.url)
+        else:
+            retval, buffer = cv2.imencode('.jpg', processed_image)
+            image_data = base64.b64encode(buffer).decode('utf-8')
+        try:
+            os.remove(upload_path)
+        except Exception as e:
+            print("DEBUG: Error removing uploaded file:", e)
+        return render_template_string(upload_template, object_type=object_type,
+                                      image_data=image_data, detection_info=detection_info)
 
 #########################################
-# Run the App
+# 5. Run the App
 #########################################
 
 if __name__ == '__main__':
-    # Ensure the app runs on 0.0.0.0 and port 7860 for Hugging Face Spaces.
+    # Runs on host 0.0.0.0 and port 7860 (adjust as needed)
     app.run(host="0.0.0.0", port=7860)
-