import os import clip import torch import logging import json import pandas as pd from PIL import Image import gradio as gr from autogluon.tabular import TabularPredictor predictor = TabularPredictor.load("AutogluonModels/ag-20240615_190835") # set logging level logging.basicConfig( level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", ) logger = logging.getLogger("AQ") CLIP_MODEL_NAME = "ViT-B/32" clip_model, preprocess = clip.load(CLIP_MODEL_NAME, device="cpu") def predict_fn(input_img): input_img = Image.fromarray(input_img.astype("uint8"), "RGB") image = preprocess(input_img).unsqueeze(0) with torch.no_grad(): image_features = clip_model.encode_image(image).numpy() input_df = pd.DataFrame(image_features[0].reshape(1, -1)) quality_score = predictor.predict(input_df).iloc[0] logger.info(f"decision: {quality_score}") decision_json = json.dumps({"quality_score": quality_score}).encode("utf-8") logger.info(f"decision_json: {decision_json}") return decision_json iface = gr.Interface( fn=predict_fn, inputs="image", outputs="text", description=""" The model returns the probability of the image being a base body. If probability > 0.9, the image can be automatically tagged as a base body. If probability < 0.2, the image can be automatically REJECTED as NOT as base body. All other cases will be submitted for moderation. Please flag if you think the decision is wrong. """, allow_flagging="manual", flagging_options=[ ": decision should be accept", ": decision should be reject", ": decision should be moderation", ], ) iface.launch()