chmadnan333/plastic-api

1from fastapi import FastAPI, File, UploadFile2from fastapi.responses import JSONResponse3from ultralytics import YOLO4import cv25import numpy as np6 7app = FastAPI()8model = YOLO("weights/best.pt")9 10def is_blurry(img, threshold=50):  # 100 se 50 kiya11    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)12    return cv2.Laplacian(gray, cv2.CV_64F).var() < threshold13 14def remove_overlapping_boxes(detections, overlap_thresh=0.4):15    if len(detections) == 0:16        return detections17    kept = []18    for det_a in detections:19        duplicate = False20        a = det_a["bbox_normalized"]21        for det_b in kept:22            b = det_b["bbox_normalized"]23            ix1 = max(a["x1"], b["x1"])24            iy1 = max(a["y1"], b["y1"])25            ix2 = min(a["x2"], b["x2"])26            iy2 = min(a["y2"], b["y2"])27            inter_w = max(0, ix2 - ix1)28            inter_h = max(0, iy2 - iy1)29            inter_area = inter_w * inter_h30            area_a = (a["x2"] - a["x1"]) * (a["y2"] - a["y1"])31            area_b = (b["x2"] - b["x1"]) * (b["y2"] - b["y1"])32            union = area_a + area_b - inter_area33            iou = inter_area / union if union > 0 else 034            if iou > overlap_thresh:35                duplicate = True36                if det_a["confidence"] > det_b["confidence"]:37                    kept.remove(det_b)38                    kept.append(det_a)39                break40        if not duplicate:41            kept.append(det_a)42    return kept43 44@app.get("/")45def home():46    return {"status": "Plastic Detection API is running!"}47 48@app.post("/detect")49async def detect(file: UploadFile = File(...)):50    contents = await file.read()51    nparr = np.frombuffer(contents, np.uint8)52    img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)53 54    if is_blurry(img):55        return JSONResponse(56            {"error": "Image blurry hai, dobara click karo"},57            status_code=40058        )59 60    img_h, img_w = img.shape[:2]61 62    results = model.predict(63        img,64        conf=0.25,   # 0.556 se 0.25 kiya ✅65        iou=0.3,66        imgsz=640,   # size fix kiya ✅67        verbose=False68    )[0]69 70    detections = []71    for box in results.boxes:72        cls_name = model.names[int(box.cls[0])]73        confidence = round(float(box.conf[0]), 4)74        x1, y1, x2, y2 = box.xyxy[0].tolist()75        detections.append({76            "class": cls_name,77            "confidence": confidence,78            "bbox_normalized": {79                "x1": round(x1 / img_w, 4),80                "y1": round(y1 / img_h, 4),81                "x2": round(x2 / img_w, 4),82                "y2": round(y2 / img_h, 4),83            }84        })85 86    detections = remove_overlapping_boxes(detections, overlap_thresh=0.4)87 88    plastic_found = any(d["class"] == "plastic" for d in detections)89    return {90        "plastic_detected": plastic_found,91        "detections": detections92    }