removing selenium

app.py

@@ -1,15 +1,12 @@
 import os
 import json
-import cv2
 import numpy as np
 import torch
-from PIL import Image
-import io
+from PIL import Image, ImageDraw
 import gradio as gr
 from openai import OpenAI
 from geopy.geocoders import Nominatim
-from folium import Map, GeoJson
-from gradio_folium import Folium
+from staticmap import StaticMap, CircleMarker, Polygon
 from diffusers import ControlNetModel, StableDiffusionControlNetInpaintPipeline
 import spaces
 
@@ -67,7 +64,7 @@ Ensure all responses are descriptive and relevant to city names only, without co
       "content": [
         {
           "type": "text",
-          "text": query
+          "text": "draw a map in coconut triangle of sri lanka: The Coconut Triangle is a region in Sri Lanka that's known for its coconut production. It's made up of the districts of Kurunegala, Puttalam, and Gampaha."
         }
       ]
     }
@@ -109,38 +106,25 @@ def generate_geojson(response):
         }]
     }
 
-# Function to compute bounds from GeoJSON
+# Generate static map image
 @spaces.GPU
-def get_bounds(geojson):
-    coordinates = []
-    for feature in geojson["features"]:
+def generate_static_map(geojson_data):
+    m = StaticMap(500, 500)
+    for feature in geojson_data["features"]:
         geom_type = feature["geometry"]["type"]
         coords = feature["geometry"]["coordinates"]
+
         if geom_type == "Point":
-            coordinates.append(coords)
+            m.add_marker(CircleMarker((coords[0], coords[1]), 'blue', 10))
         elif geom_type in ["MultiPoint", "LineString"]:
-            coordinates.extend(coords)
-        elif geom_type in ["MultiLineString", "Polygon"]:
-            for part in coords:
-                coordinates.extend(part)
-        elif geom_type == "MultiPolygon":
+            for coord in coords:
+                m.add_marker(CircleMarker((coord[0], coord[1]), 'red', 10))
+        elif geom_type in ["Polygon", "MultiPolygon"]:
             for polygon in coords:
-                for part in polygon:
-                    coordinates.extend(part)
-    lats = [coord[1] for coord in coordinates]
-    lngs = [coord[0] for coord in coordinates]
-    return [[min(lats), min(lngs)], [max(lats), max(lngs)]]
+                m.add_polygon(Polygon([(c[0], c[1]) for c in polygon], 'green', 3))
 
-# Generate map image in memory
-@spaces.GPU
-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)
-    img_data = m._to_png(5)
-    return Image.open(io.BytesIO(img_data))
+    image = m.render(zoom=10)
+    return Image.fromarray(image)
 
 # ControlNet pipeline setup
 controlnet = ControlNetModel.from_pretrained("lllyasviel/control_v11p_sd15_inpaint", torch_dtype=torch.float16)
@@ -175,13 +159,12 @@ def handle_query(query):
     geojson_data = generate_geojson(response)
 
     # Generate map image
-    map_image = generate_map_image(geojson_data)
+    map_image = generate_static_map(geojson_data)
 
     # Generate mask for ControlNet
-    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)
+    empty_map = Image.new("RGB", map_image.size, "white")
+    difference = np.array(map_image) - np.array(empty_map)
+    mask = np.any(difference != 0, axis=-1).astype(np.uint8) * 255
 
     # Convert mask to PIL Image
     mask_image = Image.fromarray(mask)

requirements.txt

@@ -10,4 +10,5 @@ spaces
 torchvision
 opencv-python
 torch
+staticmap
 selenium