File size: 7,311 Bytes
0a2704b
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
import numpy as np
import gradio as gr
from transformers import CLIPProcessor, CLIPModel
import torch
import itertools
import os

model = CLIPModel.from_pretrained("model-checkpoint")
processor = CLIPProcessor.from_pretrained("openai/clip-vit-large-patch14-336")

continents = ["Africa", "Asia", "Europe", "North America", "Oceania", "South America"]
countries_per_continent = {
    "Africa": [
        "Algeria", "Angola", "Benin", "Botswana", "Burkina Faso", "Burundi", "Cabo Verde", "Cameroon", 
        "Central African Republic", "Chad", "Comoros", "Congo", "Democratic Republic of the Congo", 
        "Djibouti", "Egypt", "Equatorial Guinea", "Eritrea", "Eswatini", "Ethiopia", "Gabon", 
        "Gambia", "Ghana", "Guinea", "Guinea-Bissau", "Ivory Coast", "Kenya", "Lesotho", "Liberia", 
        "Libya", "Madagascar", "Malawi", "Mali", "Mauritania", "Mauritius", "Morocco", "Mozambique", 
        "Namibia", "Niger", "Nigeria", "Rwanda", "Sao Tome and Principe", "Senegal", "Seychelles", 
        "Sierra Leone", "Somalia", "South Africa", "South Sudan", "Sudan", "Tanzania", "Togo", 
        "Tunisia", "Uganda", "Zambia", "Zimbabwe"
    ],
    "Asia": [
        "Afghanistan", "Armenia", "Azerbaijan", "Bahrain", "Bangladesh", "Bhutan", "Brunei", 
        "Cambodia", "China", "Cyprus", "Georgia", "India", "Indonesia", "Iran", "Iraq", 
        "Israel", "Japan", "Jordan", "Kazakhstan", "Kuwait", "Kyrgyzstan", "Laos", "Lebanon", 
        "Malaysia", "Maldives", "Mongolia", "Myanmar", "Nepal", "North Korea", "Oman", "Pakistan", 
        "Palestine", "Philippines", "Qatar", "Russia", "Saudi Arabia", "Singapore", "South Korea", 
        "Sri Lanka", "Syria", "Taiwan", "Tajikistan", "Thailand", "Timor-Leste", "Turkey", 
        "Turkmenistan", "United Arab Emirates", "Uzbekistan", "Vietnam", "Yemen"
    ],
    "Europe": [
        "Albania", "Andorra", "Armenia", "Austria", "Azerbaijan", "Belarus", "Belgium", "Bosnia and Herzegovina", 
        "Bulgaria", "Croatia", "Cyprus", "Czech Republic", "Denmark", "Estonia", "Finland", "France", 
        "Georgia", "Germany", "Greece", "Hungary", "Iceland", "Ireland", "Italy", "Kazakhstan", 
        "Kosovo", "Latvia", "Liechtenstein", "Lithuania", "Luxembourg", "Malta", "Moldova", "Monaco", 
        "Montenegro", "Netherlands", "North Macedonia", "Norway", "Poland", "Portugal", "Romania", 
        "Russia", "San Marino", "Serbia", "Slovakia", "Slovenia", "Spain", "Sweden", "Switzerland", 
        "Turkey", "Ukraine", "United Kingdom", "Vatican City"
    ],
    "North America": [
        "Antigua and Barbuda", "Bahamas", "Barbados", "Belize", "Canada", "Costa Rica", "Cuba", 
        "Dominica", "Dominican Republic", "El Salvador", "Grenada", "Guatemala", "Haiti", "Honduras", 
        "Jamaica", "Mexico", "Nicaragua", "Panama", "Saint Kitts and Nevis", "Saint Lucia", 
        "Saint Vincent and the Grenadines", "Trinidad and Tobago", "United States"
    ],
    "Oceania": [
        "Australia", "Fiji", "Kiribati", "Marshall Islands", "Micronesia", "Nauru", "New Zealand", 
        "Palau", "Papua New Guinea", "Samoa", "Solomon Islands", "Tonga", "Tuvalu", "Vanuatu"
    ],
    "South America": [
        "Argentina", "Bolivia", "Brazil", "Chile", "Colombia", "Ecuador", "Guyana", "Paraguay", 
        "Peru", "Suriname", "Uruguay", "Venezuela"
    ]
}
countries = list(set(itertools.chain.from_iterable(countries_per_continent.values())))

