Spaces:
Runtime error
Runtime error
| import numpy as np | |
| import gradio as gr | |
| from transformers import CLIPProcessor, CLIPModel | |
| import torch | |
| import itertools | |
| import os | |
| import plotly.graph_objects as go | |
| import hashlib | |
| from PIL import Image | |
| import json | |
| os.environ["PYTHONHASHSEED"] = "42" | |
| CUDA_AVAILABLE = torch.cuda.is_available() | |
| print(f"CUDA={CUDA_AVAILABLE}") | |
| device = "cuda" if CUDA_AVAILABLE else "cpu" | |
| print(f"count={torch.cuda.device_count()}") | |
| print(f"current={torch.cuda.get_device_name(torch.cuda.current_device())}") | |
| continent_model = CLIPModel.from_pretrained("jrheiner/thesis-clip-geoloc-continent", token=os.getenv("token")) | |
| country_model = CLIPModel.from_pretrained("jrheiner/thesis-clip-geoloc-country", token=os.getenv("token")) | |
| processor = CLIPProcessor.from_pretrained("jrheiner/thesis-clip-geoloc-continent", token=os.getenv("token")) | |
| continent_model = continent_model.to(device) | |
| country_model = country_model.to(device) | |
| continents = ["Africa", "Asia", "Europe", | |
| "North America", "Oceania", "South America"] | |
| countries_per_continent = { | |
| "Africa": [ | |
| "Botswana", "Eswatini", "Ghana", "Kenya", "Lesotho", "Nigeria", "Senegal", | |
| "South Africa", "Rwanda", "Uganda", "Tanzania", "Madagascar", "Djibouti", | |
| "Mali", "Libya", "Morocco", "Somalia", "Tunisia", "Egypt", "Réunion" | |
| ], | |
| "Asia": [ | |
| "Bangladesh", "Bhutan", "Cambodia", "China", "India", "Indonesia", "Israel", | |
| "Japan", "Jordan", "Kyrgyzstan", "Laos", "Malaysia", "Mongolia", "Nepal", | |
| "Palestine", "Philippines", "Singapore", "South Korea", "Sri Lanka", | |
| "Taiwan", "Thailand", "United Arab Emirates", "Vietnam", "Afghanistan", | |
| "Azerbaijan", "Cyprus", "Iran", "Syria", "Tajikistan", "Turkey", "Russia", | |
| "Pakistan", "Hong Kong" | |
| ], | |
| "Europe": [ | |
| "Albania", "Andorra", "Austria", "Belgium", "Bulgaria", "Croatia", "Czechia", | |
| "Denmark", "Estonia", "Finland", "France", "Germany", "Greece", "Hungary", | |
| "Iceland", "Ireland", "Italy", "Latvia", "Lithuania", "Luxembourg", | |
| "Montenegro", "Netherlands", "North Macedonia", "Norway", "Poland", | |
| "Portugal", "Romania", "Russia", "Serbia", "Slovakia", "Slovenia", "Spain", | |
| "Sweden", "Switzerland", "Ukraine", "United Kingdom", "Bosnia and Herzegovina", | |
| "Cyprus", "Turkey", "Greenland", "Faroe Islands" | |
| ], | |
| "North America": [ | |
| "Canada", "Dominican Republic", "Guatemala", "Mexico", "United States", | |
| "Bahamas", "Cuba", "Panama", "Puerto Rico", "Bermuda", "Greenland" | |
| ], | |
| "Oceania": [ | |
| "Australia", "New Zealand", "Fiji", "Papua New Guinea", "Solomon Islands", "Vanuatu" | |
| ], | |
| "South America": [ | |
| "Argentina", "Bolivia", "Brazil", "Chile", "Colombia", "Ecuador", "Paraguay", | |
| "Peru", "Uruguay" | |
| ] | |
| } | |
| countries = list(set(itertools.chain.from_iterable( | |
| countries_per_continent.values()))) | |
| country_to_center_coords = { | |
| "Indonesia": (-2.4833826, 117.8902853), | |
| "Egypt": (26.2540493, 29.2675469), | |
| "Dominican Republic": (19.0974031, -70.3028026), | |
| "Russia": (64.6863136, 97.7453061), | |
| "Denmark": (55.670249, 10.3333283), | |
| "Latvia": (56.8406494, 24.7537645), | |
| "Hong Kong": (22.350627, 114.1849161), | |
| "Brazil": (-10.3333333, -53.2), | |
| "Turkey": (38.9597594, 34.9249653), | |
| "Paraguay": (-23.3165935, -58.1693445), | |
| "Nigeria": (9.6000359, 7.9999721), | |
| "United Kingdom": (54.7023545, -3.2765753), | |
| "Argentina": (-34.9964963, -64.9672817), | |
| "United Arab Emirates": (24.0002488, 53.9994829), | |
| "Estonia": (58.7523778, 25.3319078), | |
| "Greenland": (69.6354163, -42.1736914), | |
| "Canada": (61.0666922, -107.991707), | |
| "Andorra": (42.5407167, 1.5732033), | |
| "Czechia": (49.7439047, 15.3381061), | |
| "Australia": (-24.7761086, 134.755), | |
| "Azerbaijan": (40.3936294, 47.7872508), | |
| "Cambodia": (12.5433216, 104.8144914), | |
| "Peru": (-6.8699697, -75.0458515), | |
| "Slovakia": (48.7411522, 19.4528646), | |
| "Réunion": (-21.130737949999997, 55.536480112992315), | |
| "France": (46.603354, 1.8883335), | |
| "Israel": (30.8124247, 34.8594762), | |
| "China": (35.000074, 104.999927), | |
| "Ecuador": (-1.3397668, -79.3666965), | |
| "Poland": (52.215933, 19.134422), | |
| "Switzerland": (46.7985624, 8.2319736), | |
| "Singapore": (1.357107, 103.8194992), | |
| "Kenya": (1.4419683, 38.4313975), | |
| "Bhutan": (27.549511, 90.5119273), | |
| "Laos": (20.0171109, 103.378253), | |
| "Vietnam": (15.9266657, 107.9650855), | |
| "Puerto Rico": (18.2247706, -66.4858295), | |
| "Germany": (51.1638175, 10.4478313), | |
| "Tanzania": (-6.5247123, 35.7878438), | |
| "Colombia": (4.099917, -72.9088133), | |
| "Italy": (42.6384261, 12.674297), | |
| "Bahamas": (24.7736546, -78.0000547), | |
| "Panama": (8.559559, -81.1308434), | |
| "Bulgaria": (42.6073975, 25.4856617), | |
| "Solomon Islands": (-8.7053941, 159.1070693851845), | |
| "Afghanistan": (33.7680065, 66.2385139), | |
| "Tajikistan": (38.6281733, 70.8156541), | |
| "Portugal": (39.6621648, -8.1353519), | |
| "Tunisia": (36.8002068, 10.1857757), | |
| "Bolivia": (-17.0568696, -64.9912286), | |
| "Malaysia": (4.5693754, 102.2656823), | |
| "Lithuania": (55.3500003, 23.7499997), | |
| "Sweden": (59.6749712, 14.5208584), | |
| "Belgium": (50.6402809, 4.6667145), | |
| "Libya": (26.8234472, 18.1236723), | |
| "Guatemala": (15.5855545, -90.345759), | |
| "India": (22.3511148, 78.6677428), | |
| "Sri Lanka": (7.5554942, 80.7137847), | |
| "New Zealand": (-41.5000831, 172.8344077), | |
| "Iceland": (64.9841821, -18.1059013), | |
| "Somalia": (8.3676771, 49.083416), | |
| "Croatia": (45.3658443, 15.6575209), | |
| "Bosnia and Herzegovina": (44.3053476, 17.5961467), | |
| "Greece": (38.9953683, 21.9877132), | |
| "Rwanda": (-1.9646631, 30.0644358), | |
| "Hungary": (47.1817585, 19.5060937), | |
| "Eswatini": (-26.5624806, 31.3991317), | |
| "Kyrgyzstan": (41.5089324, 74.724091), | |
| "Bangladesh": (23.6943117, 90.344352), | |
| "Morocco": (28.3347722, -10.371337908392647), | |
| "Finland": (63.2467777, 25.9209164), | |
| "Luxembourg": (49.6112768, 6.129799), | |
| "North Macedonia": (41.6171214, 21.7168387), | |
| "Uruguay": (-32.8755548, -56.0201525), | |
| "Chile": (-31.7613365, -71.3187697), | |
| "Spain": (39.3260685, -4.8379791), | |
| "South Korea": (36.638392, 127.6961188), | |
| "Botswana": (-23.1681782, 24.5928742), | |
| "Uganda": (1.5333554, 32.2166578), | |
| "Papua New Guinea": (-5.6816069, 144.2489081), | |
| "Mali": (16.3700359, -2.2900239), | |
| "Philippines": (12.7503486, 122.7312101), | |
| "Norway": (64.5731537, 11.52803643954819), | |
| "Thailand": (14.8971921, 100.83273), | |
| "Mongolia": (46.8651082, 103.8347844), | |
| "Japan": (36.5748441, 139.2394179), | |
| "Montenegro": (42.7044223, 19.3957785), | |
| "Austria": (47.59397, 14.12456), | |
| "Taiwan": (23.6978, 120.9605), | |
| "Netherlands": (52.2434979, 5.6343227), | |
| "Ukraine": (49.4871968, 31.2718321), | |
| "Fiji": (-18.1239696, 179.0122737), | |
| "Ghana": (8.0300284, -1.0800271), | |
| "Cuba": (23.0131338, -80.8328748), | |
| "Nepal": (28.3780464, 83.9999901), | |
| "Faroe Islands": (62.0448724, -7.0322972), | |
| "Slovenia": (46.1199444, 14.8153333), | |
| "Cyprus": (34.9174159, 32.889902651331866), | |
| "Serbia": (44.024322850000004, 21.07657433209902), | |
| "Madagascar": (-18.9249604, 46.4416422), | |
| "Pakistan": (30.3308401, 71.247499), | |
| "Syria": (34.6401861, 39.0494106), | |
| "Iran": (32.6475314, 54.5643516), | |
| "Ireland": (52.865196, -7.9794599), | |
| "South Africa": (-28.8166236, 24.991639), | |
| "Albania": (41.1529058, 20.1605717), | |
| "Lesotho": (-29.6039267, 28.3350193), | |
| "Romania": (45.9852129, 24.6859225), | |
| "Palestine": (31.947351, 35.227163), | |
| "Vanuatu": (-16.5255069, 168.1069154), | |
| "Mexico": (19.4326296, -99.1331785), | |
| "Jordan": (31.279862, 37.1297454), | |
| "Djibouti": (11.8145966, 42.8453061), | |
| "Senegal": (14.4750607, -14.4529612), | |
| "Bermuda": (32.3040273, -64.7563086), | |
| "United States": (39.7837304, -100.445882) | |
| } | |
| INTIAL_VERSUS_IMAGE = "versus_images/Europe_Germany_49.069183_10.319444_im2gps3k.jpg" | |
| INITAL_VERSUS_STATE = { | |
| "image": INTIAL_VERSUS_IMAGE, | |
| "continent": INTIAL_VERSUS_IMAGE.split("/")[-1].split("_")[0], | |
| "country": INTIAL_VERSUS_IMAGE.split("/")[-1].split("_")[1], | |
| "lat": INTIAL_VERSUS_IMAGE.split("/")[-1].split("_")[2], | |
| "lon": INTIAL_VERSUS_IMAGE.split("/")[-1].split("_")[3], | |
| "score": { | |
| "HUMAN": 0, | |
| "AI": 0 | |
| }, | |
| "idx": 0 | |
| } | |
| def predict(input_img): | |
| inputs = processor(text=[f"A photo from { | |
| geo}." for geo in continents], images=input_img, return_tensors="pt", padding=True) | |
| inputs = inputs.to(device) | |
| with torch.no_grad(): | |
| outputs = continent_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])} | |
| model_continent = continents[pred_id] | |
| predicted_continent_countries = countries_per_continent[model_continent] | |
| inputs = processor(text=[f"A photo from { | |
| geo}." for geo in predicted_continent_countries], images=input_img, return_tensors="pt", padding=True) | |
| inputs = inputs.to(device) | |
| with torch.no_grad(): | |
| outputs = country_model(**inputs) | |
| logits_per_image = outputs.logits_per_image | |
| probs = logits_per_image.softmax(dim=-1) | |
| pred_id = probs.argmax().cpu().item() | |
| model_country = predicted_continent_countries[pred_id] | |
| country_probs = {label: prob for label, prob in zip( | |
| predicted_continent_countries, probs.tolist()[0])} | |
| hash = hashlib.sha1(np.asarray(input_img).data.tobytes()).hexdigest() | |
| metadata_block = gr.Accordion(visible=False) | |
| metadata_map = None | |
| if hash in EXAMPLE_METADATA.keys(): | |
| model_result = "" | |
| if model_continent == EXAMPLE_METADATA[hash]['continent'] and model_country == EXAMPLE_METADATA[hash]['country']: | |
| model_result = "The AI 🤖 correctly guessed continent and country ✅ ✅." | |
| elif model_continent == EXAMPLE_METADATA[hash]['continent']: | |
| model_result = "The AI 🤖 only guessed the correct continent ❌ ✅." | |
| elif model_country == EXAMPLE_METADATA[hash]['country'] and model_continent != EXAMPLE_METADATA[hash]['continent']: | |
| model_result = "The AI 🤖 only guessed the correct country ✅ ❌." | |
| else: | |
| model_result = "The AI 🤖 failed to guess country and continent ❌ ❌." | |
| metadata_block = gr.Accordion(visible=True, label=f"This photo was taken in {EXAMPLE_METADATA[hash]['country']}, {EXAMPLE_METADATA[hash]['continent']}.\n{model_result}") | |
| metadata_map = make_versus_map(None, model_country, EXAMPLE_METADATA[hash]) | |
| return continent_probs, country_probs, metadata_block, metadata_map | |
| def make_versus_map(human_country, model_country, versus_state): | |
| if human_country: | |
| human_coordinates = country_to_center_coords[human_country] | |
| else: | |
| human_coordinates = (None, None) | |
| model_coordinates = country_to_center_coords[model_country] | |
| fig = go.Figure() | |
| fig.add_trace(go.Scattermapbox( | |
| lon=[versus_state["lon"]], | |
| lat=[versus_state["lat"]], | |
| text=[f"📷 Photo taken in {versus_state['country']}, { | |
| versus_state['continent']}"], | |
| mode='markers', | |
| hoverinfo='text', | |
| marker=dict(size=14, color='#0C5DA5'), | |
| showlegend=True, | |
| name="📷 Photo Location" | |
| )) | |
| if human_country == model_country: | |
| fig.add_trace(go.Scattermapbox( | |
| lat=[human_coordinates[0], model_coordinates[0]], | |
| lon=[human_coordinates[1], model_coordinates[1]], | |
| text=f"🧑 🤖 Human & AI guess {human_country}", | |
| mode='markers', | |
| hoverinfo='text', | |
| marker=dict(size=14, color='#FF9500'), | |
| showlegend=True, | |
| name="🧑 🤖 Human & AI Guess" | |
| )) | |
| else: | |
| if human_country: | |
| fig.add_trace(go.Scattermapbox( | |
| lat=[human_coordinates[0]], | |
| lon=[human_coordinates[1]], | |
| text=[f"🧑 Human guesses {human_country}"], | |
| mode='markers', | |
| hoverinfo='text', | |
| marker=dict(size=14, color='#FF9500'), | |
| showlegend=True, | |
| name="🧑 Human Guess" | |
| )) | |
| fig.add_trace(go.Scattermapbox( | |
| lat=[model_coordinates[0]], | |
| lon=[model_coordinates[1]], | |
| text=[f"🤖 AI guesses {model_country}"], | |
| mode='markers', | |
| hoverinfo='text', | |
| marker=dict(size=14, color='#474747'), | |
| showlegend=True, | |
| name="🤖 AI Guess" | |
| )) | |
| fig.update_layout( | |
| mapbox=dict( | |
| style="carto-positron", | |
| center=dict(lat=float(versus_state["lat"]), lon=float(versus_state["lon"])), | |
| zoom=2 | |
| ), | |
| margin={"r": 0, "t": 0, "l": 0, "b": 0}, | |
| legend=dict( | |
| yanchor="bottom", | |
| y=0.01, | |
| xanchor="left", | |
| x=0.01 | |
| ) | |
| ) | |
| return fig | |
| def versus_mode_inputs(input_img, human_continent, human_country, versus_state): | |
| human_points = 0 | |
| model_points = 0 | |
| if human_country == versus_state["country"]: | |
| country_result = "✅" | |
| human_points += 2 | |
| else: | |
| country_result = "❌" | |
| if human_continent == versus_state["continent"]: | |
| continent_result = "✅" | |
| human_points += 1 | |
| else: | |
| continent_result = "❌" | |
| human_result = f"The photo is from **{versus_state['country']}** { | |
| country_result} in **{versus_state['continent']}** {continent_result}" | |
| human_score_update = f"+{human_points} points" if human_points > 0 else "0 Points..." | |
| versus_state['score']['HUMAN'] += human_points | |
| continent_probs, country_probs, _,_ = predict(input_img) | |
| model_country = max(country_probs, key=country_probs.get) | |
| model_continent = max(continent_probs, key=continent_probs.get) | |
| if model_country == versus_state["country"]: | |
| model_country_result = "✅" | |
| model_points += 2 | |
| else: | |
| model_country_result = "❌" | |
| if model_continent == versus_state["continent"]: | |
| model_continent_result = "✅" | |
| model_points += 1 | |
| else: | |
| model_continent_result = "❌" | |
| model_score_update = f"+{model_points} points" if model_points > 0 else "0 Points... The model was completely wrong, it seems the world is not doomed yet." | |
| versus_state['score']['AI'] += model_points | |
| map = make_versus_map(human_country, model_country, versus_state) | |
| return f""" | |
| ## {human_result} | |
| ### The AI 🤖 thinks this photo is from **{model_country}** {model_country_result} in **{model_continent}** {model_continent_result} | |
| 🧑 {human_score_update} | |
| 🤖 {model_score_update} | |
| ### Score 🧑 {versus_state['score']['HUMAN']} : {versus_state['score']['AI']} 🤖 | |
| """, continent_probs, country_probs, map, versus_state | |
| 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 | |
| def next_versus_image(versus_state): | |
| images = get_example_images("versus_images") | |
| versus_state["idx"] += 1 | |
| if versus_state["idx"] > len(images): | |
| versus_state["idx"] = 0 | |
| versus_image = images[versus_state["idx"]] | |
| versus_state["continent"] = versus_image.split("/")[-1].split("_")[0] | |
| versus_state["country"] = versus_image.split("/")[-1].split("_")[1] | |
| versus_state["lat"] = versus_image.split("/")[-1].split("_")[2] | |
| versus_state["lon"] = versus_image.split("/")[-1].split("_")[3] | |
| versus_state["image"] = versus_image | |
| return versus_image, versus_state, None, None | |
| example_images = get_example_images("kerger-test-images") | |
| EXAMPLE_METADATA = {} | |
| for img_path in example_images: | |
| hash = hashlib.sha1(np.asarray(Image.open(img_path)).data.tobytes()).hexdigest() | |
| EXAMPLE_METADATA[hash] = { | |
| "continent": img_path.split("/")[-1].split("_")[0], | |
| "country": img_path.split("/")[-1].split("_")[1], | |
| "lat": img_path.split("/")[-1].split("_")[2], | |
| "lon": img_path.split("/")[-1].split("_")[3], | |
| } | |
| demo = gr.Blocks(title="Thesis Demo") | |
| with demo: | |
| gr.HTML(""" | |
| <h1 style="text-align: center; margin-bottom: 1rem">Image Geolocation Thesis Demo</h1> | |
| <h3> This Demo showcases the developed models and allows interacting with the optimized prototype.</h3> | |
| <p>Try the <b>"Image Geolocation Demo"</b> tab with your own images or with one of the examples. For all example image the ground truth is available and will be displayed together with the model predictions.</p> | |
| <p>In the <b>"Versus Mode"</b> tab to play against the AI, guessing the country and continent where images where taken. Images in the versus mode are from the <a href="http://graphics.cs.cmu.edu/projects/im2gps/"><code>Im2GPS</code></a> and <a href="https://arxiv.org/abs/1705.04838"><code>Im2GPS3k</code></a> geolocation literature benchmarks. Can you beat the AI? | |
| """) | |
| with gr.Accordion(label="The demo currently encompasses 116 countries from 6 continents 🌍", open=False): | |
| gr.Code(json.dumps(countries_per_continent, indent=2, ensure_ascii=False), label="countries_per_continent.json", language="json", interactive=False) | |
| 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("kerger-test-images") | |
| # example_images.extend(get_example_images("versus_images")) | |
| gr.Examples(examples=example_images, | |
| inputs=image, examples_per_page=24) | |
| with gr.Column(): | |
| with gr.Accordion(visible=False) as metadata_block: | |
| map = gr.Plot(label="Locations") | |
| with gr.Group(): | |
| continents_label = gr.Label(label="Continents") | |
| country_label = gr.Label( | |
| num_top_classes=5, label="Top countries") | |
| predict_btn.click(predict, inputs=image, outputs=[ | |
| continents_label, country_label, metadata_block, map]) | |
| with gr.Tab("Versus Mode"): | |
| versus_state = gr.State(value=INITAL_VERSUS_STATE) | |
| with gr.Row(): | |
| with gr.Column(): | |
| versus_image = gr.Image( | |
| INITAL_VERSUS_STATE["image"], interactive=False) | |
| continent_selection = gr.Radio( | |
| continents, label="Continents", info="Where was this image taken? (1 Point)") | |
| country_selection = gr.Dropdown(countries, label="Countries", info="Can you guess the exact country? (2 Points)"), | |
| with gr.Row(): | |
| next_img_btn = gr.Button("Try new image") | |
| versus_btn = gr.Button("Submit guess") | |
| with gr.Column(): | |
| versus_output = gr.Markdown() | |
| # with gr.Accordion("View Map", open=False): | |
| map = gr.Plot(label="Locations") | |
| with gr.Accordion("Full Model Output", open=False): | |
| with gr.Group(): | |
| continents_label = gr.Label(label="Continents") | |
| country_label = gr.Label( | |
| num_top_classes=5, label="Top countries") | |
| next_img_btn.click(next_versus_image, inputs=[versus_state], outputs=[ | |
| versus_image, versus_state, continent_selection, country_selection[0]]) | |
| versus_btn.click(versus_mode_inputs, inputs=[versus_image, continent_selection, country_selection[0], versus_state], outputs=[ | |
| versus_output, continents_label, country_label, map, versus_state]) | |
| if __name__ == "__main__": | |
| demo.launch(show_api=False) | |