initial commit

.gitignore

@@ -0,0 +1,3 @@
+.venv/
+model-checkpoint/
+__pycache__/

app.py

@@ -0,0 +1,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()