Spaces:
Sleeping
Sleeping
| 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) | |
| def read_root(): | |
| return {"Hello": "World"} | |
| 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)) | |
| 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)) | |