import streamlit as st
import io
from PIL import Image
import numpy as np
import cv2
# from transformers import AutoImageProcessor, AutoModelForImageClassification
from PIL import Image

import requests
from transformers import pipeline

from torchvision import transforms
import torch


def kwg(photo):
    obj_detect = pipeline("object-detection", model='hustvl/yolos-small')
    age_detect = pipeline(model='nateraw/vit-age-classifier')
    classifier = pipeline(model="openai/clip-vit-large-patch14")
    objects_detected = obj_detect(photo)
    person_box_list = []
    for obj in objects_detected:
        if obj['label'] == 'person':
            person_box_list.append(obj['box'])
    if not person_box_list:
        st.write('На фото нет людей')
    else:
        st.write(f'на фото {len(person_box_list)} персон(а)')

    ages = []
    persons_coord_list = []
    img = np.array(photo)
    for box in person_box_list:
        person_coord = [box['ymin'], box['ymax'], box['xmin'], box['xmax']]
        persons_coord_list.append(person_coord)
    person_list = []
    for coords in persons_coord_list:
        person_list.append(Image.fromarray(img[coords[0]:coords[1],coords[2]:coords[3]]))
    for person in person_list:
        age = age_detect(person)
        ages.append(age[0]['label'])
    if '0-2' in ages or '3-9' in ages or '10-19' in ages:
        st.write('На фото есть дети')
    else:
        st.write('Здесь только взрослые')
        if len(ages) == 1:
            st.write(f'И ему {ages[0]} лет')
        else:
            for j in range(len(ages)):
                st.write(f'{ages[j]} лет')
            return

    res = classifier(photo, candidate_labels=["kid with gun", "kid with toy", "kid with alcohol drink"])
        
    if res[0]['label'] == "kid with gun":
        st.write('ОБОЖЕМОЙ у РЕБЕнкА ОРУЖИЕ СДЕЛАЙТЕ ЧТО-НИБУДЬ')
    elif res[0]['label'] == "kid with alcohol drink":
        st.write('ОТДАЙ ПИВО')
    else:
        st.write('какой милый ребеночек :3')


st.set_page_config(
    page_title="Emotion App!",
    page_icon="😎",
    layout="wide"
)

st.markdown("### Привет!")
# st.write("Загрузите фото")
# text = st.text_area("Введите текст:")


file = st.file_uploader("Загрузите своё фото:", type=['png','jpeg','jpg'])
if file:
    image_data = file.getvalue()
    # Показ загруженного изображения на Web-странице средствами Streamlit
    # st.image(image_data)
    # Возврат изображения в формате PIL
    image = Image.open(io.BytesIO(image_data))
    # image = Image.open("test"+username+".jpg").convert('RGB')
    st.image(image) # показать картинку
    # preprocessor = AutoImageProcessor.from_pretrained("google/mobilenet_v2_1.0_224")
    # model = AutoModelForImageClassification.from_pretrained("google/mobilenet_v2_1.0_224")
    kwg(image)



#
# else:
#     image = Image.open("testJulifil.jpg")


# img = st.image()