TEST / app.py
nutellalpb's picture
Update app.py (#4)
9ff6f0a verified
raw
history blame
2.62 kB
import requests
# URL de l'API
url = 'https://generativelanguage.googleapis.com/v1beta/models/gemini-1.5-flash-latest:generateContent?key=AIzaSyB5WiEJf_yLMD1dMQf305EAbaPTzF_QD-I
# En-têtes de la requête
headers = {
'Content-Type': 'application/json',
}
# Corps de la requête
data = {
"contents": [
{
"parts": [
{"text": "Explain how AI works"}
]
}
]
}
# Envoyer la requête POST
response = requests.post(url, headers=headers, json=data)
# Afficher la réponse
print(response.json())
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']}")