File size: 2,325 Bytes
07dc5b2
 
 
 
567f165
a6bd9de
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
567f165
a6bd9de
 
567f165
a6bd9de
 
567f165
a6bd9de
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
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']}")