|
import streamlit as st |
|
import io |
|
from PIL import Image |
|
import numpy as np |
|
import cv2 |
|
|
|
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("### Привет!") |
|
|
|
|
|
|
|
|
|
file = st.file_uploader("Загрузите своё фото:", type=['png','jpeg','jpg']) |
|
if file: |
|
image_data = file.getvalue() |
|
|
|
|
|
|
|
image = Image.open(io.BytesIO(image_data)) |
|
|
|
st.image(image) |
|
|
|
|
|
kwg(image) |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
