server.py - develop - linariell/Honeysuckle

server.py
182 строки7 KB
Polina Vitkovskaya
fixing OCR
20 май 2026, 04:47
20 май 2026, 04:47bae406c
import cv2
import numpy as np
import base64
from fastapi import FastAPI, UploadFile, File, Form
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse, FileResponse
from scipy.signal import find_peaks
import easyocr
import re

reader = easyocr.Reader(['en'])

app = FastAPI()

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

def image_to_base64(img):
    _, buffer = cv2.imencode('.jpg', img)
    return base64.b64encode(buffer).decode('utf-8')

# --- КЛАССИЧЕСКИЙ EASYOCR (Оптимизированный по скорости) ---
def fast_extract_label(image):
    h, w = image.shape[:2]
    
    # Сжимаем фото до 1000px по большей стороне
    # Это ускорит работу EasyOCR с 3 минут до ~1 секунды
    scale = 1000.0 / max(h, w)
    if scale < 1.0:
        img_small = cv2.resize(image, (0, 0), fx=scale, fy=scale)
    else:
        img_small = image
        scale = 1.0

    # Запускаем классический EasyOCR на всем (уменьшенном) кадре
    # detail=1 возвращает: [ [координаты], 'текст', уверенность ]
    ocr_results = reader.readtext(img_small, detail=1)
    
    regex_pattern = r'([A-Za-z]{1,4})[-\s_]*([0-9OIl]{2})[-\s_]*([0-9OIl]{3})'
    
    best_text = "Образец"
    best_bbox = None
    
    # Ищем текст, подходящий под наш формат
    for bbox, text, conf in ocr_results:
        text_clean = text.replace('I', '1').replace('l', '1')
        match = re.search(regex_pattern, text_clean)
        if match:
            part1 = match.group(1)
            part2 = match.group(2).replace('O', '0').replace('o', '0')
            part3 = match.group(3).replace('O', '0').replace('o', '0')
            best_text = f"{part1}-{part2}-{part3}"
            best_bbox = bbox
            break
            
    # Если не нашли идеального совпадения, берем самый длинный кусок текста как резерв
    if best_text == "Образец" and ocr_results:
        longest = max(ocr_results, key=lambda x: len(x[1]))
        clean_fallback = re.sub(r'[^a-zA-Z0-9\-]', '', longest[1])
        if len(clean_fallback) >= 5:
            best_text = clean_fallback[:15]
            best_bbox = longest[0]

    # Вырезаем этикетку для превью, используя координаты от самого EasyOCR
    label_crop = np.zeros((50, 150, 3), dtype=np.uint8)
    if best_bbox is not None:
        pts = np.array(best_bbox, dtype=np.int32)
        x, y, cw, ch = cv2.boundingRect(pts)
        
        # Возвращаемся в масштаб оригинальной большой картинки
        X = int(x / scale)
        Y = int(y / scale)
        CW = int(cw / scale)
        CH = int(ch / scale)
        
        # Делаем отступ, чтобы захватить поля вокруг текста
        pad = 30
        X1 = max(0, X - pad)
        Y1 = max(0, Y - pad)
        X2 = min(w, X + CW + pad)
        Y2 = min(h, Y + CH + pad)
        
        if X2 > X1 and Y2 > Y1:
            label_crop = image[Y1:Y2, X1:X2]

    return best_text, label_crop

def auto_align(image):
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    edges = cv2.Canny(gray, 70, 200)
    lines = cv2.HoughLines(edges, 1, np.pi/180, 180)
    if lines is None: return image
    angles = [ (theta - np.pi/2) * 180 / np.pi for line in lines[:20] for rho, theta in [line[0]] if abs((theta - np.pi/2) * 180 / np.pi) < 15]
    if not angles: return image
    angle = np.median(angles)
    h, w = image.shape[:2]
    M = cv2.getRotationMatrix2D((w//2, h//2), angle, 1)
    return cv2.warpAffine(image, M, (w, h), borderMode=cv2.BORDER_REPLICATE)

def segment_leaf(image):
    hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
    mask = cv2.inRange(hsv, np.array([30, 40, 30]), np.array([90, 255, 255]))
    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (15, 15))
    return cv2.morphologyEx(cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel), cv2.MORPH_OPEN, cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5)))

def smart_crop(image, mask):
    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    if not contours: return image, mask, (0,0,0,0)
    largest = max(contours, key=cv2.contourArea)
    x, y, w, h = cv2.boundingRect(largest)
    return image[max(0, y-h//5):min(image.shape[0], y+h+h//5), max(0, x-w//5):min(image.shape[1], x+w+w//5)], \
           mask[max(0, y-h//5):min(mask.shape[0], y+h+h//5), max(0, x-w//5):min(mask.shape[1], x+w+w//5)], (x, y, w, h)

def compute_auto_area_ratio(roi_image, roi_mask):
    bg_mask = cv2.erode(cv2.bitwise_not(roi_mask), cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (15, 15)))
    gray = cv2.cvtColor(roi_image, cv2.COLOR_BGR2GRAY)
    thresh = cv2.adaptiveThreshold(cv2.createCLAHE(clipLimit=3.0).apply(gray), 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY_INV, 21, 5)
    roi_lines = cv2.bitwise_and(thresh, thresh, mask=bg_mask)
    
    proj_x = np.sum(roi_lines, axis=0)
    peaks, _ = find_peaks(np.correlate(proj_x - np.mean(proj_x), proj_x - np.mean(proj_x), mode='full')[len(roi_lines[0]):], distance=5, prominence=100)
    
    if len(peaks) > 1:
        step = np.median(np.diff(peaks))
        return 1.0 / (step ** 2), roi_lines
    return None, roi_lines

@app.post("/analyze")
async def analyze(file: UploadFile = File(...), rotation: int = Form(0)):
    contents = await file.read()
    img = cv2.imdecode(np.frombuffer(contents, np.uint8), cv2.IMREAD_COLOR)
    
    # Применяем поворот, если пользователь повернул картинку в интерфейсе
    if rotation == 90:
        img = cv2.rotate(img, cv2.ROTATE_90_CLOCKWISE)
    elif rotation == 180:
        img = cv2.rotate(img, cv2.ROTATE_180)
    elif rotation == 270:
        img = cv2.rotate(img, cv2.ROTATE_90_COUNTERCLOCKWISE)
    
    # 1. Выравниваем
    aligned = auto_align(img)
    
    # 2. Ищем этикетку классическим EasyOCR по сжатому фото
    detected_label, label_crop_img = fast_extract_label(aligned)
    
    # 3. Находим лист
    mask = segment_leaf(aligned)
    
    # 4. Обрезаем картинку под лист
    roi, roi_mask, _ = smart_crop(aligned, mask)
    
    # 5. Ищем сетку
    mm2_per_pixel, grid_vis_roi = compute_auto_area_ratio(roi, roi_mask)

    if mm2_per_pixel is None:
        return JSONResponse({"status": "error", "message": "Сетка не найдена"}, status_code=400)

    area_mm2 = np.sum(roi_mask == 255) * mm2_per_pixel
    return {
        "label": detected_label,
        "area_mm2": round(float(area_mm2), 2),
        "area_cm2": round(float(area_mm2 / 100), 3),
        "previews": {
            "roi": image_to_base64(roi),
            "mask": image_to_base64(roi_mask),
            "grid": image_to_base64(grid_vis_roi),
            "label_crop": image_to_base64(label_crop_img)
        }
    }

@app.get("/")
async def index():
    return FileResponse("index.html")

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)