# ======= PATCH FIRST =======
import fastcore.transform as _fct
try:
    import fasttransform
    if not hasattr(_fct, 'Pipeline'):
        from fasttransform import Pipeline
        _fct.Pipeline = Pipeline
except:
    pass
import torch
_original_torch_load = torch.load
def _patched_load(*args, **kwargs):
    kwargs['weights_only'] = False
    return _original_torch_load(*args, **kwargs)
torch.load = _patched_load
# ======= NOW import fastai =======
from fastai.vision.all import *
import gradio as gr

path = '/content/clothing-dataset-full'
def get_x(r): 
    return path + '/images_original/' + r['image']
def get_y(r): 
    return r['label_cat'].split(' ')

# ======= Load model =======
model = load_learner("clothing_classifier.pkl", cpu=True)
all_labels = model.dls.vocab

def predict(img):
    img = PILImage.create(img)
    pred, pred_idx, probs = model.predict(img)

    age_labels      = ['Kids', 'Adults']
    clothing_labels = [l for l in all_labels if l not in age_labels]

    # أفضل clothing
    clothing_probs = {l: float(probs[list(all_labels).index(l)]) for l in clothing_labels}
    best_clothing  = max(clothing_probs, key=clothing_probs.get)

    # العمر
    kids_prob   = float(probs[list(all_labels).index('Kids')])
    adults_prob = float(probs[list(all_labels).index('Adults')])
    age = 'Kids' if kids_prob > adults_prob else 'Adults'

    # النتيجة: العنوان + top 4 ملابس بس (بدون Adults/Kids)
    top_clothing = dict(sorted(clothing_probs.items(), key=lambda x: x[1], reverse=True)[:4])
    result = {f" {best_clothing} for {age}": 1.0}
    result.update(top_clothing)

    return result

demo = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil"),
    outputs=gr.Label(num_top_classes=5),
    title="👗 Clothes Classifier",
    description="Upload a clothing image and the model will classify it!"
)
demo.launch(server_name="0.0.0.0", server_port=7860)