in memory update

app.py

@@ -13,9 +13,6 @@ from gradio_folium import Folium
 from diffusers import ControlNetModel, StableDiffusionControlNetInpaintPipeline
 import spaces
 
-from selenium.webdriver.firefox.options import Options
-from selenium import webdriver
-
 # Initialize APIs
 openai_client = OpenAI(api_key=os.environ['OPENAI_API_KEY'])
 geolocator = Nominatim(user_agent="geoapi")
@@ -38,51 +35,15 @@ def process_openai_response(query):
     response = openai_client.chat.completions.create(
         model="gpt-4o-mini",
         messages=[
-    {
-      "role": "system",
-      "content": [
-        {
-          "type": "text",
-          "text": "\"input\": \"\"\"You are a skilled assistant answering geographical and historical questions. For each question, generate a structured output in JSON format, based on city names without coordinates. The response should include:\
-Answer: A concise response to the question.\
-Feature Representation: A feature type based on city names (Point, LineString, Polygon, MultiPoint, MultiLineString, MultiPolygon, GeometryCollection).\
-Description: A prompt for a diffusion model describing the what should we draw regarding that.\
-\
-Handle the following cases:\
-\
-1. **Single or Multiple Points**: Create a point or a list of points for multiple cities.\
-2. **LineString**: Create a line between two cities.\
-3. **Polygon**: Represent an area formed by three or more cities (closed). Example: Cities forming a triangle (A, B, C).\
-4. **MultiPoint, MultiLineString, MultiPolygon, GeometryCollection**: Use as needed based on the question.\
-\
-For example, if asked about cities forming a polygon, create a feature like this:\
-\
-Input: Mark an area with three cities.\
-Output: {\"input\": \"Mark an area with three cities.\", \"output\": {\"answer\": \"The cities A, B, and C form a triangle.\", \"feature_representation\": {\"type\": \"Polygon\", \"cities\": [\"A\", \"B\", \"C\"], \"properties\": {\"description\": \"satelite image of a plantation, green fill, 4k, map, detailed, greenary, plants, vegitation, high contrast\"}}}}\
-\
-Ensure all responses are descriptive and relevant to city names only, without coordinates.\
-\"}\"}"
-        }
-      ]
-    },
-    {
-      "role": "user",
-      "content": [
-        {
-          "type": "text",
-          "text": query
-        }
-      ]
-    }
-  ],
-  temperature=1,
-  max_tokens=2048,
-  top_p=1,
-  frequency_penalty=0,
-  presence_penalty=0,
-  response_format={
-    "type": "json_object"
-  }
+            {"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)
 
@@ -136,25 +97,16 @@ def get_bounds(geojson):
     lngs = [coord[0] for coord in coordinates]
     return [[min(lats), min(lngs)], [max(lats), max(lngs)]]
 
-# Generate map image
+# Generate map image in memory
 @spaces.GPU
-def save_map_image(geojson_data):
+def generate_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)
-
-    # Configure Selenium for headless operation
-    options = Options()
-    options.add_argument("--headless")  # Enable headless mode
-    driver = webdriver.Firefox(options=options)  # Ensure GeckoDriver is properly installed and in PATH
-    img_data = m._to_png(5, driver=driver)
-    driver.quit()
-
-    img = Image.open(io.BytesIO(img_data))
-    img.save('map_image.png')
-    return 'map_image.png'
+    img_data = m._to_png(5)
+    return Image.open(io.BytesIO(img_data))
 
 # ControlNet pipeline setup
 controlnet = ControlNetModel.from_pretrained("lllyasviel/control_v11p_sd15_inpaint", torch_dtype=torch.float16)
@@ -176,9 +128,7 @@ def make_inpaint_condition(init_image, mask_image):
     return init_image
 
 @spaces.GPU
-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)
+def generate_satellite_image(init_image, mask_image, prompt):
     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]
@@ -190,20 +140,22 @@ def handle_query(query):
     response = process_openai_response(query)
     geojson_data = generate_geojson(response)
 
-    # Save map image
-    map_image_path = save_map_image(geojson_data)
+    # Generate map image
+    map_image = generate_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))
+    empty_map = cv2.cvtColor(np.array(generate_map_image({"type": "FeatureCollection", "features": []})), cv2.COLOR_BGR2GRAY)
+    map_image_array = cv2.cvtColor(np.array(map_image), cv2.COLOR_BGR2GRAY)
+    difference = cv2.absdiff(empty_map, map_image_array)
     _, mask = cv2.threshold(difference, 15, 255, cv2.THRESH_BINARY)
-    cv2.imwrite("mask.png", mask)
+
+    # Convert mask to PIL Image
+    mask_image = Image.fromarray(mask)
 
     # Generate satellite image
-    satellite_image = generate_satellite_image("map_image.png", "mask.png", response['output']['feature_representation']['properties']['description'])
+    satellite_image = generate_satellite_image(map_image, mask_image, response['output']['feature_representation']['properties']['description'])
 
-    return map_image_path, satellite_image
+    return map_image, satellite_image
 
 # Gradio interface
 with gr.Blocks() as demo:
@@ -217,4 +169,4 @@ with gr.Blocks() as demo:
     submit_btn.click(handle_query, inputs=[query_input], outputs=[map_output, satellite_output])
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()