Camspection / modelAPI.py
Cpope3's picture
Update modelAPI.py
86cc47b verified
# -*- coding: utf-8 -*-
"""
Created on Fri Feb 20 00:20:10 2026
@author: Logan
"""
import os
from fastapi import FastAPI, UploadFile, File
from fastapi.responses import FileResponse
import io
from transformers import pipeline
from PIL import Image
from ultralytics import YOLO
proc_app = FastAPI(title = "Camspection Damage Detection API") #Initialize and set the title
from fastapi.middleware.cors import CORSMiddleware
proc_app.add_middleware(
CORSMiddleware,
allow_origins=["*"], # Allows any website to call the API
allow_credentials = True,
allow_methods=["*"],
allow_headers=["*"],
)
DamageClassifier = pipeline("image-classification", model = "Cpope3/DamCat")
SeverityClassifier = pipeline("image-classification", model="Cpope3/Camspection_Model") #Create the pipeline for image classification using our checkpoint (necessary for transformerse models)
@proc_app.get("/")
def home():
return FileResponse("index.html")
@proc_app.post("/detect")
async def detect(image: UploadFile = File(...)):
scan = await image.read()
image_data = Image.open(io.BytesIO(scan)).convert("RGB")
detector = YOLO('yolov8n.pt')
carCheck = detector(image_data, conf = 0.5)
YOLOClasses = [2,5,7]
detectorClasses = carCheck[0].boxes.cls.cpu().numpy()
carFound = any(cls in YOLOClasses for cls in detectorClasses)
damCheck = DamageClassifier(image_data)
damaged = damCheck[0]['label'] == "Damaged"
if not carFound:
return {
"level": "N/A",
"score": "0%",
"debug": "No vehicle detected. Please ensure the vehicle is visible in the frame."
}
if damaged: #Classify only if car is both detected -and- found to be damaged
result = SeverityClassifier(image_data)
best_match = result[0]
alt_match = result[1]
label = best_match['label']
alt_label = alt_match['label']
best_score = best_match['score']
alt_score = alt_match['score']
if (best_score - alt_score <= 0.10):
return {
"warning": "Classification within margin of error!",
"level": label,
"alternate": alt_label,
"score": f"{best_score * 100:.2f}%",
"alternate_score": f"{alt_score * 100:.2f}%"
}
else:
return {
"warning": None,
"level": label,
"score": f"{best_score * 100:.2f}%"
}
else:
# If car is not damaged according to model
return {
"level": "Whole",
"score": f"{damCheck[0]['score'] * 100:.2f}%",
"debug": "No damage detected"
}
if __name__ == "__main__":
import uvicorn
uvicorn.run(proc_app, host="0.0.0.0", port=7860) # on the actual site