Update app.py

app.py

@@ -1,11 +1,68 @@
-import io
-import base64
-import requests
+import gradio as gr
+import replicate
+import os
+import random
+import openai
 import numpy as np
 from PIL import Image
-import gradio as gr
+import requests
+import io
+import base64
 import zipfile
 
+# Set API tokens
+os.environ["REPLICATE_API_TOKEN"] = "r8_Brv0MtpmAiqrXrMrziyUXoSHuFV5hqs1Lw4Mo"
+# Initialize the Replicate client
+rep_client = replicate.Client()
+
+# Set your OpenAI API key
+OPENAI_API_KEY = "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"
+openai.api_key = OPENAI_API_KEY
+
+# Predefined prompts for the dropdown
+predefined_prompts = [
+    "Missing bolts on railway track",
+    "Cracks on railway track",
+    "Overgrown vegetation near railway track",
+    "Broken railings on railway bridge",
+    "Debris on railway track",
+    "Damaged railway platform"
+]
+
+def ask_rail_defect_question(question, model_name='ft:gpt-3.5-turbo-0125:personal::99NsSAeQ'):
+    response = openai.ChatCompletion.create(
+        model=model_name,
+        messages=[
+            {
+                "role": "system",
+                "content": "The assistant is knowledgeable about rail defects and can answer questions related to them.",
+            },
+            {
+                "role": "user",
+                "content": question,
+            }
+        ],
+    )
+    return response.choices[0].message['content']
+
+# Function to generate variations enhanced by the GPT model
+def generate_variations(base_prompt, number_of_variations):
+    locations = ["on the left side", "on the right side", "at the top", "at the bottom", "in the center"]
+    sizes = ["small", "medium", "large", "tiny", "huge"]
+    weather_conditions = ["under cold conditions", "during hot weather", "in dry weather", "in humid conditions", "under varying temperatures"]
+
+    variations = []
+    for _ in range(number_of_variations):
+        location = random.choice(locations)
+        size = random.choice(sizes)
+        weather = random.choice(weather_conditions)
+        
+        # Enhance the base prompt with the GPT model
+        enhanced_prompt = base_prompt
+        full_prompt = f"{enhanced_prompt}, with a {size} defect {location}, observed {weather}."
+        variations.append(full_prompt)
+    return variations
+
 def image_to_data_url(image):
     buffered = io.BytesIO()
     image.save(buffered, format="PNG")