Update app.py

app.py

@@ -14,7 +14,7 @@ app = Flask(__name__)
 app.secret_key = 'your_secret_key'  # Replace with a secure secret key
 
 # Global dictionary to hold job progress and results
-jobs = {}  # jobs[job_id] = {"progress": int, "result": {…}}
+jobs = {}  # jobs[job_id] = {"progress": int, "result": {...}}
 
 #########################################
 # 1. Initialize the Models
@@ -75,11 +75,12 @@ def custom_nms(preds, iou_threshold=0.3):
             filtered_preds.append(pred)
     return filtered_preds
 
-# Note the added multiplier parameter.
+# The process_image function now merges the box detection code (from your first prompt)
+# with the ArUco marker detection (and size measurement) exactly as in the code you provided now.
 def process_image(job_id, image_path, object_type, multiplier):
     try:
         jobs[job_id]['progress'] = 10
-        # Load image
+        # Load a fresh copy of the original image
         image = cv2.imread(image_path)
         if image is None:
             jobs[job_id]['progress'] = 100
@@ -97,18 +98,36 @@ def process_image(job_id, image_path, object_type, multiplier):
                 jobs[job_id]['progress'] = 100
                 jobs[job_id]['result'] = {"error": "Roboflow model not available."}
                 return
-            try:
-                results = box_model.predict(image_path, confidence=50, overlap=30).json()
-                predictions = results.get("predictions", [])
-            except Exception as e:
-                jobs[job_id]['progress'] = 100
-                jobs[job_id]['result'] = {"error": "Error during Roboflow prediction."}
-                return
 
+            # --- BOX DETECTION (from your 1st prompt) ---
+            # Upscale the image if it is small to help detect small boxes.
+            scale_factor = 1
+            if img_width < 1000 or img_height < 1000:
+                scale_factor = 2
+
+            if scale_factor > 1:
+                upscaled_image = cv2.resize(image, None, fx=scale_factor, fy=scale_factor, interpolation=cv2.INTER_LINEAR)
+                temp_path = "upscaled.jpg"
+                cv2.imwrite(temp_path, upscaled_image)
+                results = box_model.predict(temp_path, confidence=30, overlap=20).json()
+            else:
+                results = box_model.predict(image_path, confidence=30, overlap=20).json()
+
+            predictions = results.get("predictions", [])
             processed_preds = []
             for prediction in predictions:
                 try:
-                    x, y, width, height = prediction["x"], prediction["y"], prediction["width"], prediction["height"]
+                    if scale_factor > 1:
+                        x = prediction["x"] / scale_factor
+                        y = prediction["y"] / scale_factor
+                        width = prediction["width"] / scale_factor
+                        height = prediction["height"] / scale_factor
+                    else:
+                        x = prediction["x"]
+                        y = prediction["y"]
+                        width = prediction["width"]
+                        height = prediction["height"]
+
                     x1 = int(round(x - width / 2))
                     y1 = int(round(y - height / 2))
                     x2 = int(round(x + width / 2))
@@ -125,10 +144,11 @@ def process_image(job_id, image_path, object_type, multiplier):
                 except Exception as e:
                     continue
 
-            box_detections = custom_nms(processed_preds, iou_threshold=0.3)
+            # Use a higher NMS IoU threshold to avoid merging valid nearby detections.
+            box_detections = custom_nms(processed_preds, iou_threshold=0.7)
             jobs[job_id]['progress'] = 60
 
-            # (Optional) Detect ArUco marker for measurement
+            # --- ARUCO MARKER DETECTION (from the code you provided now) ---
             marker_real_width_cm = 5.0  
             try:
                 gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
@@ -153,6 +173,7 @@ def process_image(job_id, image_path, object_type, multiplier):
             except Exception as e:
                 conversion_factor = None
 
+            # --- Draw boxes and compute measurements ---
             for pred in box_detections:
                 x1, y1, x2, y2 = pred["box"]
                 label = pred["class"]
@@ -180,6 +201,7 @@ def process_image(job_id, image_path, object_type, multiplier):
                 (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)
+
         elif object_type in {"person", "car"}:
             if yolov5_model is None:
                 jobs[job_id]['progress'] = 100
@@ -213,7 +235,7 @@ def process_image(job_id, image_path, object_type, multiplier):
                 jobs[job_id]['result'] = {"error": "Error during YOLOv5 inference."}
                 return
 
-        # Build summary text (drawn on image)
+        # Build summary text drawn on the image
         detection_counts = Counter(det["class"] for det in detection_info)
         if detection_counts:
             top_text = ", ".join(f"{cls}: {count}" for cls, count in detection_counts.items())
@@ -500,7 +522,6 @@ def analyze():
     object_type = request.form.get('object_type', '').lower()
     if object_type not in {"person", "car", "box"}:
         return jsonify({"error": "Invalid object type."}), 400
-    # Retrieve the multiplier sent from the client
     try:
         multiplier = int(request.form.get('multiplier', 1))
     except ValueError: