app.py

from cProfile import label

import gradio as gr
from transformers import BlipProcessor, BlipForConditionalGeneration, pipeline
from PIL import Image
import torch
import json

# check if images are bicycle images
checker = pipeline("image-classification", model="google/vit-base-patch16-224")

# Load model của bạn
model_id = "hquan21/ai-bike-pricing"
processor = BlipProcessor.from_pretrained(model_id)
model = BlipForConditionalGeneration.from_pretrained(model_id)


TU_KHOA_XE = ['bicycle', 'bike', 'velocipede', 'mountain bike', 'tricycle']

def predict(img):

    # Kiểm tra xem ảnh có phải là ảnh xe đạp không
    try:
        results = checker(img)
        top5 = [r['label'].lower() for r in results[: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,
            "error": f"Error during image classification: {str(e)}"
        }, ensure_ascii=False)
    if not is_bike:
        return json.dumps({
            "success": False,
            "error": "Ảnh không phải là xe đạp. Vui lòng tải lên ảnh xe đạp."
        }, ensure_ascii=False)


    inputs = processor(img, return_tensors="pt")
    out = model.generate(**inputs, max_new_tokens=150)
    result_text = processor.decode(out[0], skip_special_tokens=True)
    return result_text

# Giao diện
demo = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil"),
    outputs="text",
    title="Công cụ Định giá Xe đạp",
    description="Tải ảnh lên để định giá và đưa ra lý do."
)

demo.launch()