fixed missing/colliding labels and updated global model attention

app.py

@@ -4,11 +4,7 @@ import time
 import cv2
 import numpy as np
 
-# ==========================================
-# HARDWARE SCHEDULING INTERFACES (BASELINE)
-# ==========================================
 os.environ['TF_USE_LEGACY_KERAS'] = '1'
-# Initialize clean, un-throttled thread footprints for benchmarking
 os.environ["OMP_NUM_THREADS"] = "2"
 os.environ["TF_NUM_INTRAOP_THREADS"] = "2"
 os.environ["TF_NUM_INTEROP_THREADS"] = "1"
@@ -25,7 +21,6 @@ from object_detection.builders import model_builder
 
 app = FastAPI()
 
-# Enable smooth frontend cross-origin header interceptions
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -35,27 +30,18 @@ app.add_middleware(
     expose_headers=["X-Processing-Time", "X-Model-Status"]
 )
 
-# 1. Download Private Models
 HF_TOKEN = os.getenv("HF_Token")
 REPO_ID = "SaniaE/Car_Damage_Detection"
 
-print("Downloading models from Hugging Face...")
-model_dir = snapshot_download(
-    repo_id=REPO_ID, 
-    token=HF_TOKEN, 
-    local_dir="./models_data"
-)
+model_dir = snapshot_download(repo_id=REPO_ID, token=HF_TOKEN, local_dir="./models_data")
 
 PIPELINE_CONFIG = os.path.join(model_dir, "object_detection_model/pipeline.config")
 CHECKPOINT_PATH = os.path.join(model_dir, "object_detection_model/ckpt-37")
 LABEL_MAP_PATH = os.path.join(model_dir, "object_detection_model/label_map.pbtxt")
 CNN_MODEL_PATH = os.path.join(model_dir, "cnn_filter.h5")
 
-# 3. Load Models
-print("Loading CNN Filter...")
 cnn_filter = tf.keras.models.load_model(CNN_MODEL_PATH, compile=False)
 
-print("Loading Object Detection Model...")
 configs = config_util.get_configs_from_pipeline_file(PIPELINE_CONFIG)
 detection_model = model_builder.build(model_config=configs['model'], is_training=False)
 ckpt = tf.compat.v2.train.Checkpoint(model=detection_model)
@@ -69,12 +55,15 @@ def detect_fn(image):
     detections = detection_model.postprocess(prediction_dict, shapes)
     return detections
 
-def get_top_prediction(detections):
-    """Extracts the index of the most confident detection."""
+def get_top_predictions(detections, max_predictions=3):
+    """Extracts top predictions matching criteria."""
     scores = detections['detection_scores'][0].numpy()
-    if len(scores) > 0 and scores[0] > 0.4:
-        return 0, int(detections['detection_classes'][0].numpy()[0])
-    return None, None
+    classes = detections['detection_classes'][0].numpy().astype(int)
+    valid_indices = []
+    for idx in range(min(len(scores), max_predictions)):
+        if scores[idx] > 0.4:
+            valid_indices.append((idx, classes[idx]))
+    return valid_indices
 
 @app.get("/")
 def read_root():
@@ -83,8 +72,6 @@ def read_root():
 @app.post("/predict")
 async def predict(file: UploadFile = File(...)):
     start_time = time.perf_counter()
-    
-    # Read Image
     contents = await file.read()
     image_pil = Image.open(io.BytesIO(contents)).convert("RGB")
     image_np = np.array(image_pil)
@@ -114,9 +101,14 @@ async def predict(file: UploadFile = File(...)):
                 ymin, xmin, ymax, xmax = boxes[i]
                 (left, right, top, bottom) = (xmin * width, xmax * width,
                                                ymin * height, ymax * height)
+                
+                # Draw Box
                 cv2.rectangle(image_cv, (int(left), int(top)), (int(right), int(bottom)), (255, 255, 0), 2)
+                
+                # OPTIMIZATION: Bounds-safe layout rendering to prevent clipping
                 label = f"{category_index.get(classes[i] + 1, {}).get('name', 'unknown')}: {int(scores[i]*100)}%"
-                cv2.putText(image_cv, label, (int(left), int(top) - 10),
+                text_y = int(top) - 10 if int(top) - 10 > 15 else int(top) + 20
+                cv2.putText(image_cv, label, (int(left), text_y),
                             cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 0), 2)
 
     _, buffer = cv2.imencode('.jpg', image_cv)
@@ -132,7 +124,6 @@ async def predict(file: UploadFile = File(...)):
 @app.post("/explain")
 async def explain(file: UploadFile = File(...)):
     start_time = time.perf_counter()
