Spaces:
Sleeping
Sleeping
| from fastapi import FastAPI, APIRouter, UploadFile, File, Request | |
| from fastapi.staticfiles import StaticFiles | |
| import shutil | |
| import os | |
| import cv2 | |
| import base64 | |
| from ultralytics import YOLO | |
| # ─── Load Models | |
| person_model = YOLO("yolov8n.pt") | |
| statue_model = YOLO("best.pt") | |
| # ─── App Setup | |
| app = FastAPI() | |
| UPLOAD_FOLDER = "/tmp/uploads" | |
| os.makedirs(UPLOAD_FOLDER, exist_ok=True) | |
| app.mount("/uploads", StaticFiles(directory=UPLOAD_FOLDER), name="uploads") | |
| CONF_THRESHOLD = 0.30 | |
| router = APIRouter() | |
| def draw_label(image, label, x1, y1, x2, y2, box_color, text_color): | |
| font = cv2.FONT_HERSHEY_SIMPLEX | |
| font_scale = 0.6 | |
| thickness = 2 | |
| (tw, th), _ = cv2.getTextSize(label, font, font_scale, thickness) | |
| if y1 - th - 8 >= 0: | |
| label_y1 = y1 - th - 8 | |
| label_y2 = y1 | |
| text_y = y1 - 5 | |
| else: | |
| label_y1 = y1 | |
| label_y2 = y1 + th + 8 | |
| text_y = y1 + th + 3 | |
| cv2.rectangle(image, (x1, label_y1), (x1 + tw + 4, label_y2), box_color, -1) | |
| cv2.putText(image, label, (x1 + 2, text_y), font, font_scale, text_color, thickness) | |
| def root(): | |
| return {"message": "AI API is running 🚀"} | |
| async def predict_image( | |
| request: Request, | |
| file: UploadFile = File(...) | |
| ): | |
| safe_filename = file.filename.replace(" ", "_") | |
| file_path = os.path.join(UPLOAD_FOLDER, safe_filename) | |
| with open(file_path, "wb") as buffer: | |
| shutil.copyfileobj(file.file, buffer) | |
| image = cv2.imread(file_path) | |
| if image is None: | |
| return {"error": "Invalid image"} | |
| detections = [] | |
| person_count = 0 | |
| statue_count = 0 | |
| # ── Person Detection | |
| person_results = person_model(file_path) | |
| for box in person_results[0].boxes: | |
| cls_id = int(box.cls) | |
| if cls_id != 0: | |
| continue | |
| conf = float(box.conf) | |
| if conf < CONF_THRESHOLD: | |
| continue | |
| x1, y1, x2, y2 = map(int, box.xyxy[0]) | |
| cv2.rectangle(image, (x1, y1), (x2, y2), (0, 255, 0), 2) | |
| draw_label(image, f"Person {conf:.2f}", x1, y1, x2, y2, | |
| box_color=(0, 255, 0), text_color=(0, 0, 0)) | |
| detections.append({ | |
| "type": "person", | |
| "name": "Person", | |
| "confidence": round(conf, 4), | |
| "bbox": [x1, y1, x2, y2] | |
| }) | |
| person_count += 1 | |
| # ── Statue Detection | |
| statue_results = statue_model(file_path) | |
| for box in statue_results[0].boxes: | |
| conf = float(box.conf) | |
| if conf < CONF_THRESHOLD: | |
| continue | |
| cls_id = int(box.cls) | |
| statue_name = statue_results[0].names[cls_id] | |
| x1, y1, x2, y2 = map(int, box.xyxy[0]) | |
| cv2.rectangle(image, (x1, y1), (x2, y2), (0, 0, 255), 2) | |
| draw_label(image, f"{statue_name} {conf:.2f}", x1, y1, x2, y2, | |
| box_color=(0, 0, 255), text_color=(255, 255, 255)) | |
| detections.append({ | |
| "type": "statue", | |
| "name": statue_name, | |
| "confidence": round(conf, 4), | |
| "bbox": [x1, y1, x2, y2] | |
| }) | |
| statue_count += 1 | |
| # ── Save Output Image | |
| output_filename = f"output_{safe_filename}" | |
| output_path = os.path.join(UPLOAD_FOLDER, output_filename) | |
| cv2.imwrite(output_path, image) | |
| with open(output_path, "rb") as img_file: | |
| image_base64 = base64.b64encode(img_file.read()).decode("utf-8") | |
| image_url = f"{request.base_url}uploads/{output_filename}" | |
| return { | |
| "total_count": len(detections), | |
| "persons": person_count, | |
| "statues": statue_count, | |
| "output_image_url": image_url, | |
| "output_image_base64": f"data:image/jpeg;base64,{image_base64}", | |
| "detections": detections | |
| } | |
| app.include_router(router) | |