szili2011 commited on
Commit
c07926e
·
verified ·
1 Parent(s): e3248c0

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +67 -0
app.py ADDED
@@ -0,0 +1,67 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import torch
3
+ from transformers import ViTFeatureExtractor, ViTForImageClassification, pipeline
4
+ from PIL import Image
5
+ import easyocr
6
+ from aesthetic_score import AestheticScorer
7
+
8
+ # Initialize models
9
+ # Vision Transformer for Emotion Detection
10
+ feature_extractor = ViTFeatureExtractor.from_pretrained("nateraw/vit-base-beauty")
11
+ model = ViTForImageClassification.from_pretrained("nateraw/vit-base-beauty")
12
+ emotion_detector = pipeline("image-classification", model=model, feature_extractor=feature_extractor)
13
+
14
+ # OCR for Text Detection
15
+ reader = easyocr.Reader(['en'])
16
+
17
+ # Aesthetic Scoring Model (Placeholder)
18
+ aesthetic_scorer = AestheticScorer()
19
+
20
+ # Function to analyze a single thumbnail
21
+ def analyze_single_thumbnail(image):
22
+ # Text detection using OCR
23
+ ocr_result = reader.readtext(image)
24
+ text_detected = " ".join([item[1] for item in ocr_result]) if ocr_result else "No text found"
25
+
26
+ # Emotion detection for overall appeal
27
+ emotions = emotion_detector(image)
28
+ main_emotion = max(emotions, key=lambda x: x['score'])['label'] if emotions else "Unknown"
29
+
30
+ # Aesthetic scoring
31
+ aesthetic_score = aesthetic_scorer.score(image)
32
+
33
+ return {
34
+ "Detected Text": text_detected,
35
+ "Emotion Detected": main_emotion,
36
+ "Aesthetic Score": aesthetic_score,
37
+ }
38
+
39
+ # Function to compare two thumbnails
40
+ def compare_thumbnails(image1, image2):
41
+ result1 = analyze_single_thumbnail(image1)
42
+ result2 = analyze_single_thumbnail(image2)
43
+
44
+ # Determine which thumbnail might perform better based on scores
45
+ if result1["Aesthetic Score"] > result2["Aesthetic Score"]:
46
+ better_thumbnail = "Thumbnail 1 is likely better."
47
+ elif result2["Aesthetic Score"] > result1["Aesthetic Score"]:
48
+ better_thumbnail = "Thumbnail 2 is likely better."
49
+ else:
50
+ better_thumbnail = "Both thumbnails have similar appeal."
51
+
52
+ return result1, result2, better_thumbnail
53
+
54
+ # Gradio Interface
55
+ iface = gr.Interface(
56
+ fn=compare_thumbnails,
57
+ inputs=[gr.Image(type="pil"), gr.Image(type="pil")],
58
+ outputs=[
59
+ gr.JSON(label="Thumbnail 1 Analysis"),
60
+ gr.JSON(label="Thumbnail 2 Analysis"),
61
+ gr.Textbox(label="Comparison Result"),
62
+ ],
63
+ title="YouTube Thumbnail Comparator",
64
+ description="Upload two thumbnails to compare their effectiveness based on detected text, emotions, and aesthetic score.",
65
+ )
66
+
67
+ iface.launch()