import torch
import gradio as gr
from huggingface_hub import hf_hub_download
from PIL import Image

REPO_ID = "thoucentric/Shelf_Objects_Detection_Yolov7_Pytorch"
FILENAME = "best.pt"


yolov7_custom_weights = hf_hub_download(repo_id=REPO_ID, filename=FILENAME)

model = torch.hub.load('Owaiskhan9654/yolov7-1:main',model='custom', path_or_model=yolov7_custom_weights, force_reload=True)  # Github repository https://github.com/Owaiskhan9654

def object_detection(image: gr.inputs.Image = None):
  

    results = model(image) 

    results.render()  
    count_dict = results.pandas().xyxy[0]['name'].value_counts().to_dict()


    if len(count_dict)>0:
      return Image.fromarray(results.imgs[0]),str(count_dict)
    else:
        return Image.fromarray(results.imgs[0]),'No object Found. Add more Custom classes in the training set'


title = "Yolov7 Custom"

# image = gr.inputs.Image(shape=(640, 640), image_mode="RGB", source="upload", label="Upload Image", optional=False)

inputs = gr.inputs.Image(shape=(640, 640), image_mode="RGB", source="upload", label="Upload Image", optional=False)
#     gr.inputs.Dropdown(["best.pt",], 
#                        default="best.pt", label="Model"),
#     gr.inputs.Slider(minimum=320, maximum=1280, default=640, step=32, label="Image Size"),
#     gr.inputs.Slider(minimum=0.0, maximum=1.0, default=0.25, step=0.05, label="Confidence Threshold"),
#     gr.inputs.Slider(minimum=0.0, maximum=1.0, default=0.45, step=0.05, label="IOU Threshold"),
# ]
outputs = gr.outputs.Image(type="pil", label="Output Image")
outputs_cls = gr.Label(label= "Categories Detected Proportion Statistics" )


Custom_description="<center>Custom Training Performed on Kaggle <a href='https://www.kaggle.com/code/owaiskhan9654/shelf-object-detection-yolov7-pytorch/notebook' style='text-decoration: underline' target='_blank'>Link</a> </center><br> <center>Trainable bag-of-freebies sets new state-of-the-art for real-time object detectors </center> <br> Works on around <b>140</b> general items in Stores"

Footer = (
    
    "<br><br><br><br><center><b>Item Classes it will detect(Total 140 Classes) -</b> "
    "'Drawbar box', 'Disposable cups', 'Makeup tools', 'Television', 'Toothpaste', 'Herbal tea', 'Skate', 'Coat hanger', 'Soy sauce', "
    "'Tea beverage', 'Sour Plum Soup', 'Pie', 'Chopping block', 'Refrigerator', 'Trousers', 'Oats', 'Rubber ball', 'Soap', 'Pasta', 'Juicer', "
    "'Walnut powder', 'Toothbrush', 'Chopsticks', 'Mouth wash', 'Adult socks', 'Dinner plate', 'Baby milk powder', 'Soymilk', 'Cutter', 'Hair drier', " 
    "'Electric frying pan', 'Children hats', 'Cake', 'Trash', 'Children underwear', 'Guozhen', 'Disposable bag', 'Jacket', 'Baby carriage', 'Bowl', "
    "'Baby tableware', 'Emulsion', 'Red wine', 'Mixed congee', 'Spoon', 'Dried meat', 'Dairy', 'Chewing gum', 'Cooking wine', 'Electromagnetic furnace', "
    "'Facial Cleanser', 'Sports cup', 'Quick-frozen Wonton', 'Dried fish', 'Rice cooker', 'Children shoes', 'Band aid', 'Biscuits', 'Soybean Milk machine', "
    "'Pen', 'Baby crib', 'Hair gel', 'Children Toys', 'Ice cream', 'Washing machine', 'Hot strips', 'Air conditioning fan', 'Pencil case', 'Hair conditioner'," 
    "'Razor', 'Children Socks', 'Basin', 'Chocolates', 'Shampoo', 'Soup ladle', 'Men underwear', 'Baby washing and nursing supplies', 'Noodle', 'Tampon', " 
    "'Forks', 'Liquor and Spirits', 'Bath lotion', 'Knives', 'Quick-frozen dumplings', 'Socket', 'Notebook', 'Bedding set', 'Storage box', 'Ginger Tea', " 
    "'Basketball', 'Baby Toys', 'Storage bottle', 'Instant noodles', 'Baby Furniture', 'Thermos bottle', 'Hair dye', 'Fish tofu', 'Vinegar', 'Comb', "
    "'Carbonated drinks', 'Sauce', 'Adult shoes', 'Quick-frozen Tangyuan', 'Stool', 'Football', 'Baby diapers', 'Lotus root flour', 'Air conditioner', " 
    "'Badminton', 'Knapsack', 'Adult Diapers', 'Flour', 'Sesame paste', 'Pot shovel', 'Electric kettle', 'Mug', 'Electric iron', 'Lingerie', 'Tea', " 
    "'Food box', 'Electric Hot pot', 'Baby slippers', 'Potato chips', 'Electric steaming pan', 'Rise', 'Adult hat', 'Can', 'Care Kit', 'Cotton swab', " 
    "'Baby handkerchiefs ', 'Fresh-keeping film', 'Dried beans', 'Electric fan', 'Desk lamp', 'Cocktail', 'Skincare set', 'Adult milk powder', " 
    "'Microwave Oven', 'Coffee', 'Facial mask'</center>"

    "<br><br><br><br><center>Model Trained by: Owais Ahmad Data Scientist at <b><a href=\"https://thoucentric.com/\">Thoucentric</a></b><br></center>"
    
    "<center> Model Trained Kaggle Kernel <a href=\"https://www.kaggle.com/code/owaiskhan9654/shelf-object-detection-yolov7-pytorch/notebook\">Link</a> <br></center>"
        
    
    "<center> HuggingFace🤗 Model Deployed Repository <a href=\"https://huggingface.co/thoucentric/Shelf_Objects_Detection_Yolov7_Pytorch\">Link</a> <br></center>"
    "<center>&copy; 2023 <b><a href=\"https://thoucentric.com/\">Thoucentric</a></b> </center>"

)

examples1=[["Images/Image1.jpg"],["Images/Image2.jpg"],["Images/Image3.jpg"],["Images/Image4.jpg"],["Images/Image5.jpg"],["Images/Image6.jpg"]]

Top_Title="Yolov7 🚀 Custom Trained by <a href='https://www.linkedin.com/in/owaiskhan9654/' style='text-decoration: underline' target='_blank'>Owais Ahmad </center></a> on around 140 general items in Stores"
css = ".output-image, .input-image {height: 50rem !important; width: 100% !important;}"
css = ".image-preview {height: auto !important;}"

gr.Interface(
    fn=object_detection,
    inputs=inputs,
    outputs=[outputs,outputs_cls],
    title=Top_Title,
    description=Custom_description,
    article=Footer,
    cache= False,
    allow_flagging='never',
    examples=examples1).launch()