feat: update new model

app.py

@@ -6,47 +6,74 @@ import json
 
 print("Đang tải các mô hình AI...")
 
-# 1. Load mô hình block (Phân loại ảnh vạn vật của Google)
-checker = pipeline("image-classification", model="google/vit-base-patch16-224")
-
-# 2. Load model Định giá
-model_id = "hquan21/ai-bike-pricing-up"
+checker   = pipeline("image-classification", model="google/vit-base-patch16-224")
+model_id  = "hquan21/ai-bike-pricing-gold-v2"
 processor = BlipProcessor.from_pretrained(model_id)
-model = BlipForConditionalGeneration.from_pretrained(model_id)
+model     = BlipForConditionalGeneration.from_pretrained(model_id)
+model.eval()
+
+print("Tải mô hình hoàn tất!")
 
-print("Tải mô hình hoàn tất! Sẵn sàng phục vụ.")
+TU_KHOA_XE = ['bicycle', 'bike', 'velocipede', 'mountain bike', 'tricycle']
 
 def predict(img):
-    ket_qua_nhan_dien = checker(img)
-    
-    # Lấy 5 nhãn có xác suất cao nhất
-    top_labels = [res['label'].lower() for res in ket_qua_nhan_dien[:5]]
-    
-    # Bộ từ khóa xe đạp trong từ điển ImageNet (viT model)
-    is_bicycle = any("bicycle" in label or "bike" in label or "velocipede" in label for label in top_labels)
-    
-    # Nếu đưa ảnh chó mèo, ô tô, người...
-    if not is_bicycle:
-        error_msg = {
-            "error": "Hình ảnh không chứa xe đạp, hoặc góc chụp không rõ ràng. Vui lòng chụp lại!"
-        }
-        # Trả về JSON để FE dễ dàng hiển thị Popup lỗi
-        return json.dumps(error_msg, ensure_ascii=False)
-
-    inputs = processor(img, return_tensors="pt")
-    
-    # Ép AI sinh ra tối đa 150 token (đủ dài cho JSON)
-    out = model.generate(**inputs, max_new_tokens=150)
-    result_text = processor.decode(out[0], skip_special_tokens=True)
-    
-    return result_text
+    # ── Gác cổng ──
+    try:
+        ket_qua = checker(img)
+        top5    = [r['label'].lower() for r in ket_qua[:5]]
+        is_bike = any(kw in label for label in top5 for kw in TU_KHOA_XE)
+    except Exception as e:
+        return json.dumps({
+            "success":     False,
+            "gia_de_xuat": None,
+            "ly_do":       None,
+            "error":       f"Lỗi kiểm tra ảnh: {str(e)}"
+        }, ensure_ascii=False)
+
+    if not is_bike:
+        return json.dumps({
+            "success":     False,
+            "gia_de_xuat": None,
+            "ly_do":       None,
+            "error":       "Không phải xe đạp! Vui lòng tải ảnh toàn thân xe."
+        }, ensure_ascii=False)
+
+    # ── Định giá ──
+    try:
+        with torch.no_grad():
+            inputs = processor(img, return_tensors="pt")
+            out    = model.generate(**inputs, max_new_tokens=150)
+        result = processor.decode(out[0], skip_special_tokens=True)
+    except Exception as e:
+        return json.dumps({
+            "success":     False,
+            "gia_de_xuat": None,
+            "ly_do":       None,
+            "error":       f"Lỗi định giá: {str(e)}"
+        }, ensure_ascii=False)
+
+    # ── Parse output ──
+    try:
+        data            = json.loads(result)
+        data["success"] = True
+        data["error"]   = None
+        return json.dumps(data, ensure_ascii=False)
+    except json.JSONDecodeError:
+        return json.dumps({
+            "success":     True,
+            "gia_de_xuat": None,
+            "ly_do":       result,
+            "error":       None
+        }, ensure_ascii=False)
 
 demo = gr.Interface(
-    fn=predict,
-    inputs=gr.Image(type="pil"),
-    outputs="text",
-    title="🚲 Công cụ Định giá Xe đạp AI",
-    description="Tải ảnh lên để AI định giá và đưa ra lý do. Cảnh báo: Hệ thống sẽ tự động từ chối nếu bạn tải ảnh chó, mèo, hay đồ vật khác không phải xe đạp!"
+    fn          = predict,
+    inputs      = gr.Image(type="pil"),
+    outputs     = "text",
+    title       = "Định giá Xe đạp AI",
+    description = "Tải ảnh toàn thân xe đạp. Hệ thống tự động từ chối ảnh không phải xe đạp.",
+    api_name    = "predict",
+    examples    = [],
 )
 
 demo.launch()

requirements.txt

@@ -2,3 +2,4 @@ gradio>=6.9.0
 transformers>=4.39.0
 torch>=2.2.0
 Pillow>=10.0.0
+accelerate>=0.27.0