omarkhattab28/computervision

1import os2import cv23import numpy as np4from fastapi import FastAPI, UploadFile, File5from ultralytics import YOLO6from transformers import TrOCRProcessor, VisionEncoderDecoderModel7from PIL import Image8import torch9import difflib10 11app = FastAPI()12 13# 1. تحميل الموديلات14yolo_model = YOLO("best.pt")15processor = TrOCRProcessor.from_pretrained("microsoft/trocr-base-handwritten")16trocr_model = VisionEncoderDecoderModel.from_pretrained("microsoft/trocr-base-handwritten")17 18# 2. القاموس الموسع (لأشهر الأدوية)19DRUG_DICTIONARY = [20    "Panadol", "Adol", "Paramol", "Abimol", "Catafast", "Cataflam", "Voltaren", "Brufen", "Antiflam",21    "Augmentin", "Curam", "Hibiotic", "Flumox", "Zithrokan", "Zithromax", "Clavimox", "Amoxil",22    "Antinal", "Flagyl", "Controloc", "Nexium", "Gastrolog", "Visceralgin", "Librax",23    "Congestal", "123", "Comtrex", "Night-and-day", "Telfast", "Zyrtec", "Claritine",24    "Concor", "Exforge", "Capoten", "Aspirin", "Jusprin", "Plavix", "Glucophage",25    "Neuroton", "Milga", "Vitamin C", "Feroglobin", "Betadine", "Fucidin", "Mebo",26    "Ventolin", "Farcolin", "Amrizole", "Dexamethasone", "Daflon"27]28 29def enhance_for_ocr(cv2_img):30    """تحسين جودة الكلمة المقطوعة"""31    gray = cv2.cvtColor(cv2_img, cv2.COLOR_BGR2GRAY)32    resized = cv2.resize(gray, None, fx=3, fy=3, interpolation=cv2.INTER_CUBIC)33    clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8,8))34    enhanced = clahe.apply(resized)35    return enhanced36 37def get_best_match(text):38    """تصحيح الاسم بناءً على القاموس"""39    text = text.strip().capitalize()40    if len(text) < 3: return text, False41    matches = difflib.get_close_matches(text, DRUG_DICTIONARY, n=1, cutoff=0.5)42    if matches:43        return matches[0], True44    return text, False45 46@app.post("/predict")47async def predict(image: UploadFile = File(...)):48    contents = await image.read()49    nparr = np.frombuffer(contents, np.uint8)50    img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)51 52    # 1. اكتشاف الكلمات53    results = yolo_model(img, conf=0.12, iou=0.45)54    55    # --- التعديل الجديد: ترتيب الصناديق من فوق لتحت (محور Y) ---56    # b.xyxy[0][1] هو إحداثي الـ Y العلوي لكل صندوق57    boxes = results[0].boxes58    sorted_boxes = sorted(boxes, key=lambda b: b.xyxy[0][1].item())59    60    final_output = []61    62    for box in sorted_boxes:63        x1, y1, x2, y2 = map(int, box.xyxy[0])64        word_crop = img[y1:y2, x1:x2]65        66        # 2. تحسين وقراءة الكلمة67        processed_word = enhance_for_ocr(word_crop)68        pil_img = Image.fromarray(processed_word).convert("RGB")69        70        pixel_values = processor(images=pil_img, return_tensors="pt").pixel_values71        generated_ids = trocr_model.generate(pixel_values)72        raw_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]73        74        # 3. التصحيح والمطابقة75        corrected_name, is_found = get_best_match(raw_text)76        77        final_output.append({78            "detected_word": corrected_name,79            "raw_reading": raw_text,80            "status": "Verified" if is_found else "Unknown"81        })82 83    return {84        "total_items": len(final_output),85        "results": final_output86    }87 88@app.get("/")89def home():90    return {"message": "OCR System with Sorted Output is Live!"}