TEST / app.py
ezekieel's picture
Update app.py
ca15987 verified
raw
history blame
2.33 kB
curl \
-H 'Content-Type: application/json' \
-d '{"contents":[{"parts":[{"text":"Explain how AI works"}]}]}' \
-X POST 'https://generativelanguage.googleapis.com/v1beta/models/gemini-1.5-flash-latest:generateContent?key=AIzaSyCI51d8_uJlQ5h_YLJPnta5vHePLzWZnrY'
from transformers import AutoTokenizer, AutoModel
from PIL import Image
import torch
# Charger le modèle GOT-OCR2_0 pour la reconnaissance des plaques d'immatriculation
ocr_tokenizer = AutoTokenizer.from_pretrained('ucaslcl/GOT-OCR2_0', trust_remote_code=True)
ocr_model = AutoModel.from_pretrained('ucaslcl/GOT-OCR2_0', trust_remote_code=True, low_cpu_mem_usage=True, device_map='cuda', use_safetensors=True)
ocr_model.eval().cuda()
# Charger le modèle de suivi et reconnaissance de couleur des véhicules
vehicle_model = torch.hub.load('ultralytics/yolov5', 'custom', path='sujithvamshi/Real-Time-Vehicle-Tracking-And-Colour-Recognition/best.pt')
# Fonction pour extraire la plaque d'immatriculation avec OCR
def get_license_plate(image):
# Utiliser le modèle GOT-OCR pour extraire le texte (plaque d'immatriculation)
image_tensor = ocr_tokenizer(image, return_tensors="pt").input_ids
with torch.no_grad():
output = ocr_model(image_tensor)
plate_text = ocr_tokenizer.decode(output.logits[0], skip_special_tokens=True)
return plate_text
# Fonction pour extraire la couleur du véhicule
def get_vehicle_color(image):
# Utiliser le modèle Real-Time-Vehicle-Tracking-And-Colour-Recognition pour obtenir la couleur
results = vehicle_model(image)
color_info = results.pandas().xyxy[0].color # Hypothèse: le modèle retourne une info de couleur
return color_info
# Fusionner les deux résultats
def process_image(image_path):
image = Image.open(image_path)
# 1. Extraire la plaque d'immatriculation
license_plate = get_license_plate(image)
# 2. Extraire la couleur du véhicule
vehicle_color = get_vehicle_color(image)
# 3. Retourner la fusion des résultats
return {
"license_plate": license_plate,
"vehicle_color": vehicle_color
}
# Exemple d'utilisation
image_path = "path_to_your_image.jpg"
result = process_image(image_path)
print(f"Plaque d'immatriculation: {result['license_plate']}, Couleur du véhicule: {result['vehicle_color']}")