Update app.py (#6)

- Update app.py (0458920d994b4ae10b036d171b9e880ec17f47e3)


Co-authored-by: Ben <[email protected]>

app.py

@@ -1,77 +1,30 @@
-import requests
+from transformers import AutoModel, AutoTokenizer
 
-# URL de l'API
-url = 'https://generativelanguage.googleapis.com/v1beta/models/gemini-1.5-flash-latest:generateContent?key=AIzaSyB5WiEJf_yLMD1dMQf305EAbaPTzF_QD-I
+tokenizer = AutoTokenizer.from_pretrained('ucaslcl/GOT-OCR2_0', trust_remote_code=True)
+model = AutoModel.from_pretrained('ucaslcl/GOT-OCR2_0', trust_remote_code=True, low_cpu_mem_usage=True, device_map='cuda', use_safetensors=True, pad_token_id=tokenizer.eos_token_id)
+model = model.eval().cuda()
 
-# 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"}
-            ]
-        }
-    ]
-}
+# input your test image
+image_file = 'car.jpg'
 
-# Envoyer la requête POST
-response = requests.post(url, headers=headers, json=data)
+# plain texts OCR
+res = model.chat(tokenizer, image_file, ocr_type='ocr')
 
-# Afficher la réponse
-print(response.json())
+# format texts OCR:
+# res = model.chat(tokenizer, image_file, ocr_type='format')
 
+# fine-grained OCR:
+# res = model.chat(tokenizer, image_file, ocr_type='ocr', ocr_box='')
+# res = model.chat(tokenizer, image_file, ocr_type='format', ocr_box='')
+# res = model.chat(tokenizer, image_file, ocr_type='ocr', ocr_color='')
+# res = model.chat(tokenizer, image_file, ocr_type='format', ocr_color='')
 
+# multi-crop OCR:
+# res = model.chat_crop(tokenizer, image_file, ocr_type='ocr')
+# res = model.chat_crop(tokenizer, image_file, ocr_type='format')
 
-from transformers import AutoTokenizer, AutoModel
-from PIL import Image
-import torch
+# render the formatted OCR results:
+# res = model.chat(tokenizer, image_file, ocr_type='format', render=True, save_render_file = './demo.html')
 
-# 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']}")
+print(res)