hwajjala commited on
Commit
05b1066
1 Parent(s): db84f03

Add aq model

Browse files
Files changed (2) hide show
  1. app.py +57 -0
  2. requirements.txt +9 -0
app.py ADDED
@@ -0,0 +1,57 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import clip
3
+ import torch
4
+ import logging
5
+ import json
6
+ import pandas as pd
7
+ from PIL import Image
8
+ import gradio as gr
9
+ from autogluon.tabular import TabularPredictor
10
+
11
+ predictor = TabularPredictor.load("AutogluonModels/ag-20240615_190835")
12
+ # set logging level
13
+ logging.basicConfig(
14
+ level=logging.INFO,
15
+ format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
16
+ )
17
+ logger = logging.getLogger("AQ")
18
+ CLIP_MODEL_NAME = "ViT-B/32"
19
+
20
+
21
+ clip_model, preprocess = clip.load(CLIP_MODEL_NAME, device="cpu")
22
+
23
+
24
+ def predict_fn(input_img):
25
+ input_img = Image.fromarray(input_img.astype("uint8"), "RGB")
26
+ image = preprocess(input_img).unsqueeze(0)
27
+ with torch.no_grad():
28
+ image_features = clip_model.encode_image(image).numpy()
29
+ input_df = pd.DataFrame(image_features[0].reshape(1, -1))
30
+ quality_score = predictor.predict(input_df).iloc[0]
31
+
32
+ logger.info(f"decision: {quality_score}")
33
+ decision_json = json.dumps({"quality_score": quality_score}).encode("utf-8")
34
+ logger.info(f"decision_json: {decision_json}")
35
+ return decision_json
36
+
37
+
38
+ iface = gr.Interface(
39
+ fn=predict_fn,
40
+ inputs="image",
41
+ outputs="text",
42
+ description="""
43
+ The model returns the probability of the image being a base body. If
44
+ probability > 0.9, the image can be automatically tagged as a base body. If
45
+ probability < 0.2, the image can be automatically REJECTED as NOT as base
46
+ body. All other cases will be submitted for moderation.
47
+
48
+ Please flag if you think the decision is wrong.
49
+ """,
50
+ allow_flagging="manual",
51
+ flagging_options=[
52
+ ": decision should be accept",
53
+ ": decision should be reject",
54
+ ": decision should be moderation",
55
+ ],
56
+ )
57
+ iface.launch()
requirements.txt ADDED
@@ -0,0 +1,9 @@
 
 
 
 
 
 
 
 
 
 
1
+ torch
2
+ torchvision
3
+ ftfy
4
+ regex
5
+ tqdm
6
+ autogluon
7
+ git+https://github.com/openai/CLIP.git
8
+ scikit-learn==1.3.0
9
+ scipy