Spaces:
Sleeping
Sleeping
File size: 2,031 Bytes
e4724c0 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 |
import os
import re
import requests
import json
import subprocess
from io import BytesIO
import uuid
from math import ceil
from tqdm import tqdm
from pathlib import Path
from db import Database
DB_FOLDER = Path("diffusers-gallery-data")
database = Database(DB_FOLDER)
CLASSIFIER_URL = "https://radames-aesthetic-style-nsfw-classifier.hf.space/run/inference"
ASSETS_URL = "https://d26smi9133w0oo.cloudfront.net/diffusers-gallery/"
def main():
with database.get_db() as db:
cursor = db.cursor()
cursor.execute("""
SELECT *
FROM models
""")
results = list(cursor.fetchall())
for row in tqdm(results):
row_id = row['id']
# keep json data on row_data
row_data = json.loads(row['data'])
print("updating row", row_id)
images = row_data['images']
# filter nones
images = [i for i in images if i is not None]
if len(images) > 0:
# classifying only the first image
images_urls = [ASSETS_URL + images[0]]
response = requests.post(CLASSIFIER_URL, json={"data": [
{"urls": images_urls}, # json urls: list of images urls
False, # enable/disable gallery image output
None, # single image input
None, # files input
]}).json()
# data response is array data:[[{img0}, {img1}, {img2}...], Label, Gallery],
class_data = response['data'][0][0]
class_data_parsed = {row['label']: round(
row['score'], 3) for row in class_data}
# update row data with classificator data
row_data['class'] = class_data_parsed
else:
row_data['class'] = {}
with database.get_db() as db:
cursor = db.cursor()
cursor.execute("UPDATE models SET data = ? WHERE id = ?",
[json.dumps(row_data), row_id])
db.commit()
if __name__ == "__main__":
main()
|