import os import io import cv2 import torch import base64 import numpy as np import tempfile from PIL import Image from fastapi import FastAPI, UploadFile, File, Form from fastapi.responses import JSONResponse, HTMLResponse from transformers import AutoProcessor, AutoModelForZeroShotObjectDetection from torchvision.ops import nms # <--- Added for merging duplicate boxes app = FastAPI() DEVICE = "cpu" DINO_ID = "IDEA-Research/grounding-dino-base" try: processor = AutoProcessor.from_pretrained(DINO_ID) model = AutoModelForZeroShotObjectDetection.from_pretrained(DINO_ID, low_cpu_mem_usage=True).to(DEVICE).eval() except Exception as e: print(f"Error: {e}") def img_to_b64(img_array): _, buffer = cv2.imencode('.jpg', img_array) return base64.b64encode(buffer).decode('utf-8') # --- NMS Logic for Clean Counting --- def clean_detections(boxes, scores, labels, iou_threshold=0.3): if len(boxes) == 0: return [], [], [] # Use Torch's NMS to remove overlapping boxes keep = nms(torch.tensor(boxes), torch.tensor(scores), iou_threshold) final_boxes = boxes[keep] final_labels = [labels[i] for i in keep] final_scores = scores[keep] return final_boxes, final_labels, final_scores def calculate_iou(box1, box2): xA, yA, xB, yB = max(box1[0], box2[0]), max(box1[1], box2[1]), min(box1[2], box2[2]), min(box1[3], box2[3]) inter = max(0, xB - xA) * max(0, yB - yA) area1 = (box1[2] - box1[0]) * (box1[3] - box1[1]) area2 = (box2[2] - box2[0]) * (box2[3] - box2[1]) return inter / float(area1 + area2 - inter + 1e-6) # --- DASHBOARD UI --- @app.get("/", response_class=HTMLResponse) def home(): return f""" AI Inspection Station

🛡 AI Inventory & Master Inspection

Tool 1: Precision Image Analysis

Tool 2: Dual-Video Logic Pipeline

Video: Before

Video: After

Neural Keywords

""" # --- BACKEND (ENHANCED WITH NMS) --- def run_dino(pil_img, prompt, threshold=0.15): inputs = processor(images=pil_img, text=prompt, return_tensors="pt").to(DEVICE) with torch.no_grad(): out = model(**inputs) res = processor.post_process_grounded_object_detection(out, inputs.input_ids, target_sizes=[pil_img.size[::-1]], threshold=threshold)[0] # 🧹 NMS Cleaning: Combine overlapping boxes boxes, labels, scores = clean_detections(res["boxes"], res["scores"], res["labels"]) return boxes, labels @app.post("/count-objects") async def count_objects(image1: UploadFile=File(...), prompt: str=Form(...)): img = Image.open(io.BytesIO(await image1.read())).convert("RGB") boxes, labels = run_dino(img, prompt) counts = {} for lbl in labels: if len(lbl.strip()) > 1: counts[lbl] = counts.get(lbl, 0) + 1 rep = "\\n".join([f"• {k}: {v}" for k,v in counts.items()]) return {"image": img_to_b64(cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)), "report": f"[ LOGS ]\\n{rep}"} @app.post("/damage-photo") async def damage_photo(image1: UploadFile=File(...), image2: UploadFile=File(...), prompt: str=Form(...)): b = Image.open(io.BytesIO(await image1.read())).convert("RGB") a = Image.open(io.BytesIO(await image2.read())).convert("RGB") b_bx, _ = run_dino(b, prompt) a_bx, _ = run_dino(a, prompt) new_dmg = [bx for bx in a_bx if not any(calculate_iou(bx, bb) > 0.15 for bb in b_bx)] img_out = cv2.cvtColor(np.array(a), cv2.COLOR_RGB2BGR) for bx in new_dmg: cv2.rectangle(img_out, (int(bx[0]), int(bx[1])), (int(bx[2]), int(bx[3])), (0,0,255), 3) return {"image": img_to_b64(img_out), "report": f"ALERTS:\\n• Abnormalities Found: {len(new_dmg)}"} @app.post("/video-report") async def video_report(v1: UploadFile=File(...), v2: UploadFile=File(...), v_prompt: str=Form(...)): def get_f(vf): with tempfile.NamedTemporaryFile(delete=False, suffix=".mp4") as t: t.write(vf); p = t.name cap = cv2.VideoCapture(p); ret, f = cap.read(); cap.release(); os.remove(p) return f if ret else None f1 = get_f(await v1.read()); f2 = get_f(await v2.read()) def get_data(frame): pil = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)) boxes, labels = run_dino(pil, v_prompt) cnts = {} for l in labels: if len(l.strip())>1: cnts[l] = cnts.get(l,0)+1 return cnts, boxes, labels c1, b1, l1 = get_data(f1); c2, b2, l2 = get_data(f2) # Missing Analysis missing = [] for i, bx1 in enumerate(b1): if not any(calculate_iou(bx1, bx2) > 0.3 for bx2 in b2): missing.append(l1[i]) m_dict = {} for mi in missing: m_dict[mi] = m_dict.get(mi, 0) + 1 # Format Output before_txt = "\\n".join([f"• {k}: {v}" for k,v in c1.items()]) if c1 else "None" after_txt = "\\n".join([f"• {k}: {v}" for k,v in c2.items()]) if c2 else "None" miss_txt = ", ".join([f"{k} ({v})" for k,v in m_dict.items()]) if m_dict else "None" # New Damage Detection for Red Boxes new_dmg_boxes = [bx for bx in b2 if not any(calculate_iou(bx, bb) > 0.15 for bb in b1)] report = f"""[ STAGE COMPARISON REPORT ] ----------------------------------------- REFERENCE (BEFORE): {before_txt} CURRENT (AFTER): {after_txt} SUMMARY: • Missing Items: {miss_txt} • New Issues: {len(new_dmg_boxes)} -----------------------------------------""" f2_out = f2.copy() for bx in new_dmg_boxes: cv2.rectangle(f2_out, (int(bx[0]), int(bx[1])), (int(bx[2]), int(bx[3])), (0,0,255), 3) return {"report": report, "before_img": img_to_b64(f1), "after_img": img_to_b64(f2_out)} if __name__ == "__main__": import uvicorn uvicorn.run(app, host="0.0.0.0", port=7860)