app.py

import os
import io
import time
import cv2
import numpy as np

os.environ['TF_USE_LEGACY_KERAS'] = '1'
os.environ["OMP_NUM_THREADS"] = "2"
os.environ["TF_NUM_INTRAOP_THREADS"] = "2"
os.environ["TF_NUM_INTEROP_THREADS"] = "1"

import tf_keras as keras
import tensorflow as tf
from fastapi import FastAPI, UploadFile, File
from fastapi.middleware.cors import CORSMiddleware
from PIL import Image
from huggingface_hub import snapshot_download
from fastapi.responses import StreamingResponse, Response
from object_detection.utils import label_map_util, config_util
from object_detection.builders import model_builder

app = FastAPI()

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
    expose_headers=["X-Processing-Time", "X-Model-Status"]
)

HF_TOKEN = os.getenv("HF_Token")
REPO_ID = "SaniaE/Car_Damage_Detection"

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")

cnn_filter = tf.keras.models.load_model(CNN_MODEL_PATH, compile=False)

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)
ckpt.restore(CHECKPOINT_PATH).expect_partial()
category_index = label_map_util.create_category_index_from_labelmap(LABEL_MAP_PATH)

@tf.function
def detect_fn(image):
    image, shapes = detection_model.preprocess(image)
    prediction_dict = detection_model.predict(image, shapes)
    detections = detection_model.postprocess(prediction_dict, shapes)
    return detections

def get_top_predictions(detections, max_predictions=3):
    """Extracts top predictions matching criteria."""
    scores = detections['detection_scores'][0].numpy()
    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():
    return {"status": "Model is Online", "model_repo": REPO_ID}

@app.post("/predict")
async def predict(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)
        
    image_cv = cv2.cvtColor(image_np, cv2.COLOR_RGB2BGR)
    height, width, _ = image_cv.shape

    # Step 1: CNN Filter
    img_cnn = image_pil.resize((64, 64))
    x = tf.keras.preprocessing.image.img_to_array(img_cnn)
    x = np.expand_dims(x, axis=0)
        
    cnn_pred = cnn_filter.predict(x)
    is_damage_labels = ['Clear', 'Damaged']
    status = is_damage_labels[np.argmax(cnn_pred)]

    # Step 2: Object Detection (If damaged)
    if status == 'Damaged':
        input_tensor = tf.convert_to_tensor(np.expand_dims(image_np, 0), dtype=tf.float32)
        detections = detect_fn(input_tensor)
        scores = detections['detection_scores'][0].numpy()
        classes = detections['detection_classes'][0].numpy().astype(int)
        boxes = detections['detection_boxes'][0].numpy()

        for i in range(len(scores)):
            if scores[i] > 0.4:
                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)}%"
                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)
    elapsed_time = time.perf_counter() - start_time
    print(f"[BENCHMARK] /predict turnaround: {elapsed_time:.4f}s | Status: {status}")
    
    return StreamingResponse(
        io.BytesIO(buffer.tobytes()), 
        media_type="image/jpeg",
        headers={"X-Processing-Time": f"{elapsed_time:.4f}", "X-Model-Status": status}
    )

@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)
    input_tensor = tf.convert_to_tensor(np.expand_dims(image_np, 0), dtype=tf.float32)

    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]
        
        detections = detection_model.postprocess(prediction_dict, shapes)
        valid_preds = get_top_predictions(detections, max_predictions=1)
        
        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)
    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)
    
    original_bgr = cv2.cvtColor(image_np.astype(np.uint8), cv2.COLOR_RGB2BGR)
    overlay = cv2.addWeighted(original_bgr, 0.6, heatmap, 0.4, 0)
    
    class_name = category_index.get(top_class + 1, {}).get('name', 'unknown')
    cv2.putText(overlay, f"Explaining: {class_name}", (10, 30), 
                cv2.FONT_HERSHEY_SIMPLEX, 0.8, (255, 255, 255), 2)

    _, buffer = cv2.imencode('.jpg', overlay)
    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}"})

@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)
    input_tensor = tf.convert_to_tensor(np.expand_dims(image_np, 0), dtype=tf.float32)
        
    detections = detect_fn(input_tensor)
    scores = detections['detection_scores'][0].numpy()
    classes = detections['detection_classes'][0].numpy().astype(int)
    boxes = detections['detection_boxes'][0].numpy()
    
    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]
        
        losses = []
        for _, target_class in valid_preds:
            losses.append(tf.reduce_max(raw_scores[:, target_class]))

    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)
        
    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]))
    tiled_output = np.vstack((top_row, bottom_row))

    _, buffer = cv2.imencode('.jpg', tiled_output)
    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}"})

@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)
    image_bgr = cv2.cvtColor(np.array(image_pil), cv2.COLOR_RGB2BGR)
        
    input_tensor = tf.convert_to_tensor(np.expand_dims(image_np, 0), dtype=tf.float32)

    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]
        
        foreground_scores = raw_scores[:, 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]

    # 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)
    
    cv2.putText(overlay, "Global Model Attention", (20, 40), 
                cv2.FONT_HERSHEY_SIMPLEX, 1.0, (255, 255, 0), 2)

    _, buffer = cv2.imencode('.jpg', overlay)
    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}"})