Spaces:
Paused
Paused
| 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 --- | |
| def home(): | |
| return f""" | |
| <html> | |
| <head> | |
| <title>AI Inspection Station</title> | |
| <style> | |
| body {{ font-family: sans-serif; background: #0F172A; color: white; padding: 40px; }} | |
| .card {{ background: #1E293B; border-radius: 12px; padding: 30px; margin-bottom: 25px; border: 1px solid #334155; }} | |
| .btn {{ background: #3B82F6; color: white; padding: 12px 24px; border: none; border-radius: 6px; cursor: pointer; font-weight: bold; }} | |
| .report-box {{ | |
| white-space: pre-wrap; | |
| background: #0F172A; | |
| padding: 20px; | |
| border-radius: 8px; | |
| color: #38BDF8; | |
| font-family: monospace; | |
| border-left: 4px solid #3B82F6; | |
| margin-top: 15px; | |
| line-height: 1.6; | |
| }} | |
| img {{ max-width: 100%; border-radius: 8px; margin-top: 15px; border: 1px solid #334155; }} | |
| .grid {{ display: grid; grid-template-columns: 1fr 1fr; gap: 20px; }} | |
| input, select {{ width: 100%; padding: 12px; margin: 10px 0; background: #0F172A; color: white; border: 1px solid #334155; border-radius: 6px; }} | |
| </style> | |
| </head> | |
| <body style="max-width: 1100px; margin: auto;"> | |
| <h1>🛡 AI Inventory & Master Inspection</h1> | |
| <div class="card"> | |
| <h2>Tool 1: Precision Image Analysis</h2> | |
| <div class="grid"> | |
| <div> | |
| <select id="mode" onchange="document.getElementById('f2box').style.display = (this.value == 'count' ? 'none' : 'block');"> | |
| <option value="damage">Damage Detect (Compare)</option> | |
| <option value="count">Inventory Log (Single Scan)</option> | |
| </select> | |
| <input type="file" id="f1"> | |
| <div id="f2box"><input type="file" id="f2"></div> | |
| </div> | |
| <div> | |
| <input type="text" id="prompt" value="chair. sofa. table. painting. flower."> | |
| <button class="btn" onclick="runProcess('single')">Execute Scan</button> | |
| </div> | |
| </div> | |
| <div id="single_res" style="display:none"></div> | |
| </div> | |
| <div class="card" style="border-top: 4px solid #A855F7;"> | |
| <h2>Tool 2: Dual-Video Logic Pipeline</h2> | |
| <div class="grid"> | |
| <div><label>Video: Before</label><input type="file" id="v1"></div> | |
| <div><label>Video: After</label><input type="file" id="v2"></div> | |
| </div> | |
| <label>Neural Keywords</label> | |
| <input type="text" id="v_prompt" value="chair. sofa. table. broken. fallen."> | |
| <button class="btn" style="background: #A855F7;" onclick="runProcess('video')">📋 Generate Logic Report</button> | |
| <div id="video_res" style="display:none"></div> | |
| </div> | |
| <script> | |
| async function runProcess(type) {{ | |
| const resDiv = document.getElementById(type == 'single' ? 'single_res' : 'video_res'); | |
| resDiv.style.display = 'block'; | |
| resDiv.innerHTML = "⏳ AI is processing. Verifying counts..."; | |
| const fd = new FormData(); | |
| if(type == 'single') {{ | |
| const m = document.getElementById('mode').value; | |
| fd.append('prompt', document.getElementById('prompt').value); | |
| fd.append('image1', document.getElementById('f1').files[0]); | |
| if(m=='damage') fd.append('image2', document.getElementById('f2').files[0]); | |
| const r = await fetch(m=='damage' ? '/damage-photo' : '/count-objects', {{method:'POST', body:fd}}); | |
| const d = await r.json(); | |
| resDiv.innerHTML = `<img src="data:image/jpeg;base64,${{d.image}}"><div class="report-box">${{d.report}}</div>`; | |
| }} else {{ | |
| fd.append('v1', document.getElementById('v1').files[0]); | |
| fd.append('v2', document.getElementById('v2').files[0]); | |
| fd.append('v_prompt', document.getElementById('v_prompt').value); | |
| const r = await fetch('/video-report', {{method:'POST', body:fd}}); | |
| const d = await r.json(); | |
| resDiv.innerHTML = `<div class="report-box">${{d.report}}</div><div class="grid"><img src="data:image/jpeg;base64,${{d.before_img}}"><img src="data:image/jpeg;base64,${{d.after_img}}"></div>`; | |
| }} | |
| }} | |
| </script> | |
| </body> | |
| </html> | |
| """ | |
| # --- 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 | |
| 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}"} | |
| 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)}"} | |
| 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) | |