-    
     contents = await file.read()
     image_pil = Image.open(io.BytesIO(contents)).convert("RGB")
     image_np = np.array(image_pil).astype(np.float32)
@@ -145,12 +136,13 @@ async def explain(file: UploadFile = File(...)):
         raw_scores = prediction_dict['class_predictions_with_background'][0]
         
         detections = detection_model.postprocess(prediction_dict, shapes)
-        _, top_class = get_top_prediction(detections)
+        valid_preds = get_top_predictions(detections, max_predictions=1)
         
-        if top_class is None:
+        if not valid_preds:
             elapsed_time = time.perf_counter() - start_time
             return Response(status_code=204, headers={"X-Processing-Time": f"{elapsed_time:.4f}"})
             
+        _, top_class = valid_preds[0]
         loss = tf.reduce_max(raw_scores[:, top_class])
 
     grads = tape.gradient(loss, input_tensor)
@@ -171,16 +163,11 @@ async def explain(file: UploadFile = File(...)):
     elapsed_time = time.perf_counter() - start_time
     print(f"[BENCHMARK] /explain turnaround: {elapsed_time:.4f}s")
     
-    return StreamingResponse(
-        io.BytesIO(buffer.tobytes()), 
-        media_type="image/jpeg",
-        headers={"X-Processing-Time": f"{elapsed_time:.4f}"}
-    )
+    return StreamingResponse(io.BytesIO(buffer.tobytes()), media_type="image/jpeg", headers={"X-Processing-Time": f"{elapsed_time:.4f}"})
 
 @app.post("/explain/tiled")
 async def explain_tiled(file: UploadFile = File(...)):
     start_time = time.perf_counter()
-    
     contents = await file.read()
     image_pil = Image.open(io.BytesIO(contents)).convert("RGB")
     image_np = np.array(image_pil).astype(np.float32)
@@ -193,41 +180,48 @@ async def explain_tiled(file: UploadFile = File(...)):
     
     base_image = cv2.cvtColor(image_np.astype(np.uint8), cv2.COLOR_RGB2BGR)
     h_img, w_img, _ = base_image.shape
+    
+    valid_preds = get_top_predictions(detections, max_predictions=3)
+    panels = [base_image]
+
+    for idx, (target_idx, _) in enumerate(valid_preds):
+        ymin, xmin, ymax, xmax = boxes[target_idx]
+        cv2.rectangle(base_image, (int(xmin*w_img), int(ymin*h_img)), 
+                       (int(xmax*w_img), int(ymax*h_img)), (255, 255, 0), 2)
+
+    # OPTIMIZATION: Unified single-pass Gradient Tape for all targets
+    with tf.GradientTape(persistent=True) as tape:
+        tape.watch(input_tensor)
+        image, shapes = detection_model.preprocess(input_tensor)
+        prediction_dict = detection_model.predict(image, shapes)
+        raw_scores = prediction_dict['class_predictions_with_background'][0]
         
-    for i in range(min(len(scores), 3)):
-        if scores[i] > 0.4:
-            ymin, xmin, ymax, xmax = boxes[i]
-            cv2.rectangle(base_image, (int(xmin*w_img), int(ymin*h_img)), 
-                           (int(xmax*w_img), int(ymax*h_img)), (255, 255, 0), 2)
+        losses = []
+        for _, target_class in valid_preds:
+            losses.append(tf.reduce_max(raw_scores[:, target_class]))
 
-    panels = [base_image]
+    original_bgr_raw = cv2.cvtColor(image_np.astype(np.uint8), cv2.COLOR_RGB2BGR)
+
+    for idx, loss in enumerate(losses):
+        target_class = valid_preds[idx][1]
+        grads = tape.gradient(loss, input_tensor)
+        saliency = np.max(np.abs(grads.numpy()), axis=-1)[0]
+                    
+        v_min, v_max = np.percentile(saliency, (5, 95))
+        saliency = np.clip((saliency - v_min) / (v_max - v_min + 1e-8), 0, 1)
+        heatmap = cv2.applyColorMap(np.uint8(255 * saliency), cv2.COLORMAP_JET)
+                    
+        overlay = cv2.addWeighted(original_bgr_raw.copy(), 0.6, heatmap, 0.4, 0)
+                    
+        class_name = category_index.get(target_class + 1, {}).get('name', 'unknown')
+        cv2.putText(overlay, f"Top {idx+1}: {class_name}", (10, 30), 
+                    cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2)
+        panels.append(overlay)
         
