added pydantic, slim prompt

app.py

@@ -10,6 +10,8 @@ from staticmap import StaticMap, CircleMarker, Polygon
 from diffusers import ControlNetModel, StableDiffusionControlNetInpaintPipeline
 import spaces
 import logging
+from pydantic import BaseModel, ValidationError, Field
+from typing import List, Union
 
 # Set up logging
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
@@ -19,6 +21,20 @@ logger = logging.getLogger(__name__)
 openai_client = OpenAI(api_key=os.environ['OPENAI_API_KEY'])
 geolocator = Nominatim(user_agent="geoapi")
 
+# Define Pydantic models for GeoJSON validation
+class Geometry(BaseModel):
+    type: str
+    coordinates: Union[List[List[float]], List[List[List[float]]]]
+
+class Feature(BaseModel):
+    type: str = "Feature"
+    properties: dict
+    geometry: Geometry
+
+class FeatureCollection(BaseModel):
+    type: str = "FeatureCollection"
+    features: List[Feature]
+
 # Function to fetch coordinates
 @spaces.GPU
 def get_geo_coordinates(location_name):
@@ -37,43 +53,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 (Point) or Multiple Points (MultiPoint)**: 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. **Adhere to given example format always**\
-\"}\"}"
-        }
-      ]
-    },
-    {
-      "role": "user",
-      "content": [
-        {
-          "type": "text",
-          "text": query
-        }
-      ]
-    }
-  ],
+            {
+                "role": "system",
+                "content": "You are a skilled assistant answering geographical and historical questions in JSON format."
+            },
+            {
+                "role": "user",
+                "content": query
+            }
+        ],
         temperature=1,
         max_tokens=2048,
         top_p=1,
@@ -99,48 +87,55 @@ def generate_geojson(response):
     if feature_type == "Polygon":
         coordinates.append(coordinates[0])  # Close the polygon
 
-    return {
+    geojson_data = {
         "type": "FeatureCollection",
-        "features": [{
-            "type": "Feature",
-            "properties": properties,
-            "geometry": {
-                "type": feature_type,
-                "coordinates": [coordinates] if feature_type == "Polygon" else coordinates
+        "features": [
+            {
+                "type": "Feature",
+                "properties": properties,
+                "geometry": {
+                    "type": feature_type,
+                    "coordinates": [coordinates] if feature_type == "Polygon" else coordinates
+                }
             }
-        }]
+        ]
     }
 
+    # Validate the generated GeoJSON using Pydantic
+    try:
+        validated_geojson = FeatureCollection(**geojson_data)
+        logger.info("GeoJSON validation successful.")
+        return validated_geojson.dict()
+    except ValidationError as e:
+        logger.error(f"GeoJSON validation failed: {e}")
+        raise
+
 # Generate static map image
 @spaces.GPU
 def generate_static_map(geojson_data, invisible=False):
-    # Create a static map object with specified dimensions
     m = StaticMap(600, 600)
-    #log the geojson data
     logger.info(f"GeoJSON data: {geojson_data}")
-    # Process each feature in the GeoJSON
+
     for feature in geojson_data["features"]:
         geom_type = feature["geometry"]["type"]
         coords = feature["geometry"]["coordinates"]
 
         if geom_type == "Point":
-            m.add_marker(CircleMarker((coords[0][0], coords[0][1]), '#1C00ff00' if invisible == True else 'blue', 1000))
+            m.add_marker(CircleMarker((coords[0][0], coords[0][1]), '#1C00ff00' if invisible else 'blue', 1000))
         elif geom_type in ["MultiPoint", "LineString"]:
             for coord in coords:
-                m.add_marker(CircleMarker((coord[0], coord[1]), '#1C00ff00' if invisible == True else 'blue', 1000))
+                m.add_marker(CircleMarker((coord[0], coord[1]), '#1C00ff00' if invisible else 'blue', 1000))
         elif geom_type in ["Polygon", "MultiPolygon"]:
             for polygon in coords:
-                m.add_polygon(Polygon([(c[0], c[1]) for c in polygon], '#1C00ff00' if invisible == True else 'blue', 3))
-
-    return m.render() #zoom=10
+                m.add_polygon(Polygon([(c[0], c[1]) for c in polygon], '#1C00ff00' if invisible else 'blue', 3))
 
+    return m.render()
 
 # ControlNet pipeline setup
 controlnet = ControlNetModel.from_pretrained("lllyasviel/control_v11p_sd15_inpaint", torch_dtype=torch.float16)
 pipeline = StableDiffusionControlNetInpaintPipeline.from_pretrained(
     "stable-diffusion-v1-5/stable-diffusion-inpainting", controlnet=controlnet, torch_dtype=torch.float16
 )
-# ZeroGPU compatibility
 pipeline.to('cuda')
 
 @spaces.GPU
@@ -164,66 +159,23 @@ def generate_satellite_image(init_image, mask_image, prompt):
         control_image=control_image,
         strength=0.42,
         guidance_scale=62
-        )
+    )
     return result.images[0]
 
 # Gradio UI
 @spaces.GPU
 def handle_query(query):
-    # Process OpenAI response
     response = process_openai_response(query)
     geojson_data = generate_geojson(response)
 
-    geojson_data = {
-    "type": "FeatureCollection",
-    "features": [
-        {
-            "type": "Feature",
-            "properties": {
-                "description": "satellite image of the Coconut Triangle region, green fill, 4k, map, detailed, coconut palms, lush vegetation, high contrast"
-            },
-            "geometry": {
-                "type": "Polygon",
-                "coordinates": [
-                    [
-                        [
-                            80.364908,
-                            7.4870464
-                        ],
-                        [
-                            79.82933709234904,
-                            7.981840249999999
-                        ],
-                        [
-                            79.91598756451819,
-                            7.1190247499999995
-                        ],
-                        [
-                            80.364908,
-                            7.4870464
-                        ]
-                    ]
-                ]
-            }
-        }
-    ]
-}
-
-
-    # Generate the main map image
     map_image = generate_static_map(geojson_data)
+    empty_map_image = generate_static_map(geojson_data, invisible=True)
 
-    empty_map_image = generate_static_map(geojson_data, invisible=True)  # Empty map with the same bounds
-
-    # Create the mask
     difference = np.abs(np.array(map_image.convert("RGB")) - np.array(empty_map_image.convert("RGB")))
-    threshold = 10  # Tolerance for difference
+    threshold = 10
     mask = (np.sum(difference, axis=-1) > threshold).astype(np.uint8) * 255
 
-    # Convert the mask to a PIL image
     mask_image = Image.fromarray(mask, mode="L")
-
-    # Generate the satellite image
     satellite_image = generate_satellite_image(
         empty_map_image, mask_image, response['output']['feature_representation']['properties']['description']
     )
@@ -238,7 +190,6 @@ query_options = [
     "Due to considerable rainfall in the up- and mid- stream areas of Kala Oya, the Rajanganaya reservoir is now spilling at a rate of 17,000 cubic feet per second, the department said."
 ]
 
-# Gradio interface
 with gr.Blocks() as demo:
     with gr.Row():
         selected_query = gr.Dropdown(label="Select Query", choices=query_options, value=query_options[-1])

requirements.txt

@@ -11,4 +11,4 @@ torchvision
 opencv-python
 torch
 staticmap
-selenium
+pydantic