|
|
from fastapi import FastAPI, UploadFile, File, Form
|
|
|
from fastapi.responses import JSONResponse
|
|
|
from ultralytics import YOLO
|
|
|
import numpy as np
|
|
|
from PIL import Image
|
|
|
from datetime import datetime, timedelta
|
|
|
import psycopg2
|
|
|
import io
|
|
|
|
|
|
app = FastAPI()
|
|
|
|
|
|
model = YOLO('./best v3(100).pt')
|
|
|
|
|
|
'''def init_connection():
|
|
|
return psycopg2.connect(
|
|
|
host="localhost",
|
|
|
database="water_meter",
|
|
|
user="postgres",
|
|
|
password="16ecr193@kec"
|
|
|
)
|
|
|
|
|
|
def insert_detection(conn, detection_date, water_meter_column, detected_values_str):
|
|
|
with conn.cursor() as cur:
|
|
|
created_time = datetime.now().strftime("%H:%M:%S")
|
|
|
|
|
|
cur.execute(f"""
|
|
|
INSERT INTO water_meter_readings (detection_date, {water_meter_column}, created_time)
|
|
|
VALUES (%s, %s, %s)
|
|
|
ON CONFLICT (detection_date) DO UPDATE
|
|
|
SET {water_meter_column} = EXCLUDED.{water_meter_column},
|
|
|
created_time = EXCLUDED.created_time
|
|
|
""", (detection_date, int(detected_values_str), created_time))
|
|
|
|
|
|
meter_number = water_meter_column.split('_')[-1]
|
|
|
difference_column = f'difference_{meter_number}'
|
|
|
|
|
|
yesterday = detection_date - timedelta(days=1)
|
|
|
|
|
|
cur.execute(f"""
|
|
|
SELECT {water_meter_column}
|
|
|
FROM water_meter_readings
|
|
|
WHERE detection_date = %s
|
|
|
""", (yesterday,))
|
|
|
yesterday_result = cur.fetchone()
|
|
|
|
|
|
if yesterday_result and yesterday_result[0] is not None:
|
|
|
difference = int(detected_values_str) - yesterday_result[0]
|
|
|
else:
|
|
|
difference = None
|
|
|
|
|
|
if difference is not None:
|
|
|
cur.execute(f"""
|
|
|
UPDATE water_meter_readings
|
|
|
SET {difference_column} = %s
|
|
|
WHERE detection_date = %s
|
|
|
""", (difference, detection_date))
|
|
|
else:
|
|
|
cur.execute(f"""
|
|
|
UPDATE water_meter_readings
|
|
|
SET {difference_column} = NULL
|
|
|
WHERE detection_date = %s
|
|
|
""", (detection_date,))
|
|
|
|
|
|
conn.commit()'''
|
|
|
|
|
|
@app.get("/")
|
|
|
def read_root():
|
|
|
return {"message": "Welcome to Water Meter Detection API"}
|
|
|
|
|
|
@app.post("/detect/")
|
|
|
async def detect_water_meter(
|
|
|
|
|
|
|
|
|
file: UploadFile = File(...)
|
|
|
):
|
|
|
image_data = await file.read()
|
|
|
image = Image.open(io.BytesIO(image_data))
|
|
|
img_array = np.array(image)
|
|
|
|
|
|
results = model.predict(source=img_array, conf=0.25, imgsz=640)
|
|
|
|
|
|
detected_numbers = []
|
|
|
for result in results:
|
|
|
for box in result.boxes:
|
|
|
label = result.names[box.cls[0].item()]
|
|
|
x1, _, _, _ = box.xyxy[0].tolist()
|
|
|
|
|
|
if label.isdigit():
|
|
|
detected_numbers.append((x1, label))
|
|
|
|
|
|
detected_numbers.sort(key=lambda x: x[0])
|
|
|
sorted_numbers = [num for _, num in detected_numbers]
|
|
|
detected_values_str = ''.join(sorted_numbers)
|
|
|
|
|
|
'''if detected_values_str:
|
|
|
water_meter_column = f'water_meter_{water_meter}'
|
|
|
detection_date_obj = datetime.strptime(detection_date, "%Y-%m-%d").date()
|
|
|
|
|
|
conn = init_connection()
|
|
|
insert_detection(conn, detection_date_obj, water_meter_column, detected_values_str)
|
|
|
conn.close()'''
|
|
|
|
|
|
return JSONResponse(
|
|
|
content={
|
|
|
"status": "success",
|
|
|
"detected_values": detected_values_str if sorted_numbers else "No numeric values detected."
|
|
|
}
|
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
|
import uvicorn
|
|
|
uvicorn.run(app, host="127.0.0.1", port=8000)
|
|
|
|
