Spaces:
Sleeping
Sleeping
File size: 4,608 Bytes
6efeffc 45f7be1 6efeffc 45f7be1 6efeffc 45f7be1 |
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 |
import gradio as gr
import os
import json
from openai import OpenAI
from geopy.geocoders import Nominatim
from folium import Map, GeoJson
from gradio_folium import Folium
import cv2
import numpy as np
import torch
from diffusers import ControlNetModel, StableDiffusionControlNetInpaintPipeline
from PIL import Image
import io
# Initialize APIs
openai_client = OpenAI(api_key=os.environ['OPENAI_API_KEY'])
geolocator = Nominatim(user_agent="geoapi")
# Function to fetch coordinates
def get_geo_coordinates(location_name):
try:
location = geolocator.geocode(location_name)
if location:
return [location.longitude, location.latitude]
return None
except Exception as e:
print(f"Error fetching coordinates for {location_name}: {e}")
return None
# Function to process OpenAI chat response
def process_openai_response(query):
response = openai_client.chat.completions.create(
model="gpt-4o-mini",
messages=[
{"role": "system", "content": "You are a skilled assistant answering geographical and historical questions..."},
{"role": "user", "content": query}
],
temperature=1,
max_tokens=2048,
top_p=1,
frequency_penalty=0,
presence_penalty=0,
response_format={"type": "json_object"}
)
return json.loads(response.choices[0].message.content)
# Generate GeoJSON from OpenAI response
def generate_geojson(response):
feature_type = response['output']['feature_representation']['type']
city_names = response['output']['feature_representation']['cities']
properties = response['output']['feature_representation']['properties']
coordinates = []
for city in city_names:
coord = get_geo_coordinates(city)
if coord:
coordinates.append(coord)
if feature_type == "Polygon":
coordinates.append(coordinates[0]) # Close the polygon
return {
"type": "FeatureCollection",
"features": [{
"type": "Feature",
"properties": properties,
"geometry": {
"type": feature_type,
"coordinates": [coordinates] if feature_type == "Polygon" else coordinates
}
}]
}
# Generate map image
def save_map_image(geojson_data):
m = Map()
geo_layer = GeoJson(geojson_data, name="Feature map")
geo_layer.add_to(m)
bounds = get_bounds(geojson_data)
m.fit_bounds(bounds)
img_data = m._to_png(5)
img = Image.open(io.BytesIO(img_data))
img.save('map_image.png')
return 'map_image.png'
# ControlNet pipeline setup
controlnet = ControlNetModel.from_pretrained("lllyasviel/control_v11p_sd15_inpaint", torch_dtype=torch.float16)
pipeline = StableDiffusionControlNetInpaintPipeline.from_pretrained(
"runwayml/stable-diffusion-inpainting", controlnet=controlnet, torch_dtype=torch.float16
)
pipeline.enable_model_cpu_offload()
def generate_satellite_image(init_image_path, mask_image_path, prompt):
init_image = Image.open(init_image_path)
mask_image = Image.open(mask_image_path)
control_image = make_inpaint_condition(init_image, mask_image)
result = pipeline(prompt=prompt, image=init_image, mask_image=mask_image, control_image=control_image)
return result.images[0]
# Gradio UI
def handle_query(query):
# Process OpenAI response
response = process_openai_response(query)
geojson_data = generate_geojson(response)
# Save map image
map_image_path = save_map_image(geojson_data)
# Generate mask for ControlNet
empty_map = cv2.imread("empty_map_image.png")
map_image = cv2.imread(map_image_path)
difference = cv2.absdiff(cv2.cvtColor(empty_map, cv2.COLOR_BGR2GRAY), cv2.cvtColor(map_image, cv2.COLOR_BGR2GRAY))
_, mask = cv2.threshold(difference, 15, 255, cv2.THRESH_BINARY)
cv2.imwrite("mask.png", mask)
# Generate satellite image
satellite_image = generate_satellite_image("map_image.png", "mask.png", response['output']['feature_representation']['properties']['description'])
return map_image_path, satellite_image
# Gradio interface
with gr.Blocks() as demo:
with gr.Row():
query_input = gr.Textbox(label="Enter Query")
submit_btn = gr.Button("Submit")
with gr.Row():
map_output = gr.Image(label="Map Visualization")
satellite_output = gr.Image(label="Generated Satellite Image")
submit_btn.click(handle_query, inputs=[query_input], outputs=[map_output, satellite_output])
if __name__ == "__main__":
demo.launch()
|