VERSUS_IMAGE = "versus_images/[email protected]"
VERSUS_GT = {
    "continent": "test",
    "country": "test"
}

def predict(input_img):
    inputs = processor(text=[f"A photo from {geo}." for geo in continents], images=input_img, return_tensors="pt", padding=True)
    with torch.no_grad():
        outputs = model(**inputs)
        logits_per_image = outputs.logits_per_image
        probs = logits_per_image.softmax(dim=-1)
        pred_id = probs.argmax().cpu().item()
    continent_probs = {label: prob for label, prob in zip(continents, probs.tolist()[0])}
    print(continent_probs)
    
    predicted_continent_countries = countries_per_continent[continents[pred_id]]
    inputs = processor(text=[f"A photo from {geo}." for geo in predicted_continent_countries], images=input_img, return_tensors="pt", padding=True)
    with torch.no_grad():
        outputs = model(**inputs)
        logits_per_image = outputs.logits_per_image
        probs = logits_per_image.softmax(dim=-1)
    country_probs = {label: prob for label, prob in zip(predicted_continent_countries, probs.tolist()[0])}
    print(country_probs)
    return continent_probs, country_probs

def versus_mode_inputs(input_img, human_continent, human_country):
    print(human_continent, human_country)
    continent_probs, country_probs = predict(input_img)
    return f"human guessed {human_continent} {human_country}\ngroung truth {VERSUS_GT}", continent_probs, country_probs

def next_versus_image():
    # VERSUS_GT["continent"] = "test"
    # VERSUS_GT["country"] = "test"
    return "versus_images/[email protected]"

def get_example_images(dir):
    image_extensions = (".jpg", ".jpeg", ".png")
    image_files = []
    for root, dirs, files in os.walk(dir):
        for file in files:
            if file.lower().endswith(image_extensions):
                image_files.append(os.path.join(root, file))
    return image_files

demo = gr.Blocks()
with demo:
    with gr.Tab("Image Geolocation Demo"):
        with gr.Row():
            with gr.Column():
                image = gr.Image(label="Image", type="pil", sources=["upload","clipboard"])
                predict_btn = gr.Button("Predict")
                example_images = get_example_images("examples")
                example_images.extend(get_example_images("versus_images"))
                gr.Examples(examples=example_images, inputs=image, examples_per_page=24)
            with gr.Column():
                continents_label = gr.Label(label="Continents")
                country_label = gr.Label(num_top_classes=5, label="Top countries")
                # continents_label.select(predict_country, inputs=[image, continents_label], outputs=country_label)
    predict_btn.click(predict, inputs=image, outputs=[continents_label, country_label])

    with gr.Tab("Versus Mode"):
        with gr.Row():
            with gr.Column():
                versus_image = gr.Image(VERSUS_IMAGE, interactive=False)
                continent_selection = gr.Radio(continents, label="Continents", info="Where was this image taken?")
                country_selection = gr.Dropdown(countries, label="Countries", info="Can you guess the exact country?"
                ),
                with gr.Row():
                    next_img_btn = gr.Button("Try new image")
                    versus_btn = gr.Button("Submit guess")
            with gr.Column():
                versus_output = gr.Text()
                continents_label = gr.Label(label="Continents")
                country_label = gr.Label(num_top_classes=5, label="Top countries")
        next_img_btn.click(next_versus_image, outputs=versus_image)
        versus_btn.click(versus_mode_inputs, inputs=[versus_image, continent_selection, country_selection[0]], outputs=[versus_output, continents_label, country_label])


if __name__ == "__main__":
    demo.launch()