-    for i in range(3):
-        if i < len(scores) and scores[i] > 0.4:
-            target_class = classes[i]
-                        
-            with tf.GradientTape() as tape:
-                tape.watch(input_tensor)
-                image, shapes = detection_model.preprocess(input_tensor)
-                prediction_dict = detection_model.predict(image, shapes)
-                raw_scores = prediction_dict['class_predictions_with_background'][0]
-                loss = tf.reduce_max(raw_scores[:, target_class])
-                
-            grads = tape.gradient(loss, input_tensor)
-            saliency = np.max(np.abs(grads.numpy()), axis=-1)[0]
-                        
-            v_min, v_max = np.percentile(saliency, (5, 95))
-            saliency = np.clip((saliency - v_min) / (v_max - v_min + 1e-8), 0, 1)
-            heatmap = cv2.applyColorMap(np.uint8(255 * saliency), cv2.COLORMAP_JET)
-                        
-            overlay = cv2.addWeighted(cv2.cvtColor(image_np.astype(np.uint8), cv2.COLOR_RGB2BGR), 0.6, heatmap, 0.4, 0)
-                        
-            class_name = category_index.get(target_class + 1, {}).get('name', 'unknown')
-            cv2.putText(overlay, f"Top {i+1}: {class_name}", (10, 30), 
-                        cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2)
-            panels.append(overlay)
-        else:
-            panels.append(np.zeros_like(base_image))
+    del tape # Clean up persistent allocations
+
+    while len(panels) < 4:
+        panels.append(np.zeros_like(base_image))
 
     top_row = np.hstack((panels[0], panels[1]))
     bottom_row = np.hstack((panels[2], panels[3]))
@@ -237,16 +231,11 @@ async def explain_tiled(file: UploadFile = File(...)):
     elapsed_time = time.perf_counter() - start_time
     print(f"[BENCHMARK] /explain/tiled turnaround: {elapsed_time:.4f}s")
     
-    return StreamingResponse(
-        io.BytesIO(buffer.tobytes()), 
-        media_type="image/jpeg",
-        headers={"X-Processing-Time": f"{elapsed_time:.4f}"}
-    )
+    return StreamingResponse(io.BytesIO(buffer.tobytes()), media_type="image/jpeg", headers={"X-Processing-Time": f"{elapsed_time:.4f}"})
 
 @app.post("/explain/global")
 async def explain_global(file: UploadFile = File(...)):
     start_time = time.perf_counter()
-    
     contents = await file.read()
     image_pil = Image.open(io.BytesIO(contents)).convert("RGB")
     image_np = np.array(image_pil).astype(np.float32)
@@ -259,14 +248,19 @@ async def explain_global(file: UploadFile = File(...)):
         image, shapes = detection_model.preprocess(input_tensor)
         prediction_dict = detection_model.predict(image, shapes)
         raw_scores = prediction_dict['class_predictions_with_background'][0]
+        
         foreground_scores = raw_scores[:, 1:] 
-        loss = tf.reduce_sum(tf.reduce_max(foreground_scores, axis=-1))
+        # REFINEMENT: Take top-K elements to filter out background anchor static
+        top_anchor_values, _ = tf.math.top_k(tf.reduce_max(foreground_scores, axis=-1), k=20)
+        loss = tf.reduce_sum(top_anchor_values)
 
     grads = tape.gradient(loss, input_tensor)
     saliency = np.max(np.abs(grads.numpy()), axis=-1)[0]
 
-    v_min, v_max = np.percentile(saliency, (5, 95))
+    # Dilate and blur to form clean structural contours instead of fuzz
+    v_min, v_max = np.percentile(saliency, (10, 98))
     saliency = np.clip((saliency - v_min) / (v_max - v_min + 1e-8), 0, 1)
+    saliency = cv2.GaussianBlur(saliency, (9, 9), 0)
     
     heatmap = cv2.applyColorMap(np.uint8(255 * saliency), cv2.COLORMAP_JET)
     overlay = cv2.addWeighted(image_bgr, 0.6, heatmap, 0.4, 0)
@@ -278,8 +272,4 @@ async def explain_global(file: UploadFile = File(...)):
     elapsed_time = time.perf_counter() - start_time
     print(f"[BENCHMARK] /explain/global turnaround: {elapsed_time:.4f}s")
     
-    return StreamingResponse(
-        io.BytesIO(buffer.tobytes()), 
-        media_type="image/jpeg",
-        headers={"X-Processing-Time": f"{elapsed_time:.4f}"}
-    )
+    return StreamingResponse(io.BytesIO(buffer.tobytes()), media_type="image/jpeg", headers={"X-Processing-Time": f"{elapsed_time:.4f}"})