Pujan-Dev's picture
test
f9003ec
import fastapi
import uvicorn
from fastapi import FastAPI, UploadFile, File, HTTPException
from typing import Dict
import cv2
import numpy as np
from main import NumberPlateOCR
app = FastAPI()
# Initialize the OCR pipeline
pipeline = NumberPlateOCR(use_yolo=True, verbose=False)
@app.get("/")
def read_root():
return {"Hello": "World"}
@app.post("/ocr-only/")
async def ocr_only(file: UploadFile = File(...)) -> Dict:
"""
Perform OCR on an uploaded image without YOLO detection.
Args:
file: Uploaded image file.
Returns:
JSON response with OCR results.
"""
try:
# Read image from upload
contents = await file.read()
np_img = np.frombuffer(contents, np.uint8)
image = cv2.imdecode(np_img, cv2.IMREAD_COLOR)
if image is None:
raise HTTPException(status_code=400, detail="Invalid image file")
# Process the image directly (skip YOLO)
results = pipeline.process_from_plate_image(image, show_visualization=False)
return {
"lines": results["lines"],
"multiline_text": results["multiline_text"],
"singleline_text": results["singleline_text"],
"num_lines": results["num_lines"],
"total_chars": results["total_chars"],
"confidence_stats": results["confidence_stats"]
}
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
@app.post("/yolo-ocr/")
async def yolo_ocr(file: UploadFile = File(...)) -> Dict:
"""
Perform YOLO-based plate detection and OCR on an uploaded image.
Args:
file: Uploaded image file.
Returns:
JSON response with OCR results.
"""
try:
# Read image from upload
contents = await file.read()
np_img = np.frombuffer(contents, np.uint8)
image_path = "uploaded_image.jpg"
# Save the uploaded image temporarily
with open(image_path, "wb") as f:
f.write(contents)
# Process the image using the pipeline
results = pipeline.process_image(image_path, save_contours=False, show_visualization=False)
# Format the response
response = {
"image_path": results["image_path"],
"num_plates": results["num_plates"],
"plates": []
}
for plate in results["plates"]:
response["plates"].append({
"lines": plate["lines"],
"multiline_text": plate["multiline_text"],
"singleline_text": plate["singleline_text"],
"num_lines": plate["num_lines"],
"total_chars": plate["total_chars"],
"confidence_stats": plate["confidence_stats"]
})
return response